wehub-resource-sync
90 lines
3.8 KiB
Python
90 lines
3.8 KiB
Python
import os
|
|
import argparse
|
|
import json
|
|
import sys
|
|
sys.path.insert(0, 'src/')
|
|
|
|
from tqdm import tqdm
|
|
from typing import Dict, Any
|
|
from datasets import Dataset
|
|
|
|
from evaluate_dpr_retrieval import has_answers, SimpleTokenizer, evaluate_retrieval
|
|
from data_utils import load_query_answers, load_corpus
|
|
from utils import save_json_to_file
|
|
from logger_config import logger
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description='Convert an TREC run to DPR retrieval result json.')
|
|
parser.add_argument('--data-dir', required=True, help='data dir')
|
|
parser.add_argument('--topics', required=True, help='topic name')
|
|
parser.add_argument('--topk', type=int, nargs='+', help="topk to evaluate")
|
|
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('--output', required=True, help='Output DPR Retrieval json file.')
|
|
args = parser.parse_args()
|
|
|
|
qas = load_query_answers(path=args.topics)
|
|
corpus = load_corpus(path=os.path.join(args.data_dir, 'passages.jsonl.gz'))
|
|
|
|
retrieval = {}
|
|
tokenizer = SimpleTokenizer()
|
|
|
|
predictions = []
|
|
for line in tqdm(open(args.input), mininterval=1):
|
|
question_id, doc_idx, _, score = line.strip().split('\t')[:4]
|
|
|
|
predictions.append({'question_id': question_id,
|
|
'doc_idx': int(doc_idx),
|
|
'score': score})
|
|
|
|
dataset = Dataset.from_dict({'question_id': [ex['question_id'] for ex in predictions],
|
|
'doc_idx': [ex['doc_idx'] for ex in predictions],
|
|
'score': [ex['score'] for ex in predictions]})
|
|
logger.info('Get {} predictions in total'.format(len(dataset)))
|
|
|
|
def _map_func(example: Dict[str, Any]) -> dict:
|
|
question_id, doc_idx, score = example['question_id'], example['doc_idx'], example['score']
|
|
|
|
question = qas[question_id]['query']
|
|
answers = qas[question_id]['answers']
|
|
title, text = corpus[doc_idx]['title'], corpus[doc_idx]['contents']
|
|
ctx = '{}\n{}'.format(title, text)
|
|
|
|
answer_exist = has_answers(text, answers, tokenizer, args.regex)
|
|
|
|
example['question'] = question
|
|
example['answers'] = answers
|
|
example['docid'] = doc_idx
|
|
example['has_answer'] = answer_exist
|
|
if args.store_raw:
|
|
example['text'] = ctx
|
|
|
|
return example
|
|
|
|
dataset = dataset.map(_map_func,
|
|
num_proc=min(os.cpu_count(), 16))
|
|
|
|
retrieval = {}
|
|
for ex in tqdm(dataset, mininterval=2, desc='convert to dpr format'):
|
|
question_id, question, answers = ex['question_id'], ex['question'], ex['answers']
|
|
if question_id not in retrieval:
|
|
retrieval[question_id] = {'question': question, 'answers': answers, 'contexts': []}
|
|
retrieval[question_id]['contexts'].append(
|
|
{k: ex[k] for k in ['docid', 'score', 'text', 'has_answer'] if k in ex}
|
|
)
|
|
|
|
save_json_to_file(retrieval, path=args.output)
|
|
logger.info('Convert {} to {} done'.format(args.input, args.output))
|
|
|
|
metrics = evaluate_retrieval(retrieval_file=args.output,
|
|
topk=args.topk,
|
|
regex=args.regex)
|
|
logger.info('{} recall metrics: {}'.format(
|
|
os.path.basename(args.output),
|
|
json.dumps(metrics, ensure_ascii=False, indent=4)))
|
|
|
|
base_dir, base_name = os.path.dirname(args.output), os.path.basename(args.output)
|
|
save_json_to_file(metrics, path='{}/metrics_{}'.format(base_dir, base_name))
|