Add scripts to create MS MARCO v1 Document Jsonl Collection (#905)

castorini · Dec 14, 2021 · fc54de6 · fc54de6
1 parent 11ce241
commit fc54de6
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 1 deletion.
diff --git a/scripts/msmarco-doc/convert_msmarco_doc_to_anserini_collection.py b/scripts/msmarco-doc/convert_msmarco_doc_to_anserini_collection.py
@@ -0,0 +1,32 @@
+import argparse
+import gzip
+import json
+import os
+from tqdm import tqdm
+
+
+def generate_output_dict(doc):
+    doc_id, doc_url, doc_title, doc_text = doc[0], doc[1], doc[2], doc[3]
+    doc_text = doc_text.strip()
+    doc_text = f'{doc_url}\n{doc_title}\n{doc_text}'
+    output_dict = {'id': doc_id, 'contents': doc_text}
+    return output_dict
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Convert MS MARCO V1 document ranking corpus to 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.')
+
+    args = parser.parse_args()
+
+    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')
+
+    for line in tqdm(f_corpus):
+        output_dict = generate_output_dict(line.split('\t'))
+        f_out.write(json.dumps(output_dict) + '\n')
+    print('Done!')
diff --git a/scripts/msmarco-doc/convert_msmarco_doc_to_segmented_anserini_collection.py b/scripts/msmarco-doc/convert_msmarco_doc_to_segmented_anserini_collection.py
@@ -0,0 +1,48 @@
+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!')
diff --git a/scripts/msmarco_v2/convert_passage_run_to_doc_run.py b/scripts/msmarco_v2/convert_passage_run_to_doc_run.py
@@ -75,7 +75,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-	parser = argparse.ArgumentParser(description="Conver runfile that contain psg id into runfile that contain doc id.")
+	parser = argparse.ArgumentParser(description="Convert runfile that contain psg id into runfile that contain doc id.")
 	parser.add_argument("--input", type=str, required=True, help="path to msmarco passage runfile.")
 	parser.add_argument("--id-map", type=str, required=True, help="directory that contains msmarco passage-doc id mapping .tsv files. Generated by `build_passage_to_doc_id_map.py` under the same directory.")
 	parser.add_argument("--output", type=str, required=True, help="output path to store document id runfile.")