| """
|
| 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)
|
|
|
