|
import ujson
|
|
|
|
from collections import defaultdict
|
|
from colbert.utils.runs import Run
|
|
|
|
|
|
class Metrics:
|
|
def __init__(self, mrr_depths: set, recall_depths: set, success_depths: set, total_queries=None):
|
|
self.results = {}
|
|
self.mrr_sums = {depth: 0.0 for depth in mrr_depths}
|
|
self.recall_sums = {depth: 0.0 for depth in recall_depths}
|
|
self.success_sums = {depth: 0.0 for depth in success_depths}
|
|
self.total_queries = total_queries
|
|
|
|
self.max_query_idx = -1
|
|
self.num_queries_added = 0
|
|
|
|
def add(self, query_idx, query_key, ranking, gold_positives):
|
|
self.num_queries_added += 1
|
|
|
|
assert query_key not in self.results
|
|
assert len(self.results) <= query_idx
|
|
assert len(set(gold_positives)) == len(gold_positives)
|
|
assert len(set([pid for _, pid, _ in ranking])) == len(ranking)
|
|
|
|
self.results[query_key] = ranking
|
|
|
|
positives = [i for i, (_, pid, _) in enumerate(ranking) if pid in gold_positives]
|
|
|
|
if len(positives) == 0:
|
|
return
|
|
|
|
for depth in self.mrr_sums:
|
|
first_positive = positives[0]
|
|
self.mrr_sums[depth] += (1.0 / (first_positive+1.0)) if first_positive < depth else 0.0
|
|
|
|
for depth in self.success_sums:
|
|
first_positive = positives[0]
|
|
self.success_sums[depth] += 1.0 if first_positive < depth else 0.0
|
|
|
|
for depth in self.recall_sums:
|
|
num_positives_up_to_depth = len([pos for pos in positives if pos < depth])
|
|
self.recall_sums[depth] += num_positives_up_to_depth / len(gold_positives)
|
|
|
|
def print_metrics(self, query_idx):
|
|
for depth in sorted(self.mrr_sums):
|
|
print("MRR@" + str(depth), "=", self.mrr_sums[depth] / (query_idx+1.0))
|
|
|
|
for depth in sorted(self.success_sums):
|
|
print("Success@" + str(depth), "=", self.success_sums[depth] / (query_idx+1.0))
|
|
|
|
for depth in sorted(self.recall_sums):
|
|
print("Recall@" + str(depth), "=", self.recall_sums[depth] / (query_idx+1.0))
|
|
|
|
def log(self, query_idx):
|
|
assert query_idx >= self.max_query_idx
|
|
self.max_query_idx = query_idx
|
|
|
|
Run.log_metric("ranking/max_query_idx", query_idx, query_idx)
|
|
Run.log_metric("ranking/num_queries_added", self.num_queries_added, query_idx)
|
|
|
|
for depth in sorted(self.mrr_sums):
|
|
score = self.mrr_sums[depth] / (query_idx+1.0)
|
|
Run.log_metric("ranking/MRR." + str(depth), score, query_idx)
|
|
|
|
for depth in sorted(self.success_sums):
|
|
score = self.success_sums[depth] / (query_idx+1.0)
|
|
Run.log_metric("ranking/Success." + str(depth), score, query_idx)
|
|
|
|
for depth in sorted(self.recall_sums):
|
|
score = self.recall_sums[depth] / (query_idx+1.0)
|
|
Run.log_metric("ranking/Recall." + str(depth), score, query_idx)
|
|
|
|
def output_final_metrics(self, path, query_idx, num_queries):
|
|
assert query_idx + 1 == num_queries
|
|
assert num_queries == self.total_queries
|
|
|
|
if self.max_query_idx < query_idx:
|
|
self.log(query_idx)
|
|
|
|
self.print_metrics(query_idx)
|
|
|
|
output = defaultdict(dict)
|
|
|
|
for depth in sorted(self.mrr_sums):
|
|
score = self.mrr_sums[depth] / (query_idx+1.0)
|
|
output['mrr'][depth] = score
|
|
|
|
for depth in sorted(self.success_sums):
|
|
score = self.success_sums[depth] / (query_idx+1.0)
|
|
output['success'][depth] = score
|
|
|
|
for depth in sorted(self.recall_sums):
|
|
score = self.recall_sums[depth] / (query_idx+1.0)
|
|
output['recall'][depth] = score
|
|
|
|
with open(path, 'w') as f:
|
|
ujson.dump(output, f, indent=4)
|
|
f.write('\n')
|
|
|
|
|
|
def evaluate_recall(qrels, queries, topK_pids):
|
|
if qrels is None:
|
|
return
|
|
|
|
assert set(qrels.keys()) == set(queries.keys())
|
|
recall_at_k = [len(set.intersection(set(qrels[qid]), set(topK_pids[qid]))) / max(1.0, len(qrels[qid]))
|
|
for qid in qrels]
|
|
recall_at_k = sum(recall_at_k) / len(qrels)
|
|
recall_at_k = round(recall_at_k, 3)
|
|
print("Recall @ maximum depth =", recall_at_k)
|
|
|
|
|
|
|
|
|
