wehub-resource-sync
76 lines
3.0 KiB
Python
76 lines
3.0 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import tqdm
|
|
import argparse
|
|
sys.path.insert(0, './src')
|
|
|
|
from typing import List, Dict
|
|
|
|
from utils import save_json_to_file
|
|
from logger_config import logger
|
|
from data_utils import load_qrels, load_corpus, load_queries, load_msmarco_predictions, ScoredDoc
|
|
from metrics import get_rel_threshold
|
|
|
|
parser = argparse.ArgumentParser(description='convert ms-marco predictions to a human-readable format')
|
|
parser.add_argument('--in-path', default='', type=str, metavar='N',
|
|
help='path to predictions in msmarco output format')
|
|
parser.add_argument('--split', default='dev', type=str, metavar='N',
|
|
help='which split to use')
|
|
parser.add_argument('--data-dir', default='./data/msmarco/', type=str, metavar='N',
|
|
help='data dir')
|
|
|
|
args = parser.parse_args()
|
|
logger.info('Args={}'.format(json.dumps(args.__dict__, ensure_ascii=False, indent=4)))
|
|
|
|
|
|
def main(topk: int = 10):
|
|
predictions: Dict[str, List[ScoredDoc]] = load_msmarco_predictions(path=args.in_path)
|
|
path_qrels = '{}/{}_qrels.txt'.format(args.data_dir, args.split)
|
|
qrels = load_qrels(path=path_qrels) if os.path.exists(path_qrels) else None
|
|
queries = load_queries(path='{}/{}_queries.tsv'.format(args.data_dir, args.split))
|
|
corpus = load_corpus(path='{}/passages.jsonl.gz'.format(args.data_dir))
|
|
|
|
pred_infos = []
|
|
out_path = '{}.details.json'.format(args.in_path)
|
|
rel_threshold = get_rel_threshold(qrels) if qrels else -1
|
|
for qid in tqdm.tqdm(queries):
|
|
pred_docs = []
|
|
for scored_doc in predictions[qid][:topk]:
|
|
correct = qrels is not None and scored_doc.pid in qrels[qid] and qrels[qid][scored_doc.pid] >= rel_threshold
|
|
pred_docs.append({'id': scored_doc.pid,
|
|
'contents': corpus[int(scored_doc.pid)]['contents'],
|
|
'title': corpus[int(scored_doc.pid)]['title'],
|
|
'score': scored_doc.score})
|
|
if qrels is not None:
|
|
pred_docs[-1]['correct'] = correct
|
|
if correct: break
|
|
|
|
gold_rank, gold_score = -1, -1
|
|
for idx, scored_doc in enumerate(predictions[qid]):
|
|
if qrels is None:
|
|
break
|
|
if scored_doc.pid in qrels[qid] and qrels[qid][scored_doc.pid] >= rel_threshold:
|
|
gold_rank = idx + 1
|
|
gold_score = scored_doc.score
|
|
break
|
|
|
|
pred_info = {'query_id': qid,
|
|
'query': queries[qid],
|
|
'pred_docs': pred_docs}
|
|
if qrels is not None:
|
|
pred_info.update({
|
|
'gold_docs': [corpus[int(doc_id)] for doc_id in qrels[qid] if qrels[qid][doc_id] >= rel_threshold],
|
|
'gold_score': gold_score,
|
|
'gold_rank': gold_rank
|
|
})
|
|
|
|
pred_infos.append(pred_info)
|
|
|
|
save_json_to_file(pred_infos, out_path)
|
|
logger.info('Save prediction details to {}'.format(out_path))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|