Spaces:
Runtime error
Runtime error
File size: 1,873 Bytes
d6585f5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
import argparse
import gzip
import json
import os
# Uses space==2.1.6
import spacy
from tqdm import tqdm
def generate_output_dicts(doc, nlp, max_length, stride):
doc_id, doc_url, doc_title, doc_text = doc[0], doc[1], doc[2], doc[3]
doc_text = doc_text.strip()
doc = nlp(doc_text[:10000])
sentences = [sent.string.strip() for sent in doc.sents]
output_dicts = []
for ind, pos in enumerate(range(0, len(sentences), stride)):
segment = ' '.join(sentences[pos:pos + max_length])
doc_text = f'{doc_url}\n{doc_title}\n{segment}'
output_dicts.append({'id': f'{doc_id}#{ind}', 'contents': doc_text})
if pos + max_length >= len(sentences):
break
return output_dicts
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='Convert MS MARCO V1 document ranking corpus to seg anserini\'s default jsonl collection format')
parser.add_argument('--original_docs_path', required=True, help='Original corpus file.')
parser.add_argument('--output_docs_path', required=True, help='Output file in the anserini jsonl format.')
parser.add_argument('--stride', default=5, help='Sliding-window stride')
parser.add_argument('--max_length', default=10, help='Sliding-window length')
args = parser.parse_args()
# Load spacy model
nlp = spacy.blank("en")
nlp.add_pipe(nlp.create_pipe("sentencizer"))
os.makedirs(os.path.dirname(args.output_docs_path), exist_ok=True)
f_corpus = gzip.open(args.original_docs_path, mode='rt')
f_out = open(args.output_docs_path, 'w')
print('Creating collection...')
for line in tqdm(f_corpus):
output_dicts = generate_output_dicts(line.split('\t'), nlp, args.max_length, args.stride)
for output_dict in output_dicts:
f_out.write(json.dumps(output_dict) + '\n')
print('Done!')
|