"""RAG evaluation: compare retrieval quality across chunking methods. Chunking methods (internal): (a) Graph-based -- chunks from chunk_builder.extract_chunks (b) Naive row-based -- one chunk per spreadsheet row (c) Single-cell -- each non-empty cell = one chunk (no context) Oracle methods (no GNN; chunk from whatever labels exist in the per-sheet JSON): graph_gold, graph_row_gold, graph_kg_gold — same queries and metrics as model runs. **Drive + service account:** Structure JSONs live under ``rag_eval/labeled_structure/``. For batch eval, run ``python eval_rag.py --pull-drive`` first (or use the app) so files mirror to ``data/rag_eval_labeled/``. Share the ``rag_eval`` folder with the labeler service account email (read-only is enough for download). **When this is a true "human structure" upper bound:** sheet JSON must contain real table-structure labels (``header``, ``value``, ``metadata``, ``table`` ids, etc.). The usual ``rag_eval_cache`` from raw xlsx download is often **unlabeled**, so oracle then does *not* represent gold structure—only "chunk without running the GNN." To measure oracle vs models, use labeled exports for the same ``sheet_id`` keys as your RAG eval held-out sheets (replace or augment cache files, or use ``data/labeled/``). External baseline methods (via rag_baselines.py): pandas, BeautifulSoup, lxml, unstructured, LlamaSheets, Calamari Evaluation dimensions: 1. Retrieval quality: Recall@K, MRR (answer value in retrieved chunks) 2. Source cell accuracy: Source Recall@K, Source MRR (correct cell retrieved) 3. LLM-as-a-judge: answer generation + correctness scoring 4. (Learned graph methods) Structure vs human labels: levels 1–4 via ``evaluate_sheet``, saved to ``data/rag_eval_structure_metrics.json`` (see ``rag_structure_rag_analysis.py`` for joins). 5. (Learned graph methods) ``predict_sheet`` wall time per held-out sheet; aggregates (mean/p95 ms per graph node, throughput) in ``inference_timing`` on the same JSON entry. First sheet may include one-off CPU/GPU warmup — interpret percentiles accordingly. When structure metrics are recomputed, the new entry is merged with the previous one so ``inference_timing_per_sheet`` (and other extra keys) are not dropped unless this run saves a fresh per-sheet list with ``--save-structure-details``. Usage: python eval_rag.py # run on rag_eval_dataset.json python eval_rag.py --queries path/to.json # custom query file python eval_rag.py --methods graph,pandas # specific methods only """ from __future__ import annotations import json import logging import os import time from datetime import datetime, timezone from pathlib import Path from typing import Any, Callable, Optional import numpy as np import config import rag_eval_storage from chunk_builder import extract_chunks, extract_row_chunks, chunk_to_text from evaluate import evaluate_sheet, naive_flat_chunks from label_utils import is_data_label logger = logging.getLogger("eval_rag") def _load_embed_model(): from embed_text import _get_model return _get_model() def _embed_texts(model, texts: list[str], batch_size: int = 64) -> np.ndarray: """Embed a list of texts using E5. Returns (N, 768) float32 array.""" prefixed = [f"query: {t}" if t else "" for t in texts] vecs = model.encode( prefixed, batch_size=batch_size, show_progress_bar=False, normalize_embeddings=True, ) return np.array(vecs, dtype=np.float32) # ═══════════════════════════════════════════════════════════════════════════ # Internal chunking methods (return chunks with positional metadata) # ═══════════════════════════════════════════════════════════════════════════ # Sheets above this cell count use naive chunking directly (avoids OS OOM kill # which cannot be caught by Python try/except). Sheets below this threshold # attempt GNN prediction with a fallback to naive on RuntimeError/MemoryError. GNN_CELL_LIMIT = 3000 def _rag_grid_cell_counts(grid: dict) -> tuple[int, int]: """Return (n_json_cells, n_graph_nodes). ``n_graph_nodes`` matches ``predict_sheet`` nodes: cells with non-empty ``value``. ``n_json_cells`` is ``len(cells)`` on the grid dict. """ cells = grid.get("cells") or [] n_json = len(cells) n_graph = sum(1 for c in cells if (c.get("value") or "") != "") return n_json, n_graph def _aggregate_predict_timing( per_sheet: list[dict[str, Any]], *, n_skipped_limit: int = 0, stats_rows: list[dict[str, Any]] | None = None, ) -> dict[str, Any]: """Aggregate wall-clock ``predict_sheet`` timings into JSON-serialisable metrics. ``per_sheet`` is the full list of attempts (for counts). ``stats_rows``, if set, is the subset used for mean / percentiles (e.g. exclude the first sheet for warmup). """ rows_stats = stats_rows if stats_rows is not None else per_sheet ok_rows = [r for r in rows_stats if r.get("success")] n_failed = sum(1 for r in per_sheet if not r.get("skipped_limit") and not r.get("success")) def _ms_per_node(r: dict) -> float: n = max(int(r.get("n_graph_nodes") or 0), 1) return float(r["elapsed_s"]) * 1000.0 / n def _ms_per_json(r: dict) -> float: n = max(int(r.get("n_json_cells") or 0), 1) return float(r["elapsed_s"]) * 1000.0 / n ms_nodes = [_ms_per_node(r) for r in ok_rows] ms_json = [_ms_per_json(r) for r in ok_rows] total_s = sum(float(r["elapsed_s"]) for r in ok_rows) total_nodes = sum(int(r.get("n_graph_nodes") or 0) for r in ok_rows) def _pctl(arr: list[float], q: float) -> float | None: if not arr: return None return float(np.percentile(np.array(arr, dtype=np.float64), q)) out: dict[str, Any] = { "n_sheets_attempted": len(per_sheet), "n_sheets_used_for_percentiles": len(rows_stats), "n_sheets_predict_ok": len(ok_rows), "n_sheets_predict_failed": n_failed, "n_sheets_skipped_gnn_cell_limit": int(n_skipped_limit), "predict_total_s": round(total_s, 4) if ok_rows else None, "total_graph_nodes": int(total_nodes) if ok_rows else None, "throughput_nodes_per_s": round(total_nodes / total_s, 2) if ok_rows and total_s > 0 else None, "mean_ms_per_graph_node": round(float(np.mean(ms_nodes)), 3) if ms_nodes else None, "p50_ms_per_graph_node": (round(v, 3) if (v := _pctl(ms_nodes, 50)) is not None else None), "p95_ms_per_graph_node": (round(v, 3) if (v := _pctl(ms_nodes, 95)) is not None else None), "mean_ms_per_json_cell": round(float(np.mean(ms_json)), 3) if ms_json else None, "p95_ms_per_json_cell": (round(v, 3) if (v := _pctl(ms_json, 95)) is not None else None), } return out def _predict_and_label( grid_data: dict, model=None, *, sheet_id: str | None = None, pred_cache: dict[str, list[dict] | None] | None = None, ) -> dict: """Run GNN prediction on raw grid data and return a labeled copy. Merges predicted labels + table IDs back into the cell dicts so that extract_chunks() / build_knowledge_graph() can operate on them. Returns grid_data unchanged if prediction fails (OOM, etc.). If ``pred_cache`` is provided with ``sheet_id``, uses cached predictions from :func:`_fill_prediction_cache` (single inference per sheet per architecture). """ n_cells = len(grid_data.get("cells", [])) from predict import predict_sheet if pred_cache is not None and sheet_id is not None and sheet_id in pred_cache: predictions = pred_cache[sheet_id] else: try: predictions = predict_sheet(grid_data, model_override=model) except (RuntimeError, MemoryError) as e: logger.warning("GNN prediction failed on sheet with %d cells: %s", n_cells, e) import gc; gc.collect() return grid_data if not predictions: return grid_data pred_map = {(p["row"], p["col"]): p for p in predictions} labeled_cells = [] for c in grid_data["cells"]: cell = dict(c) pred = pred_map.get((c["row"], c["col"])) if pred: cell["label"] = pred["label"] cell["table"] = pred.get("table", 1) cell["confidence"] = pred.get("confidence", 0.0) else: cell["label"] = "empty" cell["table"] = 0 labeled_cells.append(cell) labeled = dict(grid_data) labeled["cells"] = labeled_cells return labeled def _fill_prediction_cache( sheet_data: dict[str, dict], model, log: Callable[[str], None] | None = None, *, timing_exclude_first_sheet: bool = False, ) -> tuple[dict[str, list[dict] | None], dict[str, Any]]: """One ``predict_sheet`` per held-out sheet (under GNN_CELL_LIMIT). Returns ``(pred_cache, timing_bundle)`` where ``timing_bundle`` has keys ``aggregate`` (from :func:`_aggregate_predict_timing`) and ``per_sheet`` (list of per-attempt records for optional persistence / debugging). """ from predict import predict_sheet cache: dict[str, list[dict] | None] = {} per_sheet: list[dict[str, Any]] = [] n_skipped_limit = 0 for sheet_id, data in sheet_data.items(): n_json, n_graph = _rag_grid_cell_counts(data) n_cells = len(data.get("cells", [])) if n_cells > GNN_CELL_LIMIT: cache[sheet_id] = None n_skipped_limit += 1 continue t0 = time.perf_counter() try: p = predict_sheet(data, model_override=model) elapsed = time.perf_counter() - t0 cache[sheet_id] = p if p else None success = bool(p) per_sheet.append({ "sheet_id": sheet_id, "elapsed_s": round(elapsed, 6), "n_json_cells": n_json, "n_graph_nodes": n_graph, "success": success, "skipped_limit": False, "ms_per_graph_node": round(elapsed * 1000.0 / max(n_graph, 1), 4), "ms_per_json_cell": round(elapsed * 1000.0 / max(n_json, 1), 4), }) except (RuntimeError, MemoryError) as e: elapsed = time.perf_counter() - t0 if log: log(f" predict_sheet failed {sheet_id[:40]}... ({n_cells} cells): {e}") logger.warning("predict_sheet failed on %s: %s", sheet_id, e) import gc; gc.collect() cache[sheet_id] = None per_sheet.append({ "sheet_id": sheet_id, "elapsed_s": round(elapsed, 6), "n_json_cells": n_json, "n_graph_nodes": n_graph, "success": False, "skipped_limit": False, "error": str(e), "ms_per_graph_node": round(elapsed * 1000.0 / max(n_graph, 1), 4), "ms_per_json_cell": round(elapsed * 1000.0 / max(n_json, 1), 4), }) stats_rows = per_sheet if timing_exclude_first_sheet and len(per_sheet) > 1: stats_rows = per_sheet[1:] agg = _aggregate_predict_timing( per_sheet, n_skipped_limit=n_skipped_limit, stats_rows=stats_rows, ) bundle = {"aggregate": agg, "per_sheet": per_sheet} if log: if per_sheet: a = agg log( f" predict_sheet timing: total={a.get('predict_total_s')}s " f"ok={a.get('n_sheets_predict_ok')}/{a.get('n_sheets_attempted')} " f"mean_ms/node={a.get('mean_ms_per_graph_node')} " f"p95_ms/node={a.get('p95_ms_per_graph_node')} " f"nodes/s={a.get('throughput_nodes_per_s')} " f"skipped_limit={a.get('n_sheets_skipped_gnn_cell_limit')}" ) else: log( f" predict_sheet timing: no predict attempts " f"(skipped GNN_CELL_LIMIT={n_skipped_limit} / {len(sheet_data)} sheets)" ) return cache, bundle def _aggregate_structure_sheet_results(rows: list[dict]) -> dict[str, float | None]: """Macro mean (equal weight per sheet) over ``evaluate_sheet`` outputs.""" def _mean(getter: Callable[[dict], float | None]) -> float | None: vals = [getter(r) for r in rows] nums = [float(v) for v in vals if v is not None and isinstance(v, (int, float))] if not nums: return None return sum(nums) / len(nums) if not rows: return {} out: dict[str, float | None] = {} out["level1_macro_f1"] = _mean(lambda r: r.get("level1", {}).get("macro_f1")) out["level1_weighted_f1"] = _mean(lambda r: r.get("level1", {}).get("weighted_f1")) out["level2_macro_f1"] = _mean(lambda r: r.get("level2", {}).get("macro_f1")) out["level3_joint_accuracy"] = _mean(lambda r: r.get("level3", {}).get("joint_accuracy")) out["level3_ari"] = _mean(lambda r: r.get("level3", {}).get("ari")) l4k = ( "chunk_f1", "chunk_precision", "chunk_recall", "soft_chunk_f1", "chunk_answerability", "soft_chunk_answerability", "soft_mean_iou", ) for k in l4k: out[f"level4_{k}"] = _mean(lambda r, kk=k: r.get("level4", {}).get(kk)) return out def _compute_structure_vs_gold_for_arch( sheet_data: dict[str, dict], pred_cache: dict[str, list[dict] | None], *, save_per_sheet: bool = False, ) -> dict[str, Any]: """Run ``evaluate_sheet`` on sheets that have human RAG structure labels.""" per_eval: list[dict] = [] sheet_ids_used: list[str] = [] skipped_no_gold: list[str] = [] skipped_no_pred: list[str] = [] for sheet_id, data in sheet_data.items(): if not _grid_has_rag_structure_labels(data): skipped_no_gold.append(sheet_id) continue preds = pred_cache.get(sheet_id) if not preds: skipped_no_pred.append(sheet_id) continue ev = evaluate_sheet(data, preds) if not ev: skipped_no_pred.append(sheet_id) continue sheet_ids_used.append(sheet_id) per_eval.append(ev) agg = _aggregate_structure_sheet_results(per_eval) if per_eval else {} meta: dict[str, Any] = { "n_sheets_with_gold": len(sheet_ids_used), "sheet_ids_evaluated": sheet_ids_used, "skipped_no_structure_labels": skipped_no_gold, "skipped_no_predictions_or_empty_eval": skipped_no_pred, "aggregates": agg, } if save_per_sheet and per_eval: meta["_per_sheet"] = [ {"sheet_id": sid, "levels": row} for sid, row in zip(sheet_ids_used, per_eval) ] return meta def _load_rag_structure_metrics_doc() -> dict[str, Any]: p = config.RAG_EVAL_STRUCTURE_METRICS if not p.exists(): return {} try: with open(p, encoding="utf-8") as f: d = json.load(f) return d if isinstance(d, dict) else {} except Exception: return {} def _write_rag_structure_metrics_doc(doc: dict[str, Any]) -> None: path = config.RAG_EVAL_STRUCTURE_METRICS path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w", encoding="utf-8") as f: json.dump(doc, f, indent=2, default=str) def _should_skip_structure_resume( method_key_base: str, resume: bool, existing: dict[str, Any], force: bool, ) -> bool: if force or not resume: return False entry = existing.get(method_key_base) if not isinstance(entry, dict): return False if entry.get("error"): return False agg = entry.get("aggregates") if isinstance(agg, dict) and agg.get("level4_chunk_f1") is not None: return True if isinstance(agg, dict) and entry.get("n_sheets_with_gold", 0) == 0: return True return False def _graph_chunks( grid_data: dict, model=None, *, sheet_id: str = "", pred_cache: dict[str, list[dict] | None] | None = None, ) -> list[dict]: """Run GNN prediction → structured chunking (value + header paths).""" labeled = _predict_and_label( grid_data, model=model, sheet_id=sheet_id or None, pred_cache=pred_cache, ) chunks = extract_chunks(labeled) result = [] for c in chunks: text = chunk_to_text(c) if text: pos = c.get("position", {}) result.append({ "text": text, "method": "graph", "row": pos.get("row"), "col": pos.get("col"), }) if not result: return _naive_chunks(labeled) return result def _graph_row_chunks( grid_data: dict, model=None, *, sheet_id: str = "", pred_cache: dict[str, list[dict] | None] | None = None, ) -> list[dict]: """Run GNN prediction → row-based structured chunking.""" labeled = _predict_and_label( grid_data, model=model, sheet_id=sheet_id or None, pred_cache=pred_cache, ) chunks = extract_row_chunks(labeled) result = [] for c in chunks: text = c.get("text", "") if text: pos = c.get("position", {}) result.append({ "text": text, "method": "graph_row", "row": pos.get("row"), "col": pos.get("col"), }) if not result: return _naive_chunks(labeled) return result def _graph_kg_chunks( grid_data: dict, model=None, *, sheet_id: str = "", pred_cache: dict[str, list[dict] | None] | None = None, ) -> list[dict]: """Run GNN prediction → knowledge graph → graph-walk chunking.""" from chunk_builder import build_knowledge_graph, generate_chunks_from_graph labeled = _predict_and_label( grid_data, model=model, sheet_id=sheet_id or None, pred_cache=pred_cache, ) sheet_name = labeled.get("sheet_name", "") G = build_knowledge_graph(labeled["cells"], sheet_name) kg_chunks = generate_chunks_from_graph(G, sheet_name) result = [] for c in kg_chunks: if c.get("text"): result.append({ "text": c["text"], "method": "graph_kg", "row": None, "col": None, }) if not result: return _naive_chunks(labeled) return result # ── Oracle chunking (human labels from cached JSON; no _predict_and_label) ── def _graph_chunks_gold(grid_data: dict) -> list[dict]: """Structured graph chunks using labels in ``grid_data`` as-is (oracle / upper bound).""" chunks = extract_chunks(dict(grid_data)) result = [] for c in chunks: text = chunk_to_text(c) if text: pos = c.get("position", {}) result.append({ "text": text, "method": "graph_gold", "row": pos.get("row"), "col": pos.get("col"), }) if not result: return _naive_chunks(dict(grid_data)) return result def _graph_row_chunks_gold(grid_data: dict) -> list[dict]: """Row chunks from human labels only (oracle).""" chunks = extract_row_chunks(dict(grid_data)) result = [] for c in chunks: text = c.get("text", "") if text: pos = c.get("position", {}) result.append({ "text": text, "method": "graph_row_gold", "row": pos.get("row"), "col": pos.get("col"), }) if not result: return _naive_chunks(dict(grid_data)) return result def _graph_kg_chunks_gold(grid_data: dict) -> list[dict]: """KG walk chunks from human labels only (oracle).""" from chunk_builder import build_knowledge_graph, generate_chunks_from_graph labeled = dict(grid_data) sheet_name = labeled.get("sheet_name", "") or labeled.get("source", {}).get("sheet_name", "") G = build_knowledge_graph(labeled["cells"], sheet_name) kg_chunks = generate_chunks_from_graph(G, sheet_name) result = [] for c in kg_chunks: if c.get("text"): result.append({ "text": c["text"], "method": "graph_kg_gold", "row": None, "col": None, }) if not result: return _naive_chunks(labeled) return result def _naive_chunks(labeled_data: dict) -> list[dict]: cells = labeled_data["cells"] sheet = labeled_data.get("sheet_name", labeled_data.get("source", {}).get("sheet_name", "")) return naive_flat_chunks(cells, sheet) def _single_cell_chunks(labeled_data: dict) -> list[dict]: cells = labeled_data["cells"] chunks = [] for c in cells: val = c.get("value", "").strip() if val: chunks.append({"text": val, "method": "single_cell", "row": c["row"], "col": c["col"]}) return chunks # ═══════════════════════════════════════════════════════════════════════════ # Retrieval strategies # ═══════════════════════════════════════════════════════════════════════════ RETRIEVAL_STRATEGIES = ["vector", "bm25", "hybrid_rrf", "hybrid_rrf_rerank"] DEFAULT_RETRIEVAL_STRATEGIES = ["vector", "bm25", "hybrid_rrf"] def _build_faiss_index(embeddings: np.ndarray): """Build a FAISS inner-product index.""" import faiss dim = embeddings.shape[1] index = faiss.IndexFlatIP(dim) index.add(embeddings) return index def _retrieve_vector(index, query_vec: np.ndarray, k: int = 10) -> list[int]: """Return top-K indices via dense vector similarity.""" D, I = index.search(query_vec.reshape(1, -1), k) return [int(idx) for idx in I[0] if idx >= 0] def _tokenize(text: str) -> list[str]: """Simple whitespace + lowercasing tokenizer for BM25.""" import re return re.findall(r"\w+", text.lower()) def _build_bm25_index(texts: list[str]): """Build a BM25 index over tokenized chunk texts.""" from rank_bm25 import BM25Okapi tokenized = [_tokenize(t) for t in texts] return BM25Okapi(tokenized) def _retrieve_bm25(bm25_index, query: str, k: int = 10) -> list[int]: """Return top-K indices via BM25 keyword matching.""" scores = bm25_index.get_scores(_tokenize(query)) top_k = np.argsort(scores)[::-1][:k] return [int(idx) for idx in top_k if scores[idx] > 0] def _retrieve_hybrid_rrf( faiss_index, query_vec: np.ndarray, bm25_index, query_text: str, k: int = 10, rrf_k: int = 60, ) -> list[int]: """Reciprocal Rank Fusion of vector + BM25 results.""" vec_ids = _retrieve_vector(faiss_index, query_vec, k=k * 2) bm25_ids = _retrieve_bm25(bm25_index, query_text, k=k * 2) rrf_scores: dict[int, float] = {} for rank, idx in enumerate(vec_ids): rrf_scores[idx] = rrf_scores.get(idx, 0) + 1.0 / (rrf_k + rank + 1) for rank, idx in enumerate(bm25_ids): rrf_scores[idx] = rrf_scores.get(idx, 0) + 1.0 / (rrf_k + rank + 1) ranked = sorted(rrf_scores, key=rrf_scores.get, reverse=True) return ranked[:k] def _retrieve_hybrid_rrf_rerank( faiss_index, query_vec: np.ndarray, bm25_index, query_text: str, chunk_embs: np.ndarray, k: int = 10, pool_factor: int = 4, rrf_k: int = 60, ) -> list[int]: """RRF fusion for a wide candidate pool, then re-order by dense similarity.""" pool = max(k * pool_factor, k) vec_ids = _retrieve_vector(faiss_index, query_vec, k=pool) bm25_ids = _retrieve_bm25(bm25_index, query_text, k=pool) rrf_scores: dict[int, float] = {} for rank, idx in enumerate(vec_ids): rrf_scores[idx] = rrf_scores.get(idx, 0) + 1.0 / (rrf_k + rank + 1) for rank, idx in enumerate(bm25_ids): rrf_scores[idx] = rrf_scores.get(idx, 0) + 1.0 / (rrf_k + rank + 1) candidates = sorted(rrf_scores.keys(), key=lambda i: rrf_scores[i], reverse=True)[:pool] if not candidates: return [] q = query_vec.reshape(-1) scored = [(float(np.dot(chunk_embs[idx], q)), idx) for idx in candidates] scored.sort(reverse=True) return [idx for _, idx in scored[:k]] # ═══════════════════════════════════════════════════════════════════════════ # Source cell matching # ═══════════════════════════════════════════════════════════════════════════ def _chunk_contains_source(chunk: dict, source_cells: list[dict]) -> bool: """Check if a chunk covers any of the expected source cells. Two matching strategies (tried in order): 1. **Value match** — if source cells carry a ``value`` field, check whether that value string appears anywhere in the chunk text. This works for *all* chunking methods, even those without positional metadata (pandas, unstructured, beautifulsoup …). 2. **Position match** — fall back to row/col matching when chunk has positional metadata. """ chunk_text = chunk.get("text", "") c_row = chunk.get("row") c_col = chunk.get("col") for src in source_cells: src_val = src.get("value", "") if src_val and src_val in chunk_text: return True s_row, s_col = src.get("row"), src.get("col") if c_row is not None and c_row == s_row: if c_col is None or c_col == s_col: return True return False # ═══════════════════════════════════════════════════════════════════════════ # LLM-as-a-Judge # ═══════════════════════════════════════════════════════════════════════════ _ANSWER_GEN_PROMPT = """You are answering questions about data in spreadsheets. Given the following spreadsheet excerpts, answer the question. Include the specific values and numbers from the data that support your answer. If you cannot find the answer in the provided excerpts, say "NOT FOUND". Spreadsheet excerpts: {chunks} Question: {question} Answer:""" _JUDGE_PROMPT = """You are evaluating whether a generated answer correctly answers the same question as an expected answer. Both answers refer to data extracted from a spreadsheet. Expected answer: {expected} Generated answer: {generated} Rate the generated answer on a scale of 1-5: 1 = Completely wrong, contradicts the expected answer, or unrelated 2 = Partially related but key facts are wrong or misleading 3 = Answers a different aspect of the question, or gives only tangentially related information 4 = Correctly answers the core question but with less detail or supporting data than the expected answer 5 = Correctly answers the question with equivalent or sufficient detail Important: if the question is yes/no or asks "which one", and the generated answer gives the correct yes/no/choice, that is at least a 4 even if it omits supporting numbers. The core answer matters most. Respond with ONLY a single integer (1-5).""" JUDGE_MODELS = [ ("vertex_google", "gemini-2.0-flash"), ("vertex_google", "gemini-2.0-flash-lite"), ("vertex_google", "gemini-2.5-flash"), ("openai", "gpt-4o-mini"), ("openai", "gpt-4.1-nano"), ] _VERTEX_PROJECT = os.environ.get("GOOGLE_CLOUD_PROJECT", "acoustic-atom-386613") _VERTEX_LOCATION = os.environ.get("GOOGLE_CLOUD_LOCATION", "global") _openai_client = None _vertex_google_client = None _anthropic_client = None def _has_vertex_credentials() -> bool: """Check if Google Cloud / Vertex AI credentials are available.""" if os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"): return True try: import google.auth google.auth.default() return True except Exception: return False def _get_available_judges() -> list[tuple[str, str]]: """Return the subset of JUDGE_MODELS whose credentials are available.""" has_vertex = _has_vertex_credentials() has_openai = bool(os.environ.get("OPENAI_API_KEY")) has_anthropic = bool(os.environ.get("ANTHROPIC_API_KEY")) available = [] for provider, model in JUDGE_MODELS: if provider == "vertex_google" and has_vertex: available.append((provider, model)) elif provider == "anthropic" and has_anthropic: available.append((provider, model)) elif provider == "openai" and has_openai: available.append((provider, model)) return available def _call_openai(prompt: str, model: str) -> str: global _openai_client import openai if _openai_client is None: _openai_client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"]) resp = _openai_client.chat.completions.create( model=model, messages=[{"role": "user", "content": prompt}], max_tokens=256, temperature=0, ) return resp.choices[0].message.content.strip() def _call_vertex_google(prompt: str, model: str) -> str: global _vertex_google_client from google import genai from google.genai.types import GenerateContentConfig if _vertex_google_client is None: _vertex_google_client = genai.Client( vertexai=True, project=_VERTEX_PROJECT, location=_VERTEX_LOCATION, ) resp = _vertex_google_client.models.generate_content( model=model, contents=prompt, config=GenerateContentConfig(max_output_tokens=256, temperature=0), ) text = resp.text if text is None: raise RuntimeError(f"Empty response from {model}") return text.strip() def _call_anthropic(prompt: str, model: str) -> str: global _anthropic_client import anthropic if _anthropic_client is None: _anthropic_client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"]) resp = _anthropic_client.messages.create( model=model, max_tokens=256, messages=[{"role": "user", "content": prompt}], ) return resp.content[0].text.strip() _CALL_DISPATCH = { "openai": _call_openai, "vertex_google": _call_vertex_google, "anthropic": _call_anthropic, } def _call_llm(prompt: str, provider: str, model: str) -> str: """Call a single LLM by provider and model name.""" fn = _CALL_DISPATCH.get(provider) if fn is None: raise ValueError(f"Unknown provider: {provider}") return fn(prompt, model) def _generate_answer(question: str, chunks: list[str]) -> str: """Generate an answer using the first available judge model.""" chunk_text = "\n---\n".join(chunks[:5]) prompt = _ANSWER_GEN_PROMPT.format(chunks=chunk_text, question=question) judges = _get_available_judges() if not judges: raise RuntimeError( "No LLM credentials found. Set GOOGLE_APPLICATION_CREDENTIALS " "(for Vertex AI), ANTHROPIC_API_KEY, or OPENAI_API_KEY." ) provider, model = judges[0] return _call_llm(prompt, provider, model) def _parse_judge_score(response: str) -> int: """Extract a 1-5 integer from a judge response.""" try: score = int(response.strip()) return max(1, min(5, score)) except (ValueError, TypeError): import re match = re.search(r"[1-5]", response) return int(match.group()) if match else 1 def _judge_answer(expected: str, generated: str) -> dict: """Score answer correctness using all available judge models. Returns dict with per-model scores and the averaged score. """ if not generated or generated == "NOT FOUND": judges = _get_available_judges() per_model = {f"{p}/{m}": 1 for p, m in judges} if judges else {} return {"score": 1, "per_model": per_model} prompt = _JUDGE_PROMPT.format(expected=expected, generated=generated) judges = _get_available_judges() if not judges: raise RuntimeError("No LLM credentials found.") scores: dict[str, int] = {} for provider, model in judges: key = f"{provider}/{model}" try: response = _call_llm(prompt, provider, model) scores[key] = _parse_judge_score(response) except Exception as e: logger.warning("Judge %s failed: %s", key, e) if not scores: return {"score": 1, "per_model": {}} avg = sum(scores.values()) / len(scores) return {"score": round(avg, 2), "per_model": scores} # ═══════════════════════════════════════════════════════════════════════════ # Unified evaluation # ═══════════════════════════════════════════════════════════════════════════ def evaluate_method_v2( chunks: list[dict], queries: list[dict], embed_model, retrieval: str = "vector", ks: tuple[int, ...] = (1, 5, 10), use_llm_judge: bool = False, save_details: bool = False, _prebuilt_faiss=None, _prebuilt_bm25=None, _prebuilt_chunk_embs=None, ) -> dict: """Evaluate retrieval quality with source accuracy and optional LLM judge. Args: retrieval: one of "vector", "bm25", "hybrid_rrf", "hybrid_rrf_rerank" save_details: if True, ``result["_details"]`` will contain per-query data (retrieved chunks, generated answer, per-judge scores). _prebuilt_*: optional pre-computed indices to avoid rebuilding across retrieval strategies for the same set of chunks. """ texts = [c.get("text", "") for c in chunks] if not texts: result = {f"recall@{k}": 0.0 for k in ks} result.update({f"source_recall@{k}": 0.0 for k in ks}) result.update({"mrr": 0.0, "source_mrr": 0.0, "n_chunks": 0}) if use_llm_judge: result.update({"judge_score": 0.0, "judge_binary_acc": 0.0}) return result needs_vector = retrieval in ("vector", "hybrid_rrf", "hybrid_rrf_rerank") needs_bm25 = retrieval in ("bm25", "hybrid_rrf", "hybrid_rrf_rerank") chunk_embs = _prebuilt_chunk_embs faiss_index = _prebuilt_faiss bm25_index = _prebuilt_bm25 if needs_vector and chunk_embs is None: chunk_embs = _embed_texts(embed_model, texts) if needs_vector and faiss_index is None: faiss_index = _build_faiss_index(chunk_embs) if needs_bm25 and bm25_index is None: bm25_index = _build_bm25_index(texts) q_texts = [q["question"] for q in queries] q_embs = None if needs_vector: q_embs = _embed_texts(embed_model, q_texts) recalls = {k: 0 for k in ks} source_recalls = {k: 0 for k in ks} rr_sum = 0.0 source_rr_sum = 0.0 judge_scores = [] per_query_details: list[dict] = [] n_total = len(queries) for qi, q in enumerate(queries): expected_answer = q.get("expected_answer", q.get("expected_value", "")) source_cells = q.get("source_cells", []) if retrieval == "vector": valid_indices = _retrieve_vector(faiss_index, q_embs[qi], k=max(ks)) elif retrieval == "bm25": valid_indices = _retrieve_bm25(bm25_index, q["question"], k=max(ks)) elif retrieval == "hybrid_rrf": valid_indices = _retrieve_hybrid_rrf( faiss_index, q_embs[qi], bm25_index, q["question"], k=max(ks), ) elif retrieval == "hybrid_rrf_rerank": if chunk_embs is None: raise ValueError("hybrid_rrf_rerank requires chunk embeddings") valid_indices = _retrieve_hybrid_rrf_rerank( faiss_index, q_embs[qi], bm25_index, q["question"], chunk_embs, k=max(ks), ) else: raise ValueError(f"Unknown retrieval strategy: {retrieval}") retrieved_texts = [texts[idx] for idx in valid_indices] retrieved_chunks = [chunks[idx] for idx in valid_indices] first_answer_hit = None for rank, text in enumerate(retrieved_texts, start=1): if expected_answer and expected_answer in text: if first_answer_hit is None: first_answer_hit = rank break for k in ks: if any(expected_answer and expected_answer in t for t in retrieved_texts[:k]): recalls[k] += 1 if first_answer_hit is not None: rr_sum += 1.0 / first_answer_hit source_hit_rank = None if source_cells: first_source_hit = None for rank, chunk in enumerate(retrieved_chunks, start=1): if _chunk_contains_source(chunk, source_cells): if first_source_hit is None: first_source_hit = rank break source_hit_rank = first_source_hit for k in ks: if any(_chunk_contains_source(c, source_cells) for c in retrieved_chunks[:k]): source_recalls[k] += 1 if first_source_hit is not None: source_rr_sum += 1.0 / first_source_hit generated_answer = None verdict = None if use_llm_judge: try: generated_answer = _generate_answer(q["question"], retrieved_texts[:5]) verdict = _judge_answer(expected_answer, generated_answer) judge_scores.append(verdict) if (qi + 1) % 10 == 0 or qi == n_total - 1: running_avg = sum(v["score"] for v in judge_scores) / len(judge_scores) print(f" [judge] {qi+1}/{n_total} queries scored (running avg: {running_avg:.2f})") except Exception as e: logger.warning("LLM judge failed for query %d: %s", qi, e) judge_scores.append({"score": 1, "per_model": {}}) if save_details: detail = { "query_id": q.get("id", f"q_{qi}"), "question": q["question"], "expected_answer": expected_answer, "sheet_id": q.get("sheet_id", ""), "query_difficulty": q.get("query_difficulty"), "sheet_complexity": q.get("sheet_complexity"), "retrieved_chunks": [ {"rank": r + 1, "text": retrieved_texts[r], "row": retrieved_chunks[r].get("row"), "col": retrieved_chunks[r].get("col"), "contains_answer": bool(expected_answer and expected_answer in retrieved_texts[r]), "contains_source": _chunk_contains_source(retrieved_chunks[r], source_cells) if source_cells else None} for r in range(len(retrieved_texts)) ], "answer_hit_rank": first_answer_hit, "source_hit_rank": source_hit_rank, "generated_answer": generated_answer, "judge_verdict": verdict, } per_query_details.append(detail) n_q = max(len(queries), 1) result = {f"recall@{k}": recalls[k] / n_q for k in ks} result["mrr"] = rr_sum / n_q result.update({f"source_recall@{k}": source_recalls[k] / n_q for k in ks}) result["source_mrr"] = source_rr_sum / n_q result["n_chunks"] = len(texts) if use_llm_judge and judge_scores: avg_scores = [v["score"] for v in judge_scores] result["judge_score"] = sum(avg_scores) / len(avg_scores) result["judge_binary_acc"] = sum(1 for s in avg_scores if s >= 4) / len(avg_scores) all_model_keys: set[str] = set() for v in judge_scores: all_model_keys.update(v.get("per_model", {}).keys()) for mk in sorted(all_model_keys): model_scores = [v["per_model"][mk] for v in judge_scores if mk in v.get("per_model", {})] if model_scores: result[f"judge_{mk}"] = round(sum(model_scores) / len(model_scores), 3) if save_details: result["_details"] = per_query_details return result # ═══════════════════════════════════════════════════════════════════════════ # V2 runner (used by the API) # ═══════════════════════════════════════════════════════════════════════════ def _resolve_xlsx_path(query: dict, sheet_id: str) -> str | None: """Find the xlsx file path for a query's sheet. Checks (in order): query's spreadsheet_file, temp downloads, raw dir. """ if query.get("spreadsheet_file") and Path(query["spreadsheet_file"]).exists(): return query["spreadsheet_file"] drive_fid = query.get("drive_file_id", "") if drive_fid: tmp_path = config.TEMP_DIR / f"{drive_fid}.xlsx" if tmp_path.exists(): return str(tmp_path) try: import drive_client drive_client.download_file(drive_fid, str(tmp_path)) return str(tmp_path) except Exception: pass raw_dir = config.DATA_DIR / "raw" if raw_dir.exists(): for sub in raw_dir.iterdir(): if sub.is_dir(): for xlsx in sub.glob("*.xlsx"): sid = drive_fid or sheet_id.rsplit("_", 1)[0] if sid and sid in xlsx.stem: return str(xlsx) return None def _adapter_sheet_key(xlsx_path: str, sheet_name: str) -> tuple[str, str]: """Canonical key so the same workbook+sheet is not chunked twice.""" try: p = str(Path(xlsx_path).resolve()) except OSError: p = os.path.normpath(xlsx_path) return (p, sheet_name) def _load_sheet_data(sheet_id: str) -> dict | None: """Load grid data for RAG eval (RAG structure labels from Drive/local, else raw cache).""" return rag_eval_storage.load_rag_eval_grid_for_sheet(sheet_id) def _grid_has_rag_structure_labels(grid: dict | None) -> bool: """True if the loaded grid JSON has at least one non-trivial cell label (oracle / human structure).""" if not grid or not isinstance(grid, dict): return False for c in grid.get("cells") or []: lb = (c.get("label") or "").strip().lower() if lb and lb not in ("empty", "unlabeled"): return True return False def _load_rag_existing_aggregates() -> dict[str, Any]: """Load ``rag_eval_results.json`` if present (for resume / merge).""" path = config.DATA_DIR / "rag_eval_results.json" if not path.exists(): return {} try: with open(path, encoding="utf-8") as f: return json.load(f) except Exception: return {} def _should_skip_rag_resume( key: str, resume: bool, existing_agg: dict[str, Any], log_fn: Callable[[str], None] | None, *, need_judge: bool = False, ) -> bool: """Skip re-eval if this key already has a successful checkpoint (no ``error``). If ``need_judge`` is True (LLM-as-a-judge requested), only skip when retrieval metrics *and* judge fields are present — otherwise a prior ``--no-judge`` run would block adding judge scores. """ if not resume: return False val = existing_agg.get(key) if not isinstance(val, dict): return False if val.get("error"): return False has_retrieval = any(k in val for k in ("recall@1", "mrr")) has_judge = any(k in val for k in ("judge_score", "judge_binary_acc")) if need_judge: if has_retrieval and has_judge: if log_fn: log_fn(f" SKIP {key} (resume — already in rag_eval_results.json)") return True return False if has_retrieval or has_judge: if log_fn: log_fn(f" SKIP {key} (resume — already in rag_eval_results.json)") return True return False def _write_rag_results_batch( batch: dict[str, dict], *, save_details: bool, log_fn: Callable[[str], None] | None, existing_agg: dict[str, Any], ) -> None: """Merge ``batch`` into ``rag_eval_results.json`` on disk; optionally ``rag_eval_details.json``. Call after each method/strategy (or error batch) so a crash does not lose progress. """ if not batch: return path = config.DATA_DIR / "rag_eval_results.json" existing: dict[str, Any] = {} if path.exists(): try: with open(path, encoding="utf-8") as f: existing = json.load(f) except Exception: existing = {} details_delta: dict[str, Any] = {} for key, val in batch.items(): if isinstance(val, dict) and "_details" in val: details_delta[key] = val["_details"] val = {k: v for k, v in val.items() if k != "_details"} existing[key] = val existing_agg[key] = val path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w", encoding="utf-8") as f: json.dump(existing, f, indent=2, default=str) if log_fn: log_fn(f" [checkpoint] saved {len(batch)} key(s) → {path.name}") if save_details and details_delta: dpath = config.DATA_DIR / "rag_eval_details.json" ex_d: dict[str, Any] = {} if dpath.exists(): try: with open(dpath, encoding="utf-8") as f: ex_d = json.load(f) except Exception: ex_d = {} ex_d.update(details_delta) with open(dpath, "w", encoding="utf-8") as f: json.dump(ex_d, f, indent=2, default=str) if log_fn: log_fn(f" [checkpoint] details {len(details_delta)} key(s) → {dpath.name}") def run_rag_evaluation_v2( methods: list[str] | None = None, use_llm_judge: bool = True, save_details: bool = False, retrieval_strategies: list[str] | None = None, log_fn: Callable[[str], None] | None = None, architectures: list[str] | None = None, resume: bool = True, *, resume_structure: bool = True, force_structure_metrics: bool = False, save_structure_details: bool = False, timing_exclude_first_sheet: bool = False, ) -> dict: """Run full RAG evaluation using the rag_eval_dataset.json. Evaluates internal chunking methods (graph, naive, single_cell) and external baselines from rag_baselines.py. Sheet grids are loaded from ``rag_eval_cache`` then ``labeled`` (see ``_load_sheet_data``). Training data and RAG-eval data are separate concerns: RAG eval provides Q&A; chunking for ``graph_*`` uses the GNN on those grids. Oracle ``*_gold`` methods skip the GNN and use labels in the loaded JSON only, which is a meaningful ceiling **only** if those files actually contain human structure labels for the eval sheets. Args: retrieval_strategies: subset of RETRIEVAL_STRATEGIES to run (vector, bm25, hybrid_rrf, hybrid_rrf_rerank). Defaults to all except rerank. save_details: if True, merge per-query retrieval rows into ``rag_eval_details.json``. resume: if True (default), skip method/strategy keys already present in ``rag_eval_results.json`` with metrics (no ``error``), so you can restart after a crash. resume_structure: skip rewriting ``rag_eval_structure_metrics.json`` when an entry exists. force_structure_metrics: recompute structure-vs-gold even if a checkpoint exists. save_structure_details: include per-sheet ``_per_sheet`` in the structure metrics JSON. timing_exclude_first_sheet: if True, mean / percentiles for ``inference_timing`` exclude the first attempted ``predict_sheet`` call (reduces one-off warmup skew); counts still reflect all sheets. """ def log(msg: str): print(msg) if log_fn: log_fn(msg) strategies = retrieval_strategies or list(DEFAULT_RETRIEVAL_STRATEGIES) existing_agg = _load_rag_existing_aggregates() if resume and existing_agg: log( f"Resume: {len(existing_agg)} key(s) on disk — skipping rows that already " f"have metrics. Pass resume=False to recompute everything." ) dataset = rag_eval_storage.load_rag_eval_dataset() if not dataset.get("queries"): log("No RAG eval queries found. Add Q&As at /rag-eval or upload rag_eval_dataset.json to the Drive RAG folder.") return {} queries = dataset.get("queries", []) held_out = set(dataset.get("held_out_sheets", [])) log(f"Loaded {len(queries)} queries across {len(held_out)} held-out sheets") if methods is None: methods = [ "graph", "graph_kg", "naive", "single_cell", "pandas", "beautifulsoup", "lxml", ] available_judges = _get_available_judges() if use_llm_judge and not available_judges: log("No LLM credentials found — set GOOGLE_APPLICATION_CREDENTIALS, ANTHROPIC_API_KEY, or OPENAI_API_KEY — skipping LLM judge") use_llm_judge = False elif use_llm_judge: judge_names = [f"{p}/{m}" for p, m in available_judges] log(f"LLM judges ({len(available_judges)}): {', '.join(judge_names)}") sheet_data: dict[str, dict] = {} for sheet_id in held_out: data = _load_sheet_data(sheet_id) if data: sheet_data[sheet_id] = data log(f"Loaded {len(sheet_data)} held-out sheet(s) with data") structure_metrics_doc = _load_rag_structure_metrics_doc() # Enrich source_cells with actual cell values from grid data so that # _chunk_contains_source can match by value for methods without row/col. for q in queries: sid = q.get("sheet_id", "") grid = sheet_data.get(sid) if not grid: continue cell_map = {(c["row"], c["col"]): c for c in grid.get("cells", [])} for sc in q.get("source_cells", []): if "value" not in sc: cell = cell_map.get((sc.get("row"), sc.get("col"))) if cell: sc["value"] = cell.get("value", "") log("Loading embedding model...") embed_model = _load_embed_model() # Discover trained model architectures. Build a map of arch -> list # of candidate checkpoint paths (sorted by answerability desc). # Models are loaded lazily one-at-a-time during evaluation to avoid OOM. from predict import load_model_from_path, ARCH_REGISTRY experiments_dir = config.DATA_DIR / "experiments" arch_candidates: dict[str, list[Path]] = {} if experiments_dir.exists(): raw_candidates: dict[str, list[Path]] = {} for exp_dir in sorted(experiments_dir.iterdir()): if not exp_dir.is_dir(): continue ckpt_path = exp_dir / "best_model.pt" if not ckpt_path.exists(): continue arch = exp_dir.name.rsplit("_seed", 1)[0] if architectures and arch not in architectures: continue raw_candidates.setdefault(arch, []).append(exp_dir) for arch, dirs in raw_candidates.items(): def _score(d): rp = d / "results.json" if rp.exists(): try: r = json.load(open(rp)) return (r.get("chunk_answerability_mean", 0), r.get("chunk_f1_mean", 0)) except Exception: pass return (0, 0) dirs.sort(key=_score, reverse=True) arch_candidates[arch] = [d / "best_model.pt" for d in dirs] log(f" Found model: {arch} ({len(dirs)} checkpoint(s), best: {dirs[0].name})") if not arch_candidates: log(" WARNING: No experiment models found — graph methods will use default model") default_path = config.MODELS_DIR / "best_model.pt" if default_path.exists(): arch_candidates["default"] = [default_path] log(f" {len(arch_candidates)} architecture(s) available for prediction") internal_methods = {"graph", "graph_kg", "naive", "single_cell"} try: from rag_baselines import ADAPTERS except ImportError: log("Could not import rag_baselines — external methods unavailable") ADAPTERS = {} results: dict[str, Any] = {} # Oracle methods use human labels only (no checkpoint). Model methods are per-architecture. ORACLE_GRAPH_METHODS = frozenset({"graph_gold", "graph_row_gold", "graph_kg_gold"}) MODEL_GRAPH_METHODS = frozenset({"graph", "graph_kg", "graph_row"}) graph_methods = ORACLE_GRAPH_METHODS | MODEL_GRAPH_METHODS non_graph_methods = [m for m in methods if m not in graph_methods] requested_graph_methods = [m for m in methods if m in graph_methods] if sheet_data and ORACLE_GRAPH_METHODS.intersection(set(requested_graph_methods)): n_oracle_ready = sum( 1 for sid, d in sheet_data.items() if _grid_has_rag_structure_labels(d) ) log( f"Oracle (graph_*_gold): {n_oracle_ready}/{len(sheet_data)} held-out sheets have " f"human cell labels in JSON — others fall back to cache-only grids (weak oracle)." ) for gm in requested_graph_methods: # ── Oracle: human-labeled chunking (upper bound vs model predictions) ── if gm in ORACLE_GRAPH_METHODS: method_key_base = gm log(f"\nBuilding chunks for method: {method_key_base} (oracle / human labels)") all_chunks: list[dict] = [] try: for si, (sheet_id, data) in enumerate(sheet_data.items()): n_cells = len(data.get("cells", [])) if n_cells > GNN_CELL_LIMIT: fallback = _naive_chunks(data) all_chunks.extend(fallback) log( f" [{si+1}/{len(sheet_data)}] {sheet_id[:30]}... ({n_cells} cells) " f"-> naive fallback ({len(fallback)} chunks)" ) continue if gm == "graph_gold": all_chunks.extend(_graph_chunks_gold(data)) elif gm == "graph_row_gold": all_chunks.extend(_graph_row_chunks_gold(data)) else: all_chunks.extend(_graph_kg_chunks_gold(data)) if (si + 1) % 5 == 0: log(f" [{si+1}/{len(sheet_data)}] {len(all_chunks)} chunks so far") except Exception as e: log(f" ERROR building chunks for {method_key_base}: {e}") err_batch = { f"{method_key_base}/{strat}": {"n_chunks": 0, "error": str(e)} for strat in strategies } results.update(err_batch) _write_rag_results_batch( err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg, ) continue log(f" {len(all_chunks)} chunks") if not all_chunks: err_batch = { f"{method_key_base}/{strat}": {"n_chunks": 0, "error": "no chunks produced"} for strat in strategies } results.update(err_batch) _write_rag_results_batch( err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg, ) continue texts = [c.get("text", "") for c in all_chunks] log(f" Embedding {len(texts)} chunks for {method_key_base}...") chunk_embs = _embed_texts(embed_model, texts) faiss_index = _build_faiss_index(chunk_embs) bm25_index = _build_bm25_index(texts) for strat in strategies: key = f"{method_key_base}/{strat}" if _should_skip_rag_resume( key, resume, existing_agg, log, need_judge=use_llm_judge, ): results[key] = dict(existing_agg[key]) continue log(f" Evaluating {key}...") r = evaluate_method_v2( all_chunks, queries, embed_model, retrieval=strat, use_llm_judge=use_llm_judge, save_details=save_details, _prebuilt_faiss=faiss_index, _prebuilt_bm25=bm25_index, _prebuilt_chunk_embs=chunk_embs, ) results[key] = r _write_rag_results_batch( {key: r}, save_details=save_details, log_fn=log, existing_agg=existing_agg, ) for k, v in sorted(r.items()): if k.startswith("_"): continue if isinstance(v, float): log(f" {k}: {v:.4f}") else: log(f" {k}: {v}") continue # ── Model-based graph methods: one run per architecture ── for arch_name, ckpt_paths in arch_candidates.items(): method_key_base = f"{gm}_{arch_name}" log(f"\nBuilding chunks for method: {method_key_base}") # Try each checkpoint until one loads and produces chunks model = None for ckpt_path in ckpt_paths: model = load_model_from_path(ckpt_path) if model: log(f" Loaded {ckpt_path.parent.name}") break if not model: log(f" ERROR: no loadable checkpoint for {arch_name}") err_batch = { f"{method_key_base}/{strat}": {"n_chunks": 0, "error": "model load failed"} for strat in strategies } results.update(err_batch) _write_rag_results_batch( err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg, ) continue log(" predict_sheet cache (one inference per sheet)...") pred_cache, timing_bundle = _fill_prediction_cache( sheet_data, model, log, timing_exclude_first_sheet=timing_exclude_first_sheet, ) inf_agg = timing_bundle.get("aggregate") or {} inf_per = timing_bundle.get("per_sheet") or [] def _attach_inference_timing(entry: dict[str, Any]) -> dict[str, Any]: out = dict(entry) out["inference_timing"] = inf_agg if save_structure_details and inf_per: out["inference_timing_per_sheet"] = list(inf_per) # If this run did not emit per-sheet timing, keep any existing list from ``entry``. return out if not _should_skip_structure_resume( method_key_base, resume_structure, structure_metrics_doc, force_structure_metrics, ): try: smeta = _compute_structure_vs_gold_for_arch( sheet_data, pred_cache, save_per_sheet=save_structure_details, ) prev_entry = structure_metrics_doc.get(method_key_base) prev_dict = prev_entry if isinstance(prev_entry, dict) else {} entry = { **prev_dict, "method_base": method_key_base, "chunking": gm, "computed_at": datetime.now(timezone.utc).isoformat(), "rag_eval_query_count": len(queries), "held_out_sheet_count": len(sheet_data), **smeta, } entry.pop("structure_metrics_error", None) structure_metrics_doc[method_key_base] = _attach_inference_timing(entry) _write_rag_structure_metrics_doc(structure_metrics_doc) log( f" Structure vs gold: n_sheets={smeta['n_sheets_with_gold']} → " f"{config.RAG_EVAL_STRUCTURE_METRICS.name}" ) except Exception as e: log(f" WARNING: structure metrics failed for {method_key_base}: {e}") logger.exception("structure metrics") prev = structure_metrics_doc.get(method_key_base) if not isinstance(prev, dict): prev = {} fallback = { **prev, "method_base": method_key_base, "chunking": gm, "computed_at": datetime.now(timezone.utc).isoformat(), "structure_metrics_error": str(e), } structure_metrics_doc[method_key_base] = _attach_inference_timing(fallback) _write_rag_structure_metrics_doc(structure_metrics_doc) else: log(f" SKIP structure metrics for {method_key_base} (resume_structure)") prev = structure_metrics_doc.get(method_key_base) if not isinstance(prev, dict): prev = {} merged = {**prev, "method_base": method_key_base, "chunking": gm} merged["inference_timing_updated_at"] = datetime.now(timezone.utc).isoformat() structure_metrics_doc[method_key_base] = _attach_inference_timing(merged) _write_rag_structure_metrics_doc(structure_metrics_doc) log(f" Updated inference_timing only → {config.RAG_EVAL_STRUCTURE_METRICS.name}") all_chunks = [] try: for si, (sheet_id, data) in enumerate(sheet_data.items()): n_cells = len(data.get("cells", [])) if n_cells > GNN_CELL_LIMIT: fallback = _naive_chunks(data) all_chunks.extend(fallback) log(f" [{si+1}/{len(sheet_data)}] {sheet_id[:30]}... ({n_cells} cells) -> naive fallback ({len(fallback)} chunks)") continue try: if gm == "graph": all_chunks.extend( _graph_chunks( data, model=model, sheet_id=sheet_id, pred_cache=pred_cache, ) ) elif gm == "graph_row": all_chunks.extend( _graph_row_chunks( data, model=model, sheet_id=sheet_id, pred_cache=pred_cache, ) ) elif gm == "graph_kg": all_chunks.extend( _graph_kg_chunks( data, model=model, sheet_id=sheet_id, pred_cache=pred_cache, ) ) except (RuntimeError, MemoryError) as sheet_err: log(f" [{si+1}/{len(sheet_data)}] GNN failed on {sheet_id[:30]}... ({n_cells} cells): {sheet_err}") log(f" -> falling back to naive chunking for this sheet") import gc; gc.collect() fallback = _naive_chunks(data) all_chunks.extend(fallback) if (si + 1) % 5 == 0: log(f" [{si+1}/{len(sheet_data)}] {len(all_chunks)} chunks so far") except Exception as e: log(f" ERROR building chunks for {method_key_base}: {e}") err_batch = { f"{method_key_base}/{strat}": {"n_chunks": 0, "error": str(e)} for strat in strategies } results.update(err_batch) _write_rag_results_batch( err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg, ) del model import gc; gc.collect() continue del model import gc; gc.collect() log(f" {len(all_chunks)} chunks") if not all_chunks: err_batch = { f"{method_key_base}/{strat}": {"n_chunks": 0, "error": "no chunks produced"} for strat in strategies } results.update(err_batch) _write_rag_results_batch( err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg, ) continue texts = [c.get("text", "") for c in all_chunks] log(f" Embedding {len(texts)} chunks for {method_key_base}...") chunk_embs = _embed_texts(embed_model, texts) faiss_index = _build_faiss_index(chunk_embs) bm25_index = _build_bm25_index(texts) for strat in strategies: key = f"{method_key_base}/{strat}" if _should_skip_rag_resume( key, resume, existing_agg, log, need_judge=use_llm_judge, ): results[key] = dict(existing_agg[key]) continue log(f" Evaluating {key}...") r = evaluate_method_v2( all_chunks, queries, embed_model, retrieval=strat, use_llm_judge=use_llm_judge, save_details=save_details, _prebuilt_faiss=faiss_index, _prebuilt_bm25=bm25_index, _prebuilt_chunk_embs=chunk_embs, ) results[key] = r _write_rag_results_batch( {key: r}, save_details=save_details, log_fn=log, existing_agg=existing_agg, ) for k, v in sorted(r.items()): if k.startswith("_"): continue if isinstance(v, float): log(f" {k}: {v:.4f}") else: log(f" {k}: {v}") for method in non_graph_methods: log(f"\nBuilding chunks for method: {method}") all_chunks: list[dict] = [] method_queries = queries if method in internal_methods: for sheet_id, data in sheet_data.items(): if method == "naive": all_chunks.extend(_naive_chunks(data)) elif method == "single_cell": all_chunks.extend(_single_cell_chunks(data)) elif method in ADAPTERS: adapter_fn = ADAPTERS[method] # One chunk pass per unique (workbook path, sheet name) per held-out sheet. # Re-running the adapter for every query duplicated the corpus and skewed metrics. rep_q_by_sheet: dict[str, dict] = {} for q in queries: sid = q.get("sheet_id", "") if sid and sid not in rep_q_by_sheet: rep_q_by_sheet[sid] = q seen_sheet_keys: set[tuple[str, str]] = set() n_adapter_calls = 0 for sheet_id, _grid in sheet_data.items(): q = rep_q_by_sheet.get(sheet_id) if q is None: log(f" No representative query for {sheet_id!r}; skipping adapter pass") continue xlsx_path = _resolve_xlsx_path(q, sheet_id) sheet_name = q.get("sheet_name", "") if not xlsx_path or not sheet_name: log(f" Could not resolve xlsx for {sheet_id!r}") continue key = _adapter_sheet_key(xlsx_path, sheet_name) if key in seen_sheet_keys: continue seen_sheet_keys.add(key) try: adapter_chunks = adapter_fn(xlsx_path, sheet_name) all_chunks.extend(adapter_chunks) n_adapter_calls += 1 except Exception as e: log(f" Adapter {method} failed on {sheet_id} ({sheet_name!r}): {e}") log( f" Adapter {method}: {n_adapter_calls} unique sheet(s) " f"({len(sheet_data)} held-out sheet(s), {len(rep_q_by_sheet)} rep queries)" ) else: log(f" Unknown method: {method}, skipping") continue log(f" {len(all_chunks)} chunks") if not all_chunks: err_batch = { f"{method}/{strat}": {"n_chunks": 0, "error": "no chunks produced"} for strat in strategies } results.update(err_batch) _write_rag_results_batch( err_batch, save_details=save_details, log_fn=log, existing_agg=existing_agg, ) continue texts = [c.get("text", "") for c in all_chunks] log(f" Embedding {len(texts)} chunks for {method}...") chunk_embs = _embed_texts(embed_model, texts) faiss_index = _build_faiss_index(chunk_embs) bm25_index = _build_bm25_index(texts) for strat in strategies: key = f"{method}/{strat}" if _should_skip_rag_resume( key, resume, existing_agg, log, need_judge=use_llm_judge, ): results[key] = dict(existing_agg[key]) continue log(f" Evaluating {key}...") r = evaluate_method_v2( all_chunks, method_queries, embed_model, retrieval=strat, use_llm_judge=use_llm_judge, save_details=save_details, _prebuilt_faiss=faiss_index, _prebuilt_bm25=bm25_index, _prebuilt_chunk_embs=chunk_embs, ) results[key] = r _write_rag_results_batch( {key: r}, save_details=save_details, log_fn=log, existing_agg=existing_agg, ) for k, v in sorted(r.items()): if k.startswith("_"): continue if isinstance(v, float): log(f" {k}: {v:.4f}") else: log(f" {k}: {v}") # Separate per-query details from aggregate metrics details_data = {} agg_results = {} for key, val in results.items(): if isinstance(val, dict) and "_details" in val: details_data[key] = val.pop("_details") agg_results[key] = val out_path = config.DATA_DIR / "rag_eval_results.json" if out_path.exists(): with open(out_path) as f: existing = json.load(f) existing.update(agg_results) agg_results = existing with open(out_path, "w") as f: json.dump(agg_results, f, indent=2) log(f"\nAggregate results saved to {out_path}") if details_data: details_path = config.DATA_DIR / "rag_eval_details.json" existing_details = {} if details_path.exists(): with open(details_path) as f: existing_details = json.load(f) existing_details.update(details_data) with open(details_path, "w") as f: json.dump(existing_details, f, indent=2) log(f"Per-query details saved to {details_path}") for key in results: if key in details_data: results[key]["_details"] = details_data[key] return results # ═══════════════════════════════════════════════════════════════════════════ # Legacy v1 evaluation (backward compatible) # ═══════════════════════════════════════════════════════════════════════════ def evaluate_method( chunk_texts: list[str], queries: list[dict], embed_model, ks: tuple[int, ...] = (1, 5, 10), ) -> dict: """Evaluate retrieval quality for one chunking method (legacy v1).""" if not chunk_texts: return {f"recall@{k}": 0.0 for k in ks} | {"mrr": 0.0, "n_chunks": 0} chunk_embs = _embed_texts(embed_model, chunk_texts) index = _build_index(chunk_embs) q_texts = [q["question"] for q in queries] q_embs = _embed_texts(embed_model, q_texts) recalls = {k: 0 for k in ks} rr_sum = 0.0 for qi, q in enumerate(queries): expected = q.get("expected_value", "") _, indices = _retrieve(index, q_embs[qi], k=max(ks)) retrieved = [chunk_texts[int(idx)] for idx in indices if idx >= 0] first_hit = None for rank, text in enumerate(retrieved, start=1): if expected in text: if first_hit is None: first_hit = rank break for k in ks: top_k = retrieved[:k] if any(expected in t for t in top_k): recalls[k] += 1 if first_hit is not None: rr_sum += 1.0 / first_hit n_q = max(len(queries), 1) return { f"recall@{k}": recalls[k] / n_q for k in ks } | { "mrr": rr_sum / n_q, "n_chunks": len(chunk_texts), } def run_rag_evaluation( queries_path: Path | None = None, ) -> dict: """Run full RAG evaluation across internal methods (legacy v1).""" if queries_path is None: queries_path = config.DATA_DIR / "rag_queries.json" if not queries_path.exists(): print(f"No queries file at {queries_path}. Create it first.") print('Format: [{"question": "...", "expected_value": "...", "sheet_id": "..."}]') return {} with open(queries_path) as f: queries = json.load(f) print(f"Loaded {len(queries)} queries from {queries_path}") labeled_files = sorted(config.LABELED_DIR.glob("*.json")) if not labeled_files: print("No labeled data found.") return {} graph_texts: list[str] = [] naive_texts: list[str] = [] single_texts: list[str] = [] for fp in labeled_files: with open(fp) as f: data = json.load(f) graph_texts.extend(c["text"] for c in _graph_chunks(data)) naive_texts.extend(c["text"] for c in _naive_chunks(data)) single_texts.extend(c["text"] for c in _single_cell_chunks(data)) print(f"Chunks: graph={len(graph_texts)}, naive={len(naive_texts)}, single={len(single_texts)}") print("Loading embedding model...") embed_model = _load_embed_model() results = {} for name, texts in [("graph", graph_texts), ("naive", naive_texts), ("single_cell", single_texts)]: print(f"\nEvaluating {name}...") r = evaluate_method(texts, queries, embed_model) results[name] = r for k, v in r.items(): if isinstance(v, float): print(f" {k}: {v:.4f}") else: print(f" {k}: {v}") # Separate per-query details from aggregate metrics details_data = {} agg_results = {} for key, val in results.items(): if isinstance(val, dict) and "_details" in val: details_data[key] = val.pop("_details") agg_results[key] = val out_path = config.DATA_DIR / "rag_eval_results.json" if out_path.exists(): with open(out_path) as f: existing = json.load(f) existing.update(agg_results) agg_results = existing with open(out_path, "w") as f: json.dump(agg_results, f, indent=2) log(f"\nAggregate results saved to {out_path}") if details_data: details_path = config.DATA_DIR / "rag_eval_details.json" existing_details = {} if details_path.exists(): with open(details_path) as f: existing_details = json.load(f) existing_details.update(details_data) with open(details_path, "w") as f: json.dump(existing_details, f, indent=2) log(f"Per-query details saved to {details_path}") for key in results: if key in details_data: results[key]["_details"] = details_data[key] return results if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="RAG evaluation") parser.add_argument("--queries", type=str, help="Path to queries JSON file (legacy v1)") parser.add_argument("--methods", type=str, default=None, help="Comma-separated list of methods (v2)") parser.add_argument("--no-judge", action="store_true", help="Skip LLM-as-a-judge evaluation") parser.add_argument("--architectures", type=str, default=None, help="Comma-separated list of model architectures to evaluate") parser.add_argument("--strategies", type=str, default=None, help="Comma-separated retrieval strategies (vector,bm25,hybrid_rrf,hybrid_rrf_rerank)") parser.add_argument("--save-details", action="store_true", help="Save per-query details (retrieved chunks, judge verdicts)") parser.add_argument( "--no-resume", action="store_true", help="Re-run all method/strategy keys (ignore checkpoints in rag_eval_results.json)", ) parser.add_argument( "--pull-drive", action="store_true", help="Before eval: download rag_eval_dataset + RAG structure JSONs from Drive (OAuth or service account)", ) parser.add_argument( "--no-resume-structure", action="store_true", help="Recompute rag_eval_structure_metrics.json entries even if present", ) parser.add_argument( "--force-structure-metrics", action="store_true", help="Same as --no-resume-structure (recompute structure-vs-gold)", ) parser.add_argument( "--save-structure-details", action="store_true", help="Store per-sheet _per_sheet in rag_eval_structure_metrics.json", ) parser.add_argument( "--timing-exclude-first-sheet", action="store_true", help="For inference_timing aggregates only: drop the first predict_sheet call from mean/p95 (warmup)", ) args = parser.parse_args() if args.pull_drive: info = rag_eval_storage.prefetch_rag_eval_from_drive() print(json.dumps(info, indent=2)) if args.queries: qp = Path(args.queries) run_rag_evaluation(qp) else: method_list = args.methods.split(",") if args.methods else None arch_list = args.architectures.split(",") if args.architectures else None strat_list = args.strategies.split(",") if args.strategies else None run_rag_evaluation_v2( methods=method_list, use_llm_judge=not args.no_judge, architectures=arch_list, retrieval_strategies=strat_list, save_details=args.save_details, resume=not args.no_resume, resume_structure=not args.no_resume_structure, force_structure_metrics=args.force_structure_metrics, save_structure_details=args.save_structure_details, timing_exclude_first_sheet=args.timing_exclude_first_sheet, )