Update rag_eval_metrics.py

rag_eval_metrics.py

@@ -16,10 +16,47 @@ import pandas as pd
 import numpy as np
 
 
+# ----------------------------- Small Utils ----------------------------- #
+
+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 ----------------------------- #
 
 def read_logs(jsonl_path: Path) -> pd.DataFrame:
     rows = []
+    if (not jsonl_path.exists()) or jsonl_path.stat().st_size == 0:
+        return pd.DataFrame(columns=["question", "hits"])
+
     with open(jsonl_path, "r", encoding="utf-8") as f:
         for line in f:
             line = line.strip()
@@ -29,29 +66,35 @@ def read_logs(jsonl_path: Path) -> pd.DataFrame:
                 rec = json.loads(line)
             except Exception:
                 continue
+
+            # Extract question
             q = (((rec.get("inputs") or {}).get("question")) or "").strip()
+
+            # 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})
+
     df = pd.DataFrame(rows)
     if df.empty:
         return pd.DataFrame(columns=["question", "hits"])
-    def _pick_last_non_empty(hit_lists: List[List[dict]]) -> List[dict]:
-        for lst in reversed(hit_lists):
-            if lst:
-                return lst
-        return []
+
+    # Group by normalized question text and keep last non-empty hits list per question
     df = (
-        df.groupby(df["question"].str.casefold().str.strip(), as_index=False)
+        df.groupby(df["question"].astype(str).str.casefold().str.strip(), as_index=False)
           .agg({"question": "last", "hits": _pick_last_non_empty})
     )
     return df
@@ -60,6 +103,8 @@ def read_logs(jsonl_path: Path) -> pd.DataFrame:
 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:
@@ -68,18 +113,21 @@ def read_gold(csv_path: Path) -> pd.DataFrame:
     if q_col is None:
         raise ValueError("Gold CSV must contain a 'question' column (case-insensitive).")
 
+    # --- 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:
@@ -87,48 +135,68 @@ def read_gold(csv_path: Path) -> pd.DataFrame:
             break
 
     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})
+                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)
     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
+            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.")
 
     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")
-    gold = gold.drop_duplicates(subset=["question", "doc", "page"])
-    return gold
-
-
-def filename_key(s: str) -> str:
-    s = (s or "").strip().replace("\\", "/").split("/")[-1]
-    return s.casefold()
 
+    # Remove duplicates
+    gold = gold.drop_duplicates(subset=["question", "doc", "page"])
 
-def re_split_sc(s: str) -> List[str]:
-    import re
-    return re.split(r"[;,]", s)
+    return gold
 
 
 # ----------------------------- Metric Core ----------------------------- #
@@ -155,13 +223,16 @@ def compute_metrics_for_question(gold_docs, gold_pages, hits, k):
     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]
     hitk_doc = 1 if any(rel_bin_doc) else 0
     prec_doc = (sum(rel_bin_doc) / max(1, len(pred_docs))) if pred_docs else 0.0
     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))):
@@ -172,8 +243,13 @@ def compute_metrics_for_question(gold_docs, gold_pages, hits, k):
             gold_pairs.add((d, p_int))
 
     if gold_pairs:
-        rel_bin_page = [1 if ((d, (p if p is not None else -1)) in gold_pairs) else 0
-                        for (d, p) in [(d, (p if isinstance(p, int) else -1)) for (d, p) in pred_pairs]]
+        rel_bin_page = []
+        for (d, p) in pred_pairs:
+            if p is None or not isinstance(p, int):
+                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
@@ -204,6 +280,14 @@ 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)
@@ -225,12 +309,19 @@ def main():
         print(f"{COLOR_TEXT}❌ logs JSONL not found or empty at {logs_path}{COLOR_RESET}", file=sys.stderr)
         sys.exit(0)
 
+    # 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)
-    logs = read_logs(logs_path)
+
+    # 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)
 
     if gold.empty:
         print(f"{COLOR_TEXT}❌ Gold file contains no usable rows.{COLOR_RESET}", file=sys.stderr)
@@ -239,6 +330,7 @@ def main():
         print(f"{COLOR_TEXT}❌ Logs file contains no usable entries.{COLOR_RESET}", file=sys.stderr)
         sys.exit(0)
 
+    # Build gold dict: normalized_question -> list of (doc, page)
     gdict: Dict[str, List[Tuple[str, Optional[int]]]] = {}
     for _, r in gold.iterrows():
         q = str(r["question"]).strip()
