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)