Update rag_eval_metrics.py

rag_eval_metrics.py

@@ -3,30 +3,6 @@
 rag_eval_metrics.py
 
 Evaluate RAG retrieval quality by comparing app logs (JSONL) with a gold file (CSV).
-
-Inputs (CLI):
-  --gold_csv       Path to gold CSV.
-  --logs_jsonl     Path to app JSONL logs (rag_logs.jsonl).
-  --k              Cutoff for metrics (default: 8).
-  --out_dir        Output directory for metrics files (default: rag_artifacts).
-
-Outputs (written into out_dir):
-  - metrics_per_question.csv
-  - metrics_aggregate.json
-
-Gold CSV accepted schemas (case-insensitive headers):
-  Minimal (doc-level):
-    question, doc
-    (multiple rows per question allowed)
-  With page info (page-level optional):
-    question, doc, page
-  List-in-a-cell also supported:
-    question, relevant_docs   # semicolon/comma separated; page matching disabled in this column
-
-Notes:
-  - Matching is case-insensitive on question and doc filename.
-  - Page-level metrics only computed when GOLD includes a concrete page for that question.
-  - Logs are produced by app.py and contain 'retrieval'->'hits' with 'doc' and 'page'.
 """
 
 import argparse
@@ -43,7 +19,6 @@ import numpy as np
 # ----------------------------- IO Helpers ----------------------------- #
 
 def read_logs(jsonl_path: Path) -> pd.DataFrame:
-    """Read JSONL logs and return a DataFrame with columns: question, hits(list[dict])."""
     rows = []
     with open(jsonl_path, "r", encoding="utf-8") as f:
         for line in f:
@@ -57,13 +32,11 @@ def read_logs(jsonl_path: Path) -> pd.DataFrame:
             q = (((rec.get("inputs") or {}).get("question")) or "").strip()
             retr = (rec.get("retrieval") or {})
             hits = retr.get("hits", [])
-            # Normalize fields we need
             norm_hits = []
             for h in hits or []:
                 doc = (h.get("doc") or "").strip()
                 page = str(h.get("page") or "").strip()
                 try:
-                    # Try int page if it looks numeric
                     page_int = int(page)
                 except Exception:
                     page_int = None
@@ -72,8 +45,6 @@ def read_logs(jsonl_path: Path) -> pd.DataFrame:
     df = pd.DataFrame(rows)
     if df.empty:
         return pd.DataFrame(columns=["question", "hits"])
-    # Keep last occurrence per question (latest run), but also allow multiple – we aggregate by question
-    # For stability, group and keep the last non-empty hit list.
     def _pick_last_non_empty(hit_lists: List[List[dict]]) -> List[dict]:
         for lst in reversed(hit_lists):
             if lst:
@@ -87,13 +58,8 @@ def read_logs(jsonl_path: Path) -> pd.DataFrame:
 
 
 def read_gold(csv_path: Path) -> pd.DataFrame:
-    """Read gold CSV, normalize columns, and return rows with:
-       question(cf), question_raw, doc (lowercased filename), page (optional, int or NaN).
-    """
     df = pd.read_csv(csv_path)
-    # Normalize headers
     cols = {c.lower().strip(): c for c in df.columns}
-    # Find question column
     q_col = None
     for cand in ["question", "query", "q"]:
         if cand in cols:
@@ -102,7 +68,6 @@ def read_gold(csv_path: Path) -> pd.DataFrame:
     if q_col is None:
         raise ValueError("Gold CSV must contain a 'question' column (case-insensitive).")
 
-    # Accept either (doc[, page]) rows or a 'relevant_docs' list column
     rel_list_col = None
     for cand in ["relevant_docs", "relevant", "docs"]:
         if cand in cols:
@@ -123,22 +88,17 @@ def read_gold(csv_path: Path) -> pd.DataFrame:
 
     rows = []
     if rel_list_col and doc_col is None:
-        # Each row may contain a list of docs (comma/semicolon separated)
         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:
-                # still create an empty row (no gold docs)
                 rows.append({"question_raw": q_raw, "question": q_norm, "doc": None, "page": np.nan})
                 continue
-            # split by semicolon or comma
             parts = [p.strip() for p in re_split_sc(rel_val)]
-            # one row per doc (page-level off for list column)
             for d in parts:
                 rows.append({"question_raw": q_raw, "question": q_norm, "doc": filename_key(d), "page": np.nan})
     elif doc_col:
-        # Standard long form: one doc (+/- page) per row
         for _, r in df.iterrows():
             q_raw = str(r[q_col]).strip()
             q_norm = q_raw.casefold().strip()
@@ -153,26 +113,20 @@ def read_gold(csv_path: Path) -> pd.DataFrame:
         raise ValueError("Gold CSV must contain either a 'doc' column or a 'relevant_docs' column.")
 
     gold = pd.DataFrame(rows)
-    # drop fully empty doc rows (when no gold docs listed)
     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")
-    # Deduplicate
     gold = gold.drop_duplicates(subset=["question", "doc", "page"])
     return gold
 
 
 def filename_key(s: str) -> str:
-    """Normalize document name to just the basename, lowercased."""
-    s = (s or "").strip()
-    s = s.replace("\\", "/")
-    s = s.split("/")[-1]
+    s = (s or "").strip().replace("\\", "/").split("/")[-1]
     return s.casefold()
 
 
 def re_split_sc(s: str) -> List[str]:
-    """Split on semicolons or commas."""
     import re
     return re.split(r"[;,]", s)
 
@@ -180,7 +134,6 @@ def re_split_sc(s: str) -> List[str]:
 # ----------------------------- Metric Core ----------------------------- #
 
 def dcg_at_k(relevances: List[int]) -> float:
-    """Binary DCG with log2 discounts; ranks are 1-indexed in denominator."""
     dcg = 0.0
     for i, rel in enumerate(relevances, start=1):
         if rel > 0:
@@ -197,22 +150,11 @@ def ndcg_at_k(relevances: List[int]) -> float:
     return float(dcg / idcg)
 
 
-def compute_metrics_for_question(
-    gold_docs: List[str],
-    gold_pages: List[Optional[int]],
-    hits: List[Dict[str, Any]],
-    k: int
-) -> Dict[str, Any]:
-    """
-    Returns per-question metrics at cutoff k for:
-      - doc-level: match on doc only
-      - page-level: match on (doc,page) where page is provided in GOLD
-    """
+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 ---
     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
@@ -220,7 +162,6 @@ def compute_metrics_for_question(
     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 (only if at least one GOLD page specified) ---
     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))):
@@ -257,6 +198,12 @@ def compute_metrics_for_question(
 
 # ----------------------------- 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 main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--gold_csv", required=True, type=str)
@@ -272,28 +219,26 @@ def main():
     logs_path = Path(args.logs_jsonl)
 
     if not gold_path.exists():
-        print(f"❌ gold.csv not found at {gold_path}", file=sys.stderr)
+        print(f"{COLOR_TEXT}❌ gold.csv not found at {gold_path}{COLOR_RESET}", file=sys.stderr)
         sys.exit(0)
     if not logs_path.exists() or logs_path.stat().st_size == 0:
-        print(f"❌ logs JSONL not found or empty at {logs_path}", file=sys.stderr)
+        print(f"{COLOR_TEXT}❌ logs JSONL not found or empty at {logs_path}{COLOR_RESET}", file=sys.stderr)
         sys.exit(0)
 
-    # Load data
     try:
         gold = read_gold(gold_path)
     except Exception as e:
-        print(f"❌ Failed to read gold: {e}", file=sys.stderr)
+        print(f"{COLOR_TEXT}❌ Failed to read gold: {e}{COLOR_RESET}", file=sys.stderr)
         sys.exit(0)
     logs = read_logs(logs_path)
 
     if gold.empty:
-        print("❌ Gold file contains no usable rows.", file=sys.stderr)
+        print(f"{COLOR_TEXT}❌ Gold file contains no usable rows.{COLOR_RESET}", file=sys.stderr)
         sys.exit(0)
     if logs.empty:
-        print("❌ Logs file contains no usable entries.", file=sys.stderr)
+        print(f"{COLOR_TEXT}❌ Logs file contains no usable entries.{COLOR_RESET}", file=sys.stderr)
         sys.exit(0)
 
-    # Build gold dict: question -> list of (doc, page)
     gdict: Dict[str, List[Tuple[str, Optional[int]]]] = {}
     for _, r in gold.iterrows():
         q = str(r["question"]).strip()
@@ -301,19 +246,14 @@ def main():
         p = r["page"] if "page" in r else np.nan
         gdict.setdefault(q, []).append((d, p))
 
-    # Align on questions (casefolded)
     logs["q_norm"] = logs["question"].astype(str).str.casefold().str.strip()
     perq_rows = []
     not_in_logs, not_in_gold = [], []
 
     for q_norm, pairs in gdict.items():
-        # Pairs is list of (doc, page)
-        q_gold_variants = [q_norm]  # already normalized
-        # Find logs row with same normalized question
         row = logs[logs["q_norm"] == q_norm]
         if row.empty:
             not_in_logs.append(q_norm)
-            # Still record a row with zeros/NaNs
             gdocs = [d for (d, _) in pairs]
             gpages = [p for (_, p) in pairs]
             metrics = {
@@ -323,35 +263,21 @@ def main():
                 "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 the last row (grouping ensured one row per question)
         hits = row.iloc[0]["hits"] or []
-        # Prepare gold lists for metric function
         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
-        })
-
-    # Detect questions present in logs but not in gold (for reporting)
     gold_qs = set(gdict.keys())
     for qn in logs["q_norm"].tolist():
         if qn not in gold_qs:
             not_in_gold.append(qn)
 
     perq = pd.DataFrame(perq_rows)
-
-    # Aggregates over questions that are covered_in_logs == 1
     covered = perq[perq["covered_in_logs"] == 1].copy()
     agg = {
         "questions_total_gold": int(len(gdict)),
@@ -359,55 +285,53 @@ def main():
         "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),
-        # DOC-level
         "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,
-        # PAGE-level (skip NaNs)
         "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,
-        # Distribution hints
         "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,
-        # Listings (truncated for readability)
         "examples_missing_in_logs": list(not_in_logs[:10]),
         "examples_in_logs_not_in_gold": list(dict.fromkeys(not_in_gold))[:10],
     }
 
-    # Write outputs
     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)
 
-    # Console summary (stdout) for app display
-    print("RAG Evaluation Summary")
-    print("----------------------")
-    print(f"Gold questions: {agg['questions_total_gold']}")
-    print(f"Covered in logs: {agg['questions_covered_in_logs']}")
-    print(f"Missing in logs: {agg['questions_missing_in_logs']}")
-    print(f"In logs but not in gold: {agg['questions_in_logs_not_in_gold']}")
-    print(f"k = {agg['k']}")
-    print()
-    print(f"Doc-level:  Hit@k={_fmt(agg['mean_hit@k_doc'])}  "
+    # === 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}")
+    print(f"{COLOR_TEXT}Covered in logs: {COLOR_ACCENT}{agg['questions_covered_in_logs']}{COLOR_RESET}")
+    print(f"{COLOR_TEXT}Missing in logs: {COLOR_ACCENT}{agg['questions_missing_in_logs']}{COLOR_RESET}")
+    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'])}")
-    if agg["mean_hit@k_page"] is not None:
-        print(f"Page-level: Hit@k={_fmt(agg['mean_hit@k_page'])}  "
+          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'])}")
+              f"nDCG@k={_fmt(agg['mean_ndcg@k_page'])}{COLOR_RESET}")
     else:
-        print("Page-level: (no page labels in gold)")
+        print(f"{COLOR_TEXT}Page-level: (no page labels in gold){COLOR_RESET}")
 
     print()
-    print(f"Wrote per-question CSV → {perq_path}")
-    print(f"Wrote aggregate JSON   → {agg_path}")
+    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}")
 
 
 def _fmt(x: Any) -> str: