| import os
|
| import random
|
| import time
|
| import torch
|
| import torch.nn as nn
|
|
|
| from itertools import accumulate
|
| from math import ceil
|
|
|
| from colbert.utils.runs import Run
|
| from colbert.utils.utils import print_message
|
|
|
| from colbert.evaluation.metrics import Metrics
|
| from colbert.evaluation.ranking_logger import RankingLogger
|
| from colbert.modeling.inference import ModelInference
|
|
|
| from colbert.evaluation.slow import slow_rerank
|
|
|
|
|
| def evaluate(args):
|
| args.inference = ModelInference(args.colbert, amp=args.amp)
|
| qrels, queries, topK_pids = args.qrels, args.queries, args.topK_pids
|
|
|
| depth = args.depth
|
| collection = args.collection
|
| if collection is None:
|
| topK_docs = args.topK_docs
|
|
|
| def qid2passages(qid):
|
| if collection is not None:
|
| return [collection[pid] for pid in topK_pids[qid][:depth]]
|
| else:
|
| return topK_docs[qid][:depth]
|
|
|
| metrics = Metrics(mrr_depths={10, 100}, recall_depths={50, 200, 1000},
|
| success_depths={5, 10, 20, 50, 100, 1000},
|
| total_queries=len(queries))
|
|
|
| ranking_logger = RankingLogger(Run.path, qrels=qrels)
|
|
|
| args.milliseconds = []
|
|
|
| with ranking_logger.context('ranking.tsv', also_save_annotations=(qrels is not None)) as rlogger:
|
| with torch.no_grad():
|
| keys = sorted(list(queries.keys()))
|
| random.shuffle(keys)
|
|
|
| for query_idx, qid in enumerate(keys):
|
| query = queries[qid]
|
|
|
| print_message(query_idx, qid, query, '\n')
|
|
|
| if qrels and args.shortcircuit and len(set.intersection(set(qrels[qid]), set(topK_pids[qid]))) == 0:
|
| continue
|
|
|
| ranking = slow_rerank(args, query, topK_pids[qid], qid2passages(qid))
|
|
|
| rlogger.log(qid, ranking, [0, 1])
|
|
|
| if qrels:
|
| metrics.add(query_idx, qid, ranking, qrels[qid])
|
|
|
| for i, (score, pid, passage) in enumerate(ranking):
|
| if pid in qrels[qid]:
|
| print("\n#> Found", pid, "at position", i+1, "with score", score)
|
| print(passage)
|
| break
|
|
|
| metrics.print_metrics(query_idx)
|
| metrics.log(query_idx)
|
|
|
| print_message("#> checkpoint['batch'] =", args.checkpoint['batch'], '\n')
|
| print("rlogger.filename =", rlogger.filename)
|
|
|
| if len(args.milliseconds) > 1:
|
| print('Slow-Ranking Avg Latency =', sum(args.milliseconds[1:]) / len(args.milliseconds[1:]))
|
|
|
| print("\n\n")
|
|
|
| print("\n\n")
|
|
|
| print("\n\n")
|
|
|
| print('\n\n')
|
| if qrels:
|
| assert query_idx + 1 == len(keys) == len(set(keys))
|
| metrics.output_final_metrics(os.path.join(Run.path, 'ranking.metrics'), query_idx, len(queries))
|
| print('\n\n')
|
|
|
