src/evaluation.py

import numpy as np

def mrr_at_k(ranked_doc_indices, relevant_doc_ids, doc_ids_map, k=10):
    for rank, doc_i in enumerate(ranked_doc_indices[:k], 1):
        if doc_ids_map[doc_i] in relevant_doc_ids:
            return 1.0 / rank
    return 0.0

def ndcg_at_k(ranked_doc_indices, relevant_doc_ids, doc_ids_map, k=10):
    dcg = sum(1.0 / np.log2(r + 2) for r, doc_i in enumerate(ranked_doc_indices[:k]) 
              if doc_ids_map[doc_i] in relevant_doc_ids)
    idcg = sum(1.0 / np.log2(r + 2) for r in range(min(len(relevant_doc_ids), k)))
    return dcg / idcg if idcg > 0 else 0.0

def recall_at_k(ranked_doc_indices, relevant_doc_ids, doc_ids_map, k=1000):
    hits = sum(1 for doc_i in ranked_doc_indices[:k] if doc_ids_map[doc_i] in relevant_doc_ids)
    return hits / max(len(relevant_doc_ids), 1)