Update rag_eval_metrics.py

rag_eval_metrics.py

@@ -1,57 +1,37 @@
-#!/usr/bin/env python3
-"""
-rag_eval_metrics.py
-
-Evaluate RAG retrieval quality by comparing app logs (JSONL) with a gold file (CSV).
-"""
-
-import argparse
+# ===================== rag_eval_metrics.py (notebook version) =====================
 import json
-import os
 import sys
 from pathlib import Path
 from typing import Dict, List, Tuple, Any, Optional
 
-import pandas as pd
 import numpy as np
+import pandas as pd
 
-
-# ----------------------------- Small Utils ----------------------------- #
+# ---- small CLI-like helpers (you can wrap in main() if you want) ----
+COLOR_TITLE = "\033[94m"
+COLOR_TEXT  = "\033[34m"
+COLOR_ACCENT = "\033[36m"
+COLOR_RESET = "\033[0m"
 
 def filename_key(s: str) -> str:
     s = (s or "").strip().replace("\\", "/").split("/")[-1]
     return s.casefold()
 
-
 def re_split_sc(s: str) -> List[str]:
     import re
     return re.split(r"[;,]", s)
 
-
 def _pick_last_non_empty(hit_lists) -> List[dict]:
-    """
-    Robustly select the last non-empty hits list from a pandas Series or iterable.
-
-    This fixes the KeyError that happens when using reversed() directly on a Series
-    with a non-range index.
-    """
-    # Convert pandas Series or other iterables to a plain Python list
     try:
         values = list(hit_lists.tolist())
     except AttributeError:
         values = list(hit_lists)
-
-    # Walk from last to first, return first non-empty list-like
     for lst in reversed(values):
         if isinstance(lst, (list, tuple)) and len(lst) > 0:
             return lst
-
-    # If everything was empty / NaN
     return []
 
-
-# ----------------------------- IO Helpers ----------------------------- #
-
+# ----------------------------- IO: logs ----------------------------- #
 def read_logs(jsonl_path: Path) -> pd.DataFrame:
     rows = []
     if (not jsonl_path.exists()) or jsonl_path.stat().st_size == 0:
@@ -67,67 +47,96 @@ def read_logs(jsonl_path: Path) -> pd.DataFrame:
             except Exception:
                 continue
 
-            # Extract question
-            q = (((rec.get("inputs") or {}).get("question")) or "").strip()
+            inputs = (rec.get("inputs") or {})
+            q = (inputs.get("question") or "").strip()
+
+            w_tfidf = float(inputs.get("w_tfidf", 0.0))
+            w_bm25  = float(inputs.get("w_bm25", 0.0))
+            w_emb   = float(inputs.get("w_emb", 0.0))
+            top_k   = int(inputs.get("top_k", rec.get("k", 8)))
+            cfg_id  = inputs.get("config_id") or ""
 
-            # Extract retrieval hits (if present)
             retr = (rec.get("retrieval") or {})
             hits = retr.get("hits", [])
+
             norm_hits = []
             for h in hits or []:
                 doc = (h.get("doc") or "").strip()
                 page = str(h.get("page") or "").strip()
-
-                # Normalize page to int or None
                 try:
                     page_int = int(page)
                 except Exception:
                     page_int = None
-
                 norm_hits.append({"doc": doc, "page": page_int})
 
-            rows.append({"question": q, "hits": norm_hits})
+            rows.append(
+                {
+                    "question": q,
+                    "hits": norm_hits,
+                    "w_tfidf": w_tfidf,
+                    "w_bm25": w_bm25,
+                    "w_emb": w_emb,
+                    "top_k": top_k,
+                    "config_id": cfg_id,
+                }
+            )
 
     df = pd.DataFrame(rows)
     if df.empty:
         return pd.DataFrame(columns=["question", "hits"])
 
