Spaces:
Runtime error
Runtime error
File size: 3,759 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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
#
# 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 tqdm import tqdm
from pyserini.search import get_topics, get_topics_with_reader
from pyserini.search.lucene import LuceneSearcher
from pyserini.eval.evaluate_dpr_retrieval import has_answers, SimpleTokenizer
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Convert an TREC run to DPR retrieval result json.')
parser.add_argument('--topics', help='topic name')
parser.add_argument('--topics-file', help='path to a topics file')
parser.add_argument('--topics-reader', help='anserini TopicReader class')
parser.add_argument('--index', required=True, help='Anserini Index that contains raw')
parser.add_argument('--input', required=True, help='Input TREC run file.')
parser.add_argument('--store-raw', action='store_true', help='Store raw text of passage')
parser.add_argument('--regex', action='store_true', default=False, help="regex match")
parser.add_argument('--combine-title-text', action='store_true', help="Make context the concatenation of title and text.")
parser.add_argument('--output', required=True, help='Output DPR Retrieval json file.')
args = parser.parse_args()
if args.topics_file:
qas = get_topics_with_reader(args.topics_reader, args.topics_file)
elif args.topics:
qas = get_topics(args.topics)
else:
print("No topics file or topics name was provided")
if os.path.exists(args.index):
searcher = LuceneSearcher(args.index)
else:
searcher = LuceneSearcher.from_prebuilt_index(args.index)
if not searcher:
exit()
retrieval = {}
tokenizer = SimpleTokenizer()
with open(args.input) as f_in:
for line in tqdm(f_in.readlines()):
question_id, _, doc_id, _, score, _ = line.strip().split()
question_id = int(question_id)
question = qas[question_id]['title']
answers = qas[question_id]['answers']
if answers[0] == '"':
answers = answers[1:-1].replace('""', '"')
answers = eval(answers)
if args.combine_title_text:
passage = json.loads(searcher.doc(doc_id).raw())
ctx = passage['title'] + "\n" + passage['text']
else:
ctx = json.loads(searcher.doc(doc_id).raw())['contents']
if question_id not in retrieval:
retrieval[question_id] = {'question': question, 'answers': answers, 'contexts': []}
title, text = ctx.split('\n')
answer_exist = has_answers(text, answers, tokenizer, args.regex)
if args.store_raw:
retrieval[question_id]['contexts'].append(
{'docid': doc_id,
'score': score,
'text': ctx,
'has_answer': answer_exist}
)
else:
retrieval[question_id]['contexts'].append(
{'docid': doc_id, 'score': score, 'has_answer': answer_exist}
)
json.dump(retrieval, open(args.output, 'w'), indent=4, ensure_ascii=False)
|