|
"""
|
|
Divide a document collection into N-word/token passage spans (with wrap-around for last passage).
|
|
"""
|
|
|
|
import os
|
|
import math
|
|
import ujson
|
|
import random
|
|
|
|
from multiprocessing import Pool
|
|
from argparse import ArgumentParser
|
|
from colbert.utils.utils import print_message
|
|
|
|
Format1 = 'docid,text'
|
|
Format2 = 'docid,text,title'
|
|
Format3 = 'docid,url,title,text'
|
|
|
|
|
|
def process_page(inp):
|
|
"""
|
|
Wraps around if we split: make sure last passage isn't too short.
|
|
This is meant to be similar to the DPR preprocessing.
|
|
"""
|
|
|
|
(nwords, overlap, tokenizer), (title_idx, docid, title, url, content) = inp
|
|
|
|
if tokenizer is None:
|
|
words = content.split()
|
|
else:
|
|
words = tokenizer.tokenize(content)
|
|
|
|
words_ = (words + words) if len(words) > nwords else words
|
|
passages = [words_[offset:offset + nwords] for offset in range(0, len(words) - overlap, nwords - overlap)]
|
|
|
|
assert all(len(psg) in [len(words), nwords] for psg in passages), (list(map(len, passages)), len(words))
|
|
|
|
if tokenizer is None:
|
|
passages = [' '.join(psg) for psg in passages]
|
|
else:
|
|
passages = [' '.join(psg).replace(' ##', '') for psg in passages]
|
|
|
|
if title_idx % 100000 == 0:
|
|
print("#> ", title_idx, '\t\t\t', title)
|
|
|
|
for p in passages:
|
|
print("$$$ ", '\t\t', p)
|
|
print()
|
|
|
|
print()
|
|
print()
|
|
print()
|
|
|
|
return (docid, title, url, passages)
|
|
|
|
|
|
def main(args):
|
|
random.seed(12345)
|
|
print_message("#> Starting...")
|
|
|
|
letter = 'w' if not args.use_wordpiece else 't'
|
|
output_path = f'{args.input}.{letter}{args.nwords}_{args.overlap}'
|
|
assert not os.path.exists(output_path)
|
|
|
|
RawCollection = []
|
|
Collection = []
|
|
|
|
NumIllFormattedLines = 0
|
|
|
|
with open(args.input) as f:
|
|
for line_idx, line in enumerate(f):
|
|
if line_idx % (100*1000) == 0:
|
|
print(line_idx, end=' ')
|
|
|
|
title, url = None, None
|
|
|
|
try:
|
|
line = line.strip().split('\t')
|
|
|
|
if args.format == Format1:
|
|
docid, doc = line
|
|
elif args.format == Format2:
|
|
docid, doc, title = line
|
|
elif args.format == Format3:
|
|
docid, url, title, doc = line
|
|
|
|
RawCollection.append((line_idx, docid, title, url, doc))
|
|
except:
|
|
NumIllFormattedLines += 1
|
|
|
|
if NumIllFormattedLines % 1000 == 0:
|
|
print(f'\n[{line_idx}] NumIllFormattedLines = {NumIllFormattedLines}\n')
|
|
|
|
print()
|
|
print_message("# of documents is", len(RawCollection), '\n')
|
|
|
|
p = Pool(args.nthreads)
|
|
|
|
print_message("#> Starting parallel processing...")
|
|
|
|
tokenizer = None
|
|
if args.use_wordpiece:
|
|
from transformers import BertTokenizerFast
|
|
tokenizer = BertTokenizerFast.from_pretrained('bert-base-multilingual-uncased')
|
|
|
|
process_page_params = [(args.nwords, args.overlap, tokenizer)] * len(RawCollection)
|
|
Collection = p.map(process_page, zip(process_page_params, RawCollection))
|
|
|
|
print_message(f"#> Writing to {output_path} ...")
|
|
with open(output_path, 'w') as f:
|
|
line_idx = 1
|
|
|
|
if args.format == Format1:
|
|
f.write('\t'.join(['id', 'text']) + '\n')
|
|
elif args.format == Format2:
|
|
f.write('\t'.join(['id', 'text', 'title']) + '\n')
|
|
elif args.format == Format3:
|
|
f.write('\t'.join(['id', 'text', 'title', 'docid']) + '\n')
|
|
|
|
for docid, title, url, passages in Collection:
|
|
for passage in passages:
|
|
if args.format == Format1:
|
|
f.write('\t'.join([str(line_idx), passage]) + '\n')
|
|
elif args.format == Format2:
|
|
f.write('\t'.join([str(line_idx), passage, title]) + '\n')
|
|
elif args.format == Format3:
|
|
f.write('\t'.join([str(line_idx), passage, title, docid]) + '\n')
|
|
|
|
line_idx += 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = ArgumentParser(description="docs2passages.")
|
|
|
|
|
|
parser.add_argument('--input', dest='input', required=True)
|
|
parser.add_argument('--format', dest='format', required=True, choices=[Format1, Format2, Format3])
|
|
|
|
|
|
parser.add_argument('--use-wordpiece', dest='use_wordpiece', default=False, action='store_true')
|
|
parser.add_argument('--nwords', dest='nwords', default=100, type=int)
|
|
parser.add_argument('--overlap', dest='overlap', default=0, type=int)
|
|
|
|
|
|
parser.add_argument('--nthreads', dest='nthreads', default=28, type=int)
|
|
|
|
args = parser.parse_args()
|
|
assert args.nwords in range(50, 500)
|
|
|
|
main(args)
|
|
|
