eval/probe_embedders.py

"""Rank probe: where does the correct chunk land under different embedding models?

Compares all-MiniLM-L6-v2 (production baseline), BAAI/bge-small-en-v1.5, and
intfloat/e5-small-v2 on the 7 baseline needs_review IN-CORPUS gaps.

Prefix conventions applied (getting these wrong silently cripples the model):
  MiniLM   : no prefix on either side  (symmetric model)
  BGE      : query gets "Represent this sentence for searching relevant passages: ";
             passages have no prefix  (per BAAI bge-small-en-v1.5 model card)
  E5       : query gets "query: "; passages get "passage: "
             (per intfloat e5-small-v2 model card)

No Groq, no reranker, no production changes.

Usage (from repo root):
    .venv\\Scripts\\python.exe eval/probe_embedders.py
    [--gold PATH] [--docs DIR] [--report PATH]
"""
from __future__ import annotations

import argparse
import gc
import json
import sys
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

import numpy as np

_EVAL_DIR = Path(__file__).resolve().parent
_REPO_ROOT = _EVAL_DIR.parent
sys.path.insert(0, str(_REPO_ROOT))

from eval.pipeline_adapter import ingest_document
from eval.metrics import normalize_text
from eval.runner import _resolve_pdf

# ── Embedder registry ─────────────────────────────────────────────────────────
EMBEDDERS = [
    {
        "key": "MiniLM",
        "model_name": "all-MiniLM-L6-v2",
        "query_prefix": "",
        "passage_prefix": "",
        "prefix_note": "no prefix on either side (symmetric model; matches production)",
    },
    {
        "key": "BGE",
        "model_name": "BAAI/bge-small-en-v1.5",
        "query_prefix": "Represent this sentence for searching relevant passages: ",
        "passage_prefix": "",
        "prefix_note": (
            'query: "Represent this sentence for searching relevant passages: {q}"; '
            "passage: no prefix  (BAAI bge-small-en-v1.5 model card)"
        ),
    },
    {
        "key": "E5",
        "model_name": "intfloat/e5-small-v2",
        "query_prefix": "query: ",
        "passage_prefix": "passage: ",
        "prefix_note": (
            'query: "query: {q}"; passage: "passage: {p}"  '
            "(intfloat e5-small-v2 model card)"
        ),
    },
]

TOP_THRESHOLDS = (20, 50)


# ── helpers ───────────────────────────────────────────────────────────────────
def _text_hit_ids(hint: str, chunks: List[Dict]) -> Set[int]:
    """chunk_ids whose extracted text contains hint verbatim (text path only)."""
    if not hint:
        return set()
    norm = normalize_text(hint)
    return {c["chunk_id"] for c in chunks if norm in normalize_text(c["text"])}


def _encode(model, texts: List[str], batch_size: int = 128) -> np.ndarray:
    """Encode a list of texts; return (n, dim) float32 L2-normalised matrix."""
    embs = model.encode(
        texts,
        batch_size=batch_size,
        show_progress_bar=False,
        normalize_embeddings=True,
        convert_to_numpy=True,
    )
    return embs.astype("float32")


def _best_rank(
    q_emb: np.ndarray,       # (dim,) unit vector
    chunk_embs: np.ndarray,  # (n, dim) unit matrix
    correct_ids: Set[int],
    chunks: List[Dict],
) -> Tuple[Optional[int], int]:
    """Return (rank_of_best_correct_chunk, total_chunks).  Rank is 1-indexed."""
    sims = chunk_embs @ q_emb          # cosine sim, shape (n,)
    order = np.argsort(-sims)          # descending
    n = len(chunks)
    for rank, idx in enumerate(order, start=1):
        if chunks[int(idx)]["chunk_id"] in correct_ids:
            return rank, n
    return None, n