-    # Group by normalized question text and keep last non-empty hits list per question
+    # group by normalized question + weights + config
     df = (
-        df.groupby(df["question"].astype(str).str.casefold().str.strip(), as_index=False)
-          .agg({"question": "last", "hits": _pick_last_non_empty})
+        df.groupby(
+            [
+                df["question"].astype(str).str.casefold().str.strip(),
+                df["w_tfidf"].round(3),
+                df["w_bm25"].round(3),
+                df["w_emb"].round(3),
+                df["top_k"],
+                df["config_id"],
+            ],
+            as_index=False,
+        )
+        .agg(
+            {
+                "question": "last",
+                "hits": _pick_last_non_empty,
+                "w_tfidf": "last",
+                "w_bm25": "last",
+                "w_emb": "last",
+                "top_k": "last",
+                "config_id": "last",
+            }
+        )
     )
     return df
 
-
+# ----------------------------- IO: gold ----------------------------- #
 def read_gold(csv_path: Path) -> pd.DataFrame:
     df = pd.read_csv(csv_path)
     cols = {c.lower().strip(): c for c in df.columns}
 
-    # --- question column ---
     q_col = None
     for cand in ["question", "query", "q"]:
         if cand in cols:
             q_col = cols[cand]
             break
     if q_col is None:
-        raise ValueError("Gold CSV must contain a 'question' column (case-insensitive).")
+        raise ValueError("Gold CSV must contain a 'question' column.")
 
-    # --- possible relevant_docs (list-in-cell) column ---
     rel_list_col = None
     for cand in ["relevant_docs", "relevant", "docs"]:
         if cand in cols:
             rel_list_col = cols[cand]
             break
 
-    # --- single-doc-per-row column ---
     doc_col = None
     for cand in ["doc", "document", "file", "doc_name"]:
         if cand in cols:
             doc_col = cols[cand]
             break
 
-    # --- optional page column ---
     page_col = None
     for cand in ["page", "page_num", "page_number"]:
         if cand in cols:
@@ -136,71 +145,61 @@ def read_gold(csv_path: Path) -> pd.DataFrame:
 
     rows = []
 
-    # Case 1: relevant_docs list column (no explicit doc_col)
     if rel_list_col and doc_col is None:
         for _, r in df.iterrows():
             q_raw = str(r[q_col]).strip()
             q_norm = q_raw.casefold().strip()
-
             rel_val = str(r[rel_list_col]) if pd.notna(r[rel_list_col]) else ""
             if not rel_val:
-                rows.append({
-                    "question_raw": q_raw,
-                    "question": q_norm,
-                    "doc": None,
-                    "page": np.nan
-                })
+                rows.append(
+                    {
+                        "question_raw": q_raw,
+                        "question": q_norm,
+                        "doc": None,
+                        "page": np.nan,
+                    }
+                )
                 continue
-
             parts = [p.strip() for p in re_split_sc(rel_val)]
             for d in parts:
-                rows.append({
-                    "question_raw": q_raw,
-                    "question": q_norm,
-                    "doc": filename_key(d),
-                    "page": np.nan
-                })
-
-    # Case 2: doc/page columns (one relevant doc per row)
+                rows.append(
+                    {
+                        "question_raw": q_raw,
+                        "question": q_norm,
+                        "doc": filename_key(d),
+                        "page": np.nan,
+                    }
+                )
     elif doc_col:
         for _, r in df.iterrows():
             q_raw = str(r[q_col]).strip()
             q_norm = q_raw.casefold().strip()
-
             d = str(r[doc_col]).strip() if pd.notna(r[doc_col]) else ""
             p = r[page_col] if (page_col and pd.notna(r[page_col])) else np.nan
-
             try:
                 p = int(p)
             except Exception:
                 p = np.nan
-
-            rows.append({
-                "question_raw": q_raw,
-                "question": q_norm,
-                "doc": filename_key(d),
-                "page": p
-            })
-
+            rows.append(
+                {
+                    "question_raw": q_raw,
+                    "question": q_norm,
+                    "doc": filename_key(d),
+                    "page": p,
+                }
+            )
     else:
