File size: 2,119 Bytes
828992f |
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 |
import os
from contextlib import contextmanager
from colbert.utils.utils import print_message, NullContextManager
from colbert.utils.runs import Run
class RankingLogger():
def __init__(self, directory, qrels=None, log_scores=False):
self.directory = directory
self.qrels = qrels
self.filename, self.also_save_annotations = None, None
self.log_scores = log_scores
@contextmanager
def context(self, filename, also_save_annotations=False):
assert self.filename is None
assert self.also_save_annotations is None
filename = os.path.join(self.directory, filename)
self.filename, self.also_save_annotations = filename, also_save_annotations
print_message("#> Logging ranked lists to {}".format(self.filename))
with open(filename, 'w') as f:
self.f = f
with (open(filename + '.annotated', 'w') if also_save_annotations else NullContextManager()) as g:
self.g = g
try:
yield self
finally:
pass
def log(self, qid, ranking, is_ranked=True, print_positions=[]):
print_positions = set(print_positions)
f_buffer = []
g_buffer = []
for rank, (score, pid, passage) in enumerate(ranking):
is_relevant = self.qrels and int(pid in self.qrels[qid])
rank = rank+1 if is_ranked else -1
possibly_score = [score] if self.log_scores else []
f_buffer.append('\t'.join([str(x) for x in [qid, pid, rank] + possibly_score]) + "\n")
if self.g:
g_buffer.append('\t'.join([str(x) for x in [qid, pid, rank, is_relevant]]) + "\n")
if rank in print_positions:
prefix = "** " if is_relevant else ""
prefix += str(rank)
print("#> ( QID {} ) ".format(qid) + prefix + ") ", pid, ":", score, ' ', passage)
self.f.write(''.join(f_buffer))
if self.g:
self.g.write(''.join(g_buffer))
|