Spaces:
Runtime error
Runtime error
| # | |
| # Pyserini: Reproducible IR research with sparse and dense representations | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| import argparse | |
| import json | |
| import os | |
| from pyserini.analysis import Analyzer, get_lucene_analyzer | |
| """ | |
| append d2q prediction as an extra field to collection jsonl | |
| """ | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser( | |
| description='Converts MSMARCO\'s tsv collection to Anserini jsonl files.') | |
| parser.add_argument('--collection_path', required=True, help='MS MARCO .tsv collection file') | |
| parser.add_argument('--predictions', required=True, help='File containing predicted queries.') | |
| parser.add_argument('--output_folder', required=True, help='output folder') | |
| parser.add_argument('--max_docs_per_file', default=1000000, type=int, | |
| help='maximum number of documents in each jsonl file.') | |
| args = parser.parse_args() | |
| if not os.path.exists(args.output_folder): | |
| os.makedirs(args.output_folder) | |
| analyzer = Analyzer(get_lucene_analyzer()) | |
| print('Converting collection...') | |
| file_index = 0 | |
| new_words = 0 | |
| total_words = 0 | |
| with open(args.collection_path) as f_corpus, open(args.predictions) as f_pred: | |
| for i, (line_doc, line_pred) in enumerate(zip(f_corpus, f_pred)): | |
| # Write to a new file when the current one reaches maximum capacity. | |
| if i % args.max_docs_per_file == 0: | |
| if i > 0: | |
| output_jsonl_file.close() | |
| output_path = os.path.join(args.output_folder, f'docs{file_index:02d}.json') | |
| output_jsonl_file = open(output_path, 'w') | |
| file_index += 1 | |
| doc_json = json.loads(line_doc) | |
| pred_text = line_pred.rstrip() | |
| predict_text = pred_text + ' ' | |
| analyzed = analyzer.analyze(predict_text) | |
| for token in analyzed: | |
| assert ' ' not in token | |
| predict = ' '.join(analyzed) | |
| doc_json['predict'] = predict | |
| output_jsonl_file.write(json.dumps(doc_json) + '\n') | |
| if i % 100000 == 0: | |
| print('Converted {} docs in {} files'.format(i, file_index)) | |
| output_jsonl_file.close() | |
| print('Done!') |