-        raise ValueError("Gold CSV must contain either a 'doc' column or a 'relevant_docs' column.")
+        raise ValueError("Gold CSV must contain either a 'doc' or 'relevant_docs' column.")
 
     gold = pd.DataFrame(rows)
-
-    # Keep only rows with a valid doc (when docs exist)
     gold["has_doc"] = gold["doc"].apply(lambda x: isinstance(x, str) and len(x) > 0)
     if gold["has_doc"].any():
         gold = gold[gold["has_doc"]].copy()
     gold.drop(columns=["has_doc"], inplace=True, errors="ignore")
-
-    # Remove duplicates
     gold = gold.drop_duplicates(subset=["question", "doc", "page"])
-
     return gold
 
-
-# ----------------------------- Metric Core ----------------------------- #
-
+# ----------------------------- metrics ----------------------------- #
 def dcg_at_k(relevances: List[int]) -> float:
     dcg = 0.0
     for i, rel in enumerate(relevances, start=1):
@@ -208,7 +207,6 @@ def dcg_at_k(relevances: List[int]) -> float:
             dcg += 1.0 / np.log2(i + 1.0)
     return float(dcg)
 
-
 def ndcg_at_k(relevances: List[int]) -> float:
     dcg = dcg_at_k(relevances)
     ideal = sorted(relevances, reverse=True)
@@ -217,13 +215,11 @@ def ndcg_at_k(relevances: List[int]) -> float:
         return 0.0
     return float(dcg / idcg)
 
-
 def compute_metrics_for_question(gold_docs, gold_pages, hits, k):
     top = hits[:k] if hits else []
     pred_docs = [filename_key(h.get("doc", "")) for h in top]
     pred_pairs = [(filename_key(h.get("doc", "")), h.get("page", None)) for h in top]
 
-    # --- Doc-level metrics ---
     gold_doc_set = set([d for d in gold_docs if isinstance(d, str) and d])
 
     rel_bin_doc = [1 if d in gold_doc_set else 0 for d in pred_docs]
@@ -232,10 +228,11 @@ def compute_metrics_for_question(gold_docs, gold_pages, hits, k):
     rec_doc = (sum(rel_bin_doc) / max(1, len(gold_doc_set))) if gold_doc_set else 0.0
     ndcg_doc = ndcg_at_k(rel_bin_doc)
 
-    # --- Page-level metrics (only if gold has page labels) ---
     gold_pairs = set()
     for d, p in zip(gold_docs, gold_pages):
-        if isinstance(d, str) and d and (p is not None) and (not (isinstance(p, float) and np.isnan(p))):
+        if isinstance(d, str) and d and (p is not None) and not (
+            isinstance(p, float) and np.isnan(p)
+        ):
             try:
                 p_int = int(p)
             except Exception:
@@ -249,7 +246,6 @@ def compute_metrics_for_question(gold_docs, gold_pages, hits, k):
                 rel_bin_page.append(0)
             else:
                 rel_bin_page.append(1 if (d, p) in gold_pairs else 0)
-
         hitk_page = 1 if any(rel_bin_page) else 0
         prec_page = (sum(rel_bin_page) / max(1, len(pred_pairs))) if pred_pairs else 0.0
         rec_page = (sum(rel_bin_page) / max(1, len(gold_pairs))) if gold_pairs else 0.0
@@ -268,69 +264,56 @@ def compute_metrics_for_question(gold_docs, gold_pages, hits, k):
         "ndcg@k_page": ndcg_page,
         "n_gold_docs": int(len(gold_doc_set)),
         "n_gold_doc_pages": int(len(gold_pairs)),
-        "n_pred": int(len(pred_docs))
+        "n_pred": int(len(pred_docs)),
     }
 
-
-# ----------------------------- Orchestration ----------------------------- #
-
-# === Dark blue and accent colors ===
-COLOR_TITLE = "\033[94m"     # light blue for titles
-COLOR_TEXT = "\033[34m"      # dark blue
-COLOR_ACCENT = "\033[36m"    # cyan for metrics
-COLOR_RESET = "\033[0m"
-
-
 def _fmt(x: Any) -> str:
     try:
         return f"{float(x):.3f}"
     except Exception:
         return "-"
 
