-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
67 lines (58 loc) · 2.3 KB
/
utils.py
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
import os
import re
import random
from collections import defaultdict
import subprocess
def generate_rank(input_path, output_path):
score_dict = defaultdict(list)
for line in open(input_path):
query_id, para_id, score = line.split("\t")
score_dict[int(query_id)].append((float(score), int(para_id)))
with open(output_path, "w") as outFile:
for query_id, para_lst in score_dict.items():
random.shuffle(para_lst)
para_lst = sorted(para_lst, key=lambda x:x[0], reverse=True)
for rank_idx, (score, para_id) in enumerate(para_lst):
score = 1000 - rank_idx
outFile.write("{} 0 {} {} {} run\n".format(query_id, para_id, rank_idx+1, score))
def eval_results(run_file_path,
eval_script="./ms_marco_eval.py",
qrels="/data/users/liuyuan/msmarco-doc/dev/msmarco-docdev-qrels.tsv" ):
assert os.path.exists(eval_script) and os.path.exists(qrels)
dataset = 'marcodoc'
result = subprocess.check_output(['python', eval_script, qrels, run_file_path, dataset])
match = re.search('MRR @100: ([\d.]+)', result.decode('utf-8'))
mrr = float(match.group(1))
return mrr
def eval_passage_results(run_file_path,
eval_script="./ms_marco_eval.py",
qrels="./data/msmarco_pas/qrels.dev.small.tsv" ):
assert os.path.exists(eval_script) and os.path.exists(qrels)
dataset = 'marcopassage'
result = subprocess.check_output(['python', eval_script, qrels, run_file_path, dataset])
match = re.search('MRR @100: ([\d.]+)', result.decode('utf-8'))
mrr = float(match.group(1))
return mrr
def read_qd_pairs(QD_FILE_PATH):
qds = {}
with open(QD_FILE_PATH, 'r', encoding='utf-8') as qdsf:
for line in qdsf:
ss = line.strip().split('\t')
qid = int(ss[0])
did = int(ss[1])
if qid not in qds:
qds[qid] = []
qds[qid].append(did)
return qds
def read_qds_pairs(QD_FILE_PATH):
qds = {}
with open(QD_FILE_PATH, 'r', encoding='utf-8') as qdsf:
for line in qdsf:
ss = line.strip().split('\t')
qid = int(ss[0])
did = int(ss[1])
score = float(ss[2])
if qid not in qds:
qds[qid] = {}
qds[qid][did] = score
return qds