| 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) |
|
|
|
|
| |
|
|