-
-def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--gold_csv", required=True, type=str)
-    ap.add_argument("--logs_jsonl", required=True, type=str)
-    ap.add_argument("--k", type=int, default=8)
-    ap.add_argument("--out_dir", type=str, default="rag_artifacts")
-    args = ap.parse_args()
-
-    out_dir = Path(args.out_dir)
+# ----------------------------- main evaluation ----------------------------- #
+def evaluate_rag(
+    gold_csv: str,
+    logs_jsonl: str,
+    k: int = 8,
+    out_dir: str = "rag_artifacts",
+    group_by_weights: bool = True,
+):
+    out_dir = Path(out_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
 
-    gold_path = Path(args.gold_csv)
-    logs_path = Path(args.logs_jsonl)
+    gold_path = Path(gold_csv)
+    logs_path = Path(logs_jsonl)
 
     if not gold_path.exists():
-        print(f"{COLOR_TEXT}❌ gold.csv not found at {gold_path}{COLOR_RESET}", file=sys.stderr)
-        sys.exit(0)
+        print(f"{COLOR_TEXT}❌ gold.csv not found at {gold_path}{COLOR_RESET}")
+        return
     if not logs_path.exists() or logs_path.stat().st_size == 0:
-        print(f"{COLOR_TEXT}❌ logs JSONL not found or empty at {logs_path}{COLOR_RESET}", file=sys.stderr)
-        sys.exit(0)
+        print(f"{COLOR_TEXT}❌ logs JSONL not found or empty at {logs_path}{COLOR_RESET}")
+        return
 
-    # Read gold
     try:
         gold = read_gold(gold_path)
     except Exception as e:
-        print(f"{COLOR_TEXT}❌ Failed to read gold: {e}{COLOR_RESET}", file=sys.stderr)
-        sys.exit(0)
+        print(f"{COLOR_TEXT}❌ Failed to read gold: {e}{COLOR_RESET}")
+        return
 
-    # Read logs (with robust aggregation)
     try:
         logs = read_logs(logs_path)
     except Exception as e:
-        print(f"{COLOR_TEXT}❌ Failed to read logs: {e}{COLOR_RESET}", file=sys.stderr)
-        sys.exit(0)
+        print(f"{COLOR_TEXT}❌ Failed to read logs: {e}{COLOR_RESET}")
+        return
 
     if gold.empty:
-        print(f"{COLOR_TEXT}❌ Gold file contains no usable rows.{COLOR_RESET}", file=sys.stderr)
-        sys.exit(0)
+        print(f"{COLOR_TEXT}❌ Gold file contains no usable rows.{COLOR_RESET}")
+        return
     if logs.empty:
-        print(f"{COLOR_TEXT}❌ Logs file contains no usable entries.{COLOR_RESET}", file=sys.stderr)
-        sys.exit(0)
+        print(f"{COLOR_TEXT}❌ Logs file contains no usable entries.{COLOR_RESET}")
+        return
 
-    # Build gold dict: normalized_question -> list of (doc, page)
+    # build gold dict
     gdict: Dict[str, List[Tuple[str, Optional[int]]]] = {}
     for _, r in gold.iterrows():
         q = str(r["question"]).strip()
@@ -338,20 +321,18 @@ def main():
         p = r["page"] if "page" in r else np.nan
         gdict.setdefault(q, []).append((d, p))
 
-    # Normalize log questions for join
+    # normalize questions in logs
     logs["q_norm"] = logs["question"].astype(str).str.casefold().str.strip()
 
     perq_rows = []
     not_in_logs, not_in_gold = [], []
 
-    # For each gold question, compute metrics using logs
     for q_norm, pairs in gdict.items():
         row = logs[logs["q_norm"] == q_norm]
         gdocs = [d for (d, _) in pairs]
         gpages = [p for (_, p) in pairs]
 
         if row.empty:
-            # No logs for this gold question → zero retrieval
             not_in_logs.append(q_norm)
             metrics = {
                 "hit@k_doc": 0,
@@ -362,30 +343,27 @@ def main():
                 "precision@k_page": np.nan,
                 "recall@k_page": np.nan,
                 "ndcg@k_page": np.nan,
-                "n_gold_docs": int(len(set([d for d in gdocs if isinstance(d, str) and d]))),
-                "n_gold_doc_pages": int(len([
-                    (d, p) for (d, p) in zip(gdocs, gpages)
-                    if isinstance(d, str) and d and pd.notna(p)
-                ])),
-                "n_pred": 0
+                "n_gold_docs": int(
+                    len(set([d for d in gdocs if isinstance(d, str) and d]))
+                ),
+                "n_gold_doc_pages": int(
+                    len(
+                        [
+                            (d, p)
+                            for (d, p) in zip(gdocs, gpages)
+                            if isinstance(d, str) and d and pd.notna(p)
+                        ]
+                    )
+                ),
+                "n_pred": 0,
             }
-            perq_rows.append({
-                "question": q_norm,
-                "covered_in_logs": 0,
-                **metrics
-            })
+            perq_rows.append({"question": q_norm, "covered_in_logs": 0, **metrics})
             continue
 
-        # Use aggregated hits from read_logs
         hits = row.iloc[0]["hits"] or []
-        metrics = compute_metrics_for_question(gdocs, gpages, hits, args.k)
-        perq_rows.append({
-            "question": q_norm,
-            "covered_in_logs": 1,
-            **metrics
-        })
-
-    # Any log questions not in gold
+        metrics = compute_metrics_for_question(gdocs, gpages, hits, k)
+        perq_rows.append({"question": q_norm, "covered_in_logs": 1, **metrics})
+
     gold_qs = set(gdict.keys())
     for qn in logs["q_norm"].tolist():
         if qn not in gold_qs:
@@ -399,15 +377,23 @@ def main():
         "questions_covered_in_logs": int(covered.shape[0]),
         "questions_missing_in_logs": int(len(not_in_logs)),
         "questions_in_logs_not_in_gold": int(len(set(not_in_gold))),
-        "k": int(args.k),
+        "k": int(k),
         "mean_hit@k_doc": float(covered["hit@k_doc"].mean()) if not covered.empty else 0.0,
         "mean_precision@k_doc": float(covered["precision@k_doc"].mean()) if not covered.empty else 0.0,
         "mean_recall@k_doc": float(covered["recall@k_doc"].mean()) if not covered.empty else 0.0,
         "mean_ndcg@k_doc": float(covered["ndcg@k_doc"].mean()) if not covered.empty else 0.0,
-        "mean_hit@k_page": float(covered["hit@k_page"].dropna().mean()) if covered["hit@k_page"].notna().any() else None,
-        "mean_precision@k_page": float(covered["precision@k_page"].dropna().mean()) if covered["precision@k_page"].notna().any() else None,
-        "mean_recall@k_page": float(covered["recall@k_page"].dropna().mean()) if covered["recall@k_page"].notna().any() else None,
-        "mean_ndcg@k_page": float(covered["ndcg@k_page"].dropna().mean()) if covered["ndcg@k_page"].notna().any() else None,
+        "mean_hit@k_page": float(covered["hit@k_page"].dropna().mean())
+            if covered["hit@k_page"].notna().any()
+            else None,
+        "mean_precision@k_page": float(covered["precision@k_page"].dropna().mean())
+            if covered["precision@k_page"].notna().any()
+            else None,
+        "mean_recall@k_page": float(covered["recall@k_page"].dropna().mean())
+            if covered["recall@k_page"].notna().any()
+            else None,
+        "mean_ndcg@k_page": float(covered["ndcg@k_page"].dropna().mean())
+            if covered["ndcg@k_page"].notna().any()
+            else None,
         "avg_gold_docs_per_q": float(perq["n_gold_docs"].mean()) if not perq.empty else 0.0,
         "avg_preds_per_q": float(perq["n_pred"].mean()) if not perq.empty else 0.0,
         "examples_missing_in_logs": list(not_in_logs[:10]),
@@ -415,13 +401,12 @@ def main():
     }
 
     perq_path = out_dir / "metrics_per_question.csv"