# ── main ──────────────────────────────────────────────────────────────────────
def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--gold",   default="eval/benchmark/questions.json")
    parser.add_argument("--docs",   default="eval/benchmark/docs")
    parser.add_argument("--report", default="eval/results/baseline_report.json")
    args = parser.parse_args()

    def _abs(p: str) -> Path:
        path = Path(p)
        return path if path.is_absolute() else _REPO_ROOT / path

    questions: List[Dict] = json.loads(
        _abs(args.gold).read_text(encoding="utf-8")
    )
    by_id = {q["id"]: q for q in questions}

    report = json.loads(_abs(args.report).read_text(encoding="utf-8"))
    review_ids: List[str] = [item["id"] for item in report.get("needs_review", [])]

    if not review_ids:
        print("needs_review is empty -- nothing to probe.")
        return

    # Group review questions by document
    by_doc: Dict[str, List[str]] = {}
    for qid in review_ids:
        by_doc.setdefault(by_id[qid]["document_id"], []).append(qid)

    # ── Step 1: Ingest all PDFs, collect chunks + correct chunk info ──────────
    print("=" * 72)
    print(f"Rank probe -- {len(review_ids)} needs_review questions, "
          f"{len(EMBEDDERS)} embedders")
    print("=" * 72)
    print()
    print("Ingesting documents...", flush=True)

    doc_chunks: Dict[str, List[Dict]] = {}
    # correct_chunk_info: qid -> first correct chunk for the fairness block
    correct_chunk_info: Dict[str, Dict] = {}
    # correct_ids_map: qid -> set of chunk_ids that are correct
    correct_ids_map: Dict[str, Set[int]] = {}

    for doc_id, qids in by_doc.items():
        doc_file = by_id[qids[0]]["document_file"]
        pdf_path = _resolve_pdf(doc_file, _abs(args.docs))
        print(f"  {pdf_path.name}...", flush=True)
        chunks, _ = ingest_document(str(pdf_path))
        doc_chunks[doc_id] = chunks
        print(f"    {len(chunks)} chunks", flush=True)

        for qid in qids:
            hint = by_id[qid].get("supporting_text_hint", "")
            cids = _text_hit_ids(hint, chunks)
            correct_ids_map[qid] = cids
            # capture the first correct chunk for the fairness block
            for c in chunks:
                if c["chunk_id"] in cids:
                    correct_chunk_info[qid] = {
                        "chunk_id": c["chunk_id"],
                        "page":     c["page"],
                        "text":     c["text"],
                    }
                    break
            print(f"    [{qid}]  hint={hint[:60]!r}  "
                  f"correct_ids={sorted(cids)}", flush=True)

    # ── Step 2: For each embedder, encode docs + probe questions ──────────────
    # results[qid][key] = {"rank": int|None, "total": int}
    results: Dict[str, Dict[str, Dict]] = {qid: {} for qid in review_ids}

    from sentence_transformers import SentenceTransformer

    for cfg in EMBEDDERS:
        key   = cfg["key"]
        mname = cfg["model_name"]
        qpfx  = cfg["query_prefix"]
        ppfx  = cfg["passage_prefix"]

        print()
        print(f"--- {key}  ({mname}) ---")
        print(f"    Prefix: {cfg['prefix_note']}", flush=True)

        model = SentenceTransformer(mname)

        # Encode every document's chunks once
        doc_embs: Dict[str, np.ndarray] = {}
        for doc_id, chunks in doc_chunks.items():
            print(f"    Encoding {len(chunks)} chunks [{doc_id}]...", flush=True)
            texts = [ppfx + c["text"] for c in chunks]
            doc_embs[doc_id] = _encode(model, texts)

        # Rank each question
        for doc_id, qids in by_doc.items():
            chunks = doc_chunks[doc_id]
            chunk_embs = doc_embs[doc_id]
            for qid in qids:
                cids = correct_ids_map[qid]
                if not cids:
                    results[qid][key] = {"rank": None, "total": len(chunks)}
                    print(f"    [{qid}]  SKIP (no correct chunk in corpus)", flush=True)
                    continue

                question = by_id[qid]["question"]
                q_emb = _encode(model, [qpfx + question])[0]

                rank, total = _best_rank(q_emb, chunk_embs, cids, chunks)
                results[qid][key] = {"rank": rank, "total": total}

                flags = "  ".join(
                    f"top-{t}={'Y' if rank is not None and rank <= t else 'N'}"
                    for t in TOP_THRESHOLDS
                )
                print(f"    [{qid}]  rank={rank}/{total}  {flags}", flush=True)

        del model
        gc.collect()

    # ── Output section ────────────────────────────────────────────────────────
    SEP72  = "=" * 72
    SEP88  = "=" * 88
    keys = [cfg["key"] for cfg in EMBEDDERS]

    # -- Per-question rank table -----------------------------------------------
    print()
    print(SEP88)
    print("PER-QUESTION RANK TABLE  (rank = cosine nearest-neighbor position, 1-best)")
    print(SEP88)
    print()

    # Build header
    hdr  = f"{'id':<34} {'type':<12}"
    for k in keys:
        hdr += f"  {k+'-rank':<10} {'t20':<4} {'t50':<4}"
    print(hdr)
    print("-" * len(hdr))

    for qid in review_ids:
        q   = by_id[qid]
        row = f"{qid:<34} {q['question_type']:<12}"
        for k in keys:
            r     = results[qid].get(k, {})
            rank  = r.get("rank")
            total = r.get("total", "?")
            rank_str = f"{rank}/{total}" if rank is not None else f">{total}"
            t20 = "Y" if rank is not None and rank <= 20 else "N"
            t50 = "Y" if rank is not None and rank <= 50 else "N"
            row += f"  {rank_str:<10} {t20:<4} {t50:<4}"
        print(row)

    # -- Fairness blocks -------------------------------------------------------
    print()
    print(SEP88)
    print("FAIRNESS BLOCKS  (question / correct-chunk pairs for human review)")
    print(SEP88)

    _NUMBER_WORDS = (
        "how much", "how many", "what percentage", "what amount",
        "what figure", "which figure", "what total", "how large",
        "how big", "what size", "what number",
    )

    for qid in review_ids:
        q    = by_id[qid]
        info = correct_chunk_info.get(qid)
        print(f"\n[{qid}]  ({q['question_type']})")
        print(f"  Q : {q['question']}")
        print(f"  hint: {q.get('supporting_text_hint', '')!r}")
        if info:
            preview = info["text"][:300].replace("\n", " ")
            print(f"  Chunk  chunk_id={info['chunk_id']}  page={info['page']}:")
            print(f"    {preview!r}")
            q_lower = q["question"].lower()
            asks_number = any(w in q_lower for w in _NUMBER_WORDS)
            has_digit   = any(ch.isdigit() for ch in info["text"])
            if asks_number and not has_digit:
                flag = ("Kind mismatch: question requests a quantitative value; "
                        "correct chunk is narrative prose with no visible digits.")
            elif asks_number and has_digit:
                flag = ("Kind match: question requests a number; "
                        "correct chunk contains numeric content.")
            else:
                flag = ("Both qualitative/conceptual -- "
                        "no obvious kind mismatch; pure semantic gap.")
        else:
            flag = "No correct chunk captured."
        print(f"  Flag: {flag}")

    # -- Summary table ---------------------------------------------------------
    print()
    print(SEP88)
    print("SUMMARY TABLE")
    print(SEP88)
    print()
    print(f"  {'Embedder':<38}  {'top-20 / 7':<12}  {'top-50 / 7'}")
    print(f"  {'-'*62}")
    for cfg in EMBEDDERS:
        k = cfg["key"]
        t20 = sum(
            1 for qid in review_ids
            if (r := results[qid].get(k, {})) and r.get("rank") and r["rank"] <= 20
        )
        t50 = sum(
            1 for qid in review_ids
            if (r := results[qid].get(k, {})) and r.get("rank") and r["rank"] <= 50
        )
        label = f"{k} ({cfg['model_name']})"
        print(f"  {label:<38}  {str(t20)+'/7':<12}  {t50}/7")

    print()
    print("  Verdicts per question:")
    print(f"  {'-'*80}")
    for qid in review_ids:
        q = by_id[qid]

        def _in(k: str, thresh: int) -> bool:
            r = results[qid].get(k, {}).get("rank")
            return r is not None and r <= thresh

        miniLM_t20 = _in("MiniLM", 20)
        other_t20  = any(_in(k, 20) for k in ("BGE", "E5"))
        any_t50    = any(_in(k, 50) for k in keys)
        any_t20    = any(_in(k, 20) for k in keys)

        if not miniLM_t20 and other_t20:
            verdict = "RESCUED by stronger embedder (MiniLM misses top-20; BGE/E5 hits top-20)"
        elif not any_t50:
            verdict = "STILL MISSED by all embedders -- candidate unfair/ambiguous question"
        elif not any_t20:
            verdict = "MARGINAL (best rank is top-50 but no embedder reaches top-20)"
        else:
            verdict = "TOP-20 by at least one model (including MiniLM)"

        ranks = "  ".join(
            f"{k}={results[qid].get(k, {}).get('rank', 'N/A')}"
            for k in keys
        )
        print(f"\n  {qid}  ({q['question_type']})")
        print(f"    ranks: {ranks}")
        print(f"    => {verdict}")

    print()
    print(SEP72)


if __name__ == "__main__":
    main()