@@ -246,32 +338,54 @@ def main():
         p = r["page"] if "page" in r else np.nan
         gdict.setdefault(q, []).append((d, p))
 
+    # Normalize log questions for join
     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)
-            gdocs = [d for (d, _) in pairs]
-            gpages = [p for (_, p) in pairs]
             metrics = {
-                "hit@k_doc": 0, "precision@k_doc": 0.0, "recall@k_doc": 0.0, "ndcg@k_doc": 0.0,
-                "hit@k_page": np.nan, "precision@k_page": np.nan, "recall@k_page": np.nan, "ndcg@k_page": np.nan,
+                "hit@k_doc": 0,
+                "precision@k_doc": 0.0,
+                "recall@k_doc": 0.0,
+                "ndcg@k_doc": 0.0,
+                "hit@k_page": np.nan,
+                "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_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 []
-        gdocs = [d for (d, _) in pairs]
-        gpages = [p for (_, p) in pairs]
         metrics = compute_metrics_for_question(gdocs, gpages, hits, args.k)
-        perq_rows.append({"question": q_norm, "covered_in_logs": 1, **metrics})
+        perq_rows.append({
+            "question": q_norm,
+            "covered_in_logs": 1,
+            **metrics
+        })
 
+    # Any log questions not in gold
     gold_qs = set(gdict.keys())
     for qn in logs["q_norm"].tolist():
         if qn not in gold_qs:
@@ -279,6 +393,7 @@ def main():
 
     perq = pd.DataFrame(perq_rows)
     covered = perq[perq["covered_in_logs"] == 1].copy()
+
     agg = {
         "questions_total_gold": int(len(gdict)),
         "questions_covered_in_logs": int(covered.shape[0]),
@@ -301,6 +416,7 @@ def main():
 
     perq_path = out_dir / "metrics_per_question.csv"
     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)
@@ -314,18 +430,22 @@ def main():
     print(f"{COLOR_TEXT}In logs but not in gold: {COLOR_ACCENT}{agg['questions_in_logs_not_in_gold']}{COLOR_RESET}")
     print(f"{COLOR_TEXT}k = {COLOR_ACCENT}{agg['k']}{COLOR_RESET}\n")
 
-    print(f"{COLOR_TEXT}Doc-level:{COLOR_RESET}  "
-          f"{COLOR_ACCENT}Hit@k={_fmt(agg['mean_hit@k_doc'])}  "
-          f"Precision@k={_fmt(agg['mean_precision@k_doc'])}  "
-          f"Recall@k={_fmt(agg['mean_recall@k_doc'])}  "
-          f"nDCG@k={_fmt(agg['mean_ndcg@k_doc'])}{COLOR_RESET}")
+    print(
+        f"{COLOR_TEXT}Doc-level:{COLOR_RESET}  "
+        f"{COLOR_ACCENT}Hit@k={_fmt(agg['mean_hit@k_doc'])}  "
+        f"Precision@k={_fmt(agg['mean_precision@k_doc'])}  "
+        f"Recall@k={_fmt(agg['mean_recall@k_doc'])}  "
+        f"nDCG@k={_fmt(agg['mean_ndcg@k_doc'])}{COLOR_RESET}"
+    )
 
     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'])}  "
-              f"Precision@k={_fmt(agg['mean_precision@k_page'])}  "
-              f"Recall@k={_fmt(agg['mean_recall@k_page'])}  "
-              f"nDCG@k={_fmt(agg['mean_ndcg@k_page'])}{COLOR_RESET}")
+        print(
+            f"{COLOR_TEXT}Page-level:{COLOR_RESET} "
+            f"{COLOR_ACCENT}Hit@k={_fmt(agg['mean_hit@k_page'])}  "
+            f"Precision@k={_fmt(agg['mean_precision@k_page'])}  "
+            f"Recall@k={_fmt(agg['mean_recall@k_page'])}  "
+            f"nDCG@k={_fmt(agg['mean_ndcg@k_page'])}{COLOR_RESET}"
+        )
     else:
         print(f"{COLOR_TEXT}Page-level: (no page labels in gold){COLOR_RESET}")
 
@@ -334,12 +454,5 @@ def main():
     print(f"{COLOR_TEXT}Wrote aggregate JSON   → {COLOR_ACCENT}{agg_path}{COLOR_RESET}")
 
 
-def _fmt(x: Any) -> str:
-    try:
-        return f"{float(x):.3f}"
-    except Exception:
-        return "-"
-
-
 if __name__ == "__main__":
     main()