-    agg_path = out_dir / "metrics_aggregate.json"
+    agg_path  = out_dir / "metrics_aggregate.json"
 
     perq.to_csv(perq_path, index=False)
     with open(agg_path, "w", encoding="utf-8") as f:
         json.dump(agg, f, ensure_ascii=False, indent=2)
 
-    # === Console summary with color ===
     print(f"{COLOR_TITLE}RAG Evaluation Summary{COLOR_RESET}")
     print(f"{COLOR_TITLE}----------------------{COLOR_RESET}")
     print(f"{COLOR_TEXT}Gold questions: {COLOR_ACCENT}{agg['questions_total_gold']}{COLOR_RESET}")
@@ -438,7 +423,7 @@ def main():
         f"nDCG@k={_fmt(agg['mean_ndcg@k_doc'])}{COLOR_RESET}"
     )
 
-    if agg['mean_hit@k_page'] is not None:
+    if agg["mean_hit@k_page"] is not None:
         print(
             f"{COLOR_TEXT}Page-level:{COLOR_RESET} "
             f"{COLOR_ACCENT}Hit@k={_fmt(agg['mean_hit@k_page'])}  "
@@ -453,6 +438,48 @@ def main():
     print(f"{COLOR_TEXT}Wrote per-question CSV → {COLOR_ACCENT}{perq_path}{COLOR_RESET}")
     print(f"{COLOR_TEXT}Wrote aggregate JSON   → {COLOR_ACCENT}{agg_path}{COLOR_RESET}")
 
+    # --------- NEW: configuration-wise surface metrics ---------
+    if group_by_weights:
+        logs_short = logs[
+            ["q_norm", "w_tfidf", "w_bm25", "w_emb", "top_k", "config_id"]
+        ].drop_duplicates()
+
+        perq_for_merge = perq.copy()
+        perq_for_merge["q_norm"] = perq_for_merge["question"]
+        perq_for_merge = perq_for_merge.merge(
+            logs_short, on="q_norm", how="left"
+        )
+        perq_for_merge.drop(columns=["q_norm"], inplace=True, errors="ignore")
+
+        gb_cols = ["w_tfidf", "w_bm25", "w_emb", "top_k", "config_id"]
+        surf_rows = []
+        for key, grp in perq_for_merge.groupby(gb_cols):
+            wt, wb, we, tk, cid = key
+            cov = grp[grp["covered_in_logs"] == 1]
+            if cov.empty:
+                continue
+            surf_rows.append(
+                {
+                    "w_tfidf": float(wt),
+                    "w_bm25": float(wb),
+                    "w_emb": float(we),
+                    "top_k": int(tk),
+                    "config_id": cid,
+                    "mean_hit@k_doc": float(cov["hit@k_doc"].mean()),
+                    "mean_precision@k_doc": float(cov["precision@k_doc"].mean()),
+                    "mean_recall@k_doc": float(cov["recall@k_doc"].mean()),
+                    "mean_ndcg@k_doc": float(cov["ndcg@k_doc"].mean()),
+                }
+            )
+
+        surf_df = pd.DataFrame(surf_rows)
+        surf_path = out_dir / "metrics_by_weights.csv"
+        surf_df.to_csv(surf_path, index=False)
+        print(
+            f"{COLOR_TEXT}Wrote config surface metrics → "
+            f"{COLOR_ACCENT}{surf_path}{COLOR_RESET}"
+        )
 
-if __name__ == "__main__":
-    main()
+# Example call from notebook:
+# evaluate_rag("gold.csv", "rag_artifacts/rag_logs.jsonl", k=8, out_dir="rag_artifacts")
+print("✅ rag_eval_metrics helpers loaded.")