|
"""
|
|
Evaluate MS MARCO Passages ranking.
|
|
"""
|
|
|
|
import os
|
|
import math
|
|
import tqdm
|
|
import ujson
|
|
import random
|
|
|
|
from argparse import ArgumentParser
|
|
from collections import defaultdict
|
|
from colbert.utils.utils import print_message, file_tqdm
|
|
|
|
|
|
def main(args):
|
|
qid2positives = defaultdict(list)
|
|
qid2ranking = defaultdict(list)
|
|
qid2mrr = {}
|
|
qid2recall = {depth: {} for depth in [50, 200, 1000]}
|
|
|
|
with open(args.qrels) as f:
|
|
print_message(f"#> Loading QRELs from {args.qrels} ..")
|
|
for line in file_tqdm(f):
|
|
qid, _, pid, label = map(int, line.strip().split())
|
|
assert label == 1
|
|
|
|
qid2positives[qid].append(pid)
|
|
|
|
with open(args.ranking) as f:
|
|
print_message(f"#> Loading ranked lists from {args.ranking} ..")
|
|
for line in file_tqdm(f):
|
|
qid, pid, rank, *score = line.strip().split('\t')
|
|
qid, pid, rank = int(qid), int(pid), int(rank)
|
|
|
|
if len(score) > 0:
|
|
assert len(score) == 1
|
|
score = float(score[0])
|
|
else:
|
|
score = None
|
|
|
|
qid2ranking[qid].append((rank, pid, score))
|
|
|
|
assert set.issubset(set(qid2ranking.keys()), set(qid2positives.keys()))
|
|
|
|
num_judged_queries = len(qid2positives)
|
|
num_ranked_queries = len(qid2ranking)
|
|
|
|
if num_judged_queries != num_ranked_queries:
|
|
print()
|
|
print_message("#> [WARNING] num_judged_queries != num_ranked_queries")
|
|
print_message(f"#> {num_judged_queries} != {num_ranked_queries}")
|
|
print()
|
|
|
|
print_message(f"#> Computing MRR@10 for {num_judged_queries} queries.")
|
|
|
|
for qid in tqdm.tqdm(qid2positives):
|
|
ranking = qid2ranking[qid]
|
|
positives = qid2positives[qid]
|
|
|
|
for rank, (_, pid, _) in enumerate(ranking):
|
|
rank = rank + 1
|
|
|
|
if pid in positives:
|
|
if rank <= 10:
|
|
qid2mrr[qid] = 1.0 / rank
|
|
break
|
|
|
|
for rank, (_, pid, _) in enumerate(ranking):
|
|
rank = rank + 1
|
|
|
|
if pid in positives:
|
|
for depth in qid2recall:
|
|
if rank <= depth:
|
|
qid2recall[depth][qid] = qid2recall[depth].get(qid, 0) + 1.0 / len(positives)
|
|
|
|
assert len(qid2mrr) <= num_ranked_queries, (len(qid2mrr), num_ranked_queries)
|
|
|
|
print()
|
|
mrr_10_sum = sum(qid2mrr.values())
|
|
print_message(f"#> MRR@10 = {mrr_10_sum / num_judged_queries}")
|
|
print_message(f"#> MRR@10 (only for ranked queries) = {mrr_10_sum / num_ranked_queries}")
|
|
print()
|
|
|
|
for depth in qid2recall:
|
|
assert len(qid2recall[depth]) <= num_ranked_queries, (len(qid2recall[depth]), num_ranked_queries)
|
|
|
|
print()
|
|
metric_sum = sum(qid2recall[depth].values())
|
|
print_message(f"#> Recall@{depth} = {metric_sum / num_judged_queries}")
|
|
print_message(f"#> Recall@{depth} (only for ranked queries) = {metric_sum / num_ranked_queries}")
|
|
print()
|
|
|
|
if args.annotate:
|
|
print_message(f"#> Writing annotations to {args.output} ..")
|
|
|
|
with open(args.output, 'w') as f:
|
|
for qid in tqdm.tqdm(qid2positives):
|
|
ranking = qid2ranking[qid]
|
|
positives = qid2positives[qid]
|
|
|
|
for rank, (_, pid, score) in enumerate(ranking):
|
|
rank = rank + 1
|
|
label = int(pid in positives)
|
|
|
|
line = [qid, pid, rank, score, label]
|
|
line = [x for x in line if x is not None]
|
|
line = '\t'.join(map(str, line)) + '\n'
|
|
f.write(line)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = ArgumentParser(description="msmarco_passages.")
|
|
|
|
|
|
parser.add_argument('--qrels', dest='qrels', required=True, type=str)
|
|
parser.add_argument('--ranking', dest='ranking', required=True, type=str)
|
|
parser.add_argument('--annotate', dest='annotate', default=False, action='store_true')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.annotate:
|
|
args.output = f'{args.ranking}.annotated'
|
|
assert not os.path.exists(args.output), args.output
|
|
|
|
main(args)
|
|
|
