""" Evaluate retrieval performance (recall@K, hit-rate) for the RAG engine. Expected ground-truth file format (JSONL or JSON list): - Each item is an object with keys: "query" (str) and "relevant_ids" (list[int] or list[str]) - `relevant_ids` should match the index positions in `data/chunks.pkl` or a stable id present in the chunk dicts. Usage: python evaluate_retrieval.py --gt ground_truth.jsonl --k 10 If your chunk metadata uses custom IDs (e.g., a `chunk_id` field), pass `--id-field chunk_id`. """ from __future__ import annotations import argparse import json import os import pickle import logging from typing import List, Any import math from collections import defaultdict from rag_engine import RAGEngine from config import CHUNKS_PATH logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) def load_ground_truth(path: str) -> List[dict]: if not os.path.exists(path): raise FileNotFoundError(path) with open(path, "r", encoding="utf-8") as f: txt = f.read().strip() if not txt: return [] try: data = json.loads(txt) except json.JSONDecodeError: # fallback to JSONL data = [json.loads(l) for l in txt.splitlines() if l.strip()] return data def map_chunk_ids(chunks: List[dict], id_field: str = None) -> dict: """Return a mapping from id -> index in chunks list. If `id_field` is None, uses the integer index as the id. """ mapping = {} if id_field: for i, c in enumerate(chunks): if id_field in c: mapping[c[id_field]] = i else: for i, c in enumerate(chunks): mapping[i] = i return mapping def recall_at_k(rag: RAGEngine, gt: List[dict], k: int = 10, id_field: str | None = None) -> dict: """Compute Recall@K, Precision@K, MRR, and nDCG@K with optional per-language breakdown. Ground-truth entries may include a `lang` field (e.g. 'en' or 'hinglish'). Returns a stats dict with overall metrics and a `per_language` mapping. """ # Ensure engine initialized rag.initialize() # Load chunks and build id mapping with open(CHUNKS_PATH, "rb") as f: chunks = pickle.load(f) id2idx = map_chunk_ids(chunks, id_field) total = 0 hits = 0 sum_precision = 0.0 sum_rr = 0.0 sum_ndcg = 0.0 per_language = defaultdict(lambda: {"total": 0, "hits": 0, "sum_precision": 0.0, "sum_rr": 0.0, "sum_ndcg": 0.0}) for item in gt: query = item.get("query") or item.get("q") relevant = item.get("relevant_ids") or item.get("relevant") lang = item.get("lang", "en") if not query or not relevant: continue total += 1 per_language[lang]["total"] += 1 # Map relevant ids to indices if mapping available relevant_idxs = set() for r in relevant: if id_field and r in id2idx: relevant_idxs.add(id2idx[r]) else: try: relevant_idxs.add(int(r)) except Exception: pass # Retrieve ordered results results = rag.retrieve(query, top_k=k) retrieved_order = [] for r in results: idx = None # Prefer explicit id field when available if id_field and isinstance(r, dict) and id_field in r and r[id_field] in id2idx: idx = id2idx[r[id_field]] # Match by text content if present (handles copied dicts with score) if idx is None and isinstance(r, dict) and "text" in r: r_text = r.get("text") for i, c in enumerate(chunks): try: c_text = c.get("text") if isinstance(c, dict) else str(c) except Exception: c_text = str(c) if c_text == r_text: idx = i break # Last-resort: try direct identity/index lookup if idx is None: try: idx = chunks.index(r) except Exception: idx = None if idx is not None and idx not in retrieved_order: retrieved_order.append(idx) # Metrics for this query top_k_list = retrieved_order[:k] hits_k = sum(1 for idx in top_k_list if idx in relevant_idxs) precision_k = hits_k / k if k > 0 else 0.0 hit_flag = hits_k > 0 if hit_flag: hits += 1 per_language[lang]["hits"] += 1 # Reciprocal rank rr = 0.0 for pos, idx in enumerate(retrieved_order, start=1): if idx in relevant_idxs: rr = 1.0 / pos break # nDCG@K (binary relevance) dcg = 0.0 for pos, idx in enumerate(top_k_list, start=1): if idx in relevant_idxs: dcg += 1.0 / math.log2(pos + 1) ideal_rel = min(len(relevant_idxs), k) idcg = sum(1.0 / math.log2(i + 1) for i in range(1, ideal_rel + 1)) if ideal_rel > 0 else 0.0 ndcg = dcg / idcg if idcg > 0 else 0.0 sum_precision += precision_k sum_rr += rr sum_ndcg += ndcg per_language[lang]["sum_precision"] += precision_k per_language[lang]["sum_rr"] += rr per_language[lang]["sum_ndcg"] += ndcg # Aggregate recall = hits / total if total else 0.0 precision_at_k = sum_precision / total if total else 0.0 mrr = sum_rr / total if total else 0.0 ndcg = sum_ndcg / total if total else 0.0 # Per-language averages per_lang_stats = {} for lang, vals in per_language.items(): t = vals["total"] per_lang_stats[lang] = { "total": t, "recall_at_k": (vals["hits"] / t) if t else 0.0, "precision_at_k": (vals["sum_precision"] / t) if t else 0.0, "mrr": (vals["sum_rr"] / t) if t else 0.0, "ndcg": (vals["sum_ndcg"] / t) if t else 0.0, } return { "recall_at_k": recall, "precision_at_k": precision_at_k, "mrr": mrr, "ndcg": ndcg, "total_queries": total, "hits": hits, "per_language": per_lang_stats, } def main(): parser = argparse.ArgumentParser() parser.add_argument("--gt", required=True, help="Ground truth JSON/JSONL file") parser.add_argument("--k", type=int, default=10, help="Top-K to evaluate") parser.add_argument("--id-field", help="Field name in chunk dicts used as stable id") args = parser.parse_args() gt = load_ground_truth(args.gt) rag = RAGEngine() stats = recall_at_k(rag, gt, args.k, args.id_field) logger.info(f"Recall@{args.k}: {stats['recall_at_k']:.4f} ({stats['hits']}/{stats['total_queries']})") if __name__ == "__main__": main()