evalrx 0.1.2__py3-none-any.whl

evalrx/analysis/planner.py

@@ -0,0 +1,283 @@
+"""Generic planning utilities for M2 analysis.
+
+The planner consumes a profile of the available data and emits a deterministic
+tool plan. It is intentionally conservative: discovery may inspect everything,
+but confirmatory plans should make tested families explicit and avoid silently
+dropping columns because of their original order.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from dataclasses import dataclass, field
+from typing import Any
+
+from evalrx.analysis.profile import DatasetProfile, profile_stats_input
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AnalysisPlanItem:
+    """One planned statistical analysis."""
+
+    tool: str
+    config: dict[str, Any] = field(default_factory=dict)
+    rationale: str = ""
+    estimand: str = ""
+    family: str = "confirmatory"
+    priority: float = 0.0
+
+    def as_legacy_tuple(self) -> tuple[str, dict[str, Any], str]:
+        return self.tool, self.config, self.rationale
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "tool": self.tool,
+            "config": self.config,
+            "rationale": self.rationale,
+            "estimand": self.estimand,
+            "family": self.family,
+            "priority": self.priority,
+        }
+
+
+def _signal_priority(inp: Any, signal: str, profile: DatasetProfile) -> float:
+    col = profile.columns.get(signal)
+    sigmap = (getattr(inp, "per_case", {}) or {}).get(signal, {})
+    labels = getattr(inp, "labels", {}) or {}
+    if not sigmap:
+        return 0.0
+    n_labeled = max(1, len(labels))
+    coverage = min(1.0, len(sigmap) / n_labeled)
+    values = [float(v) for v in sigmap.values()]
+    unique = len(set(values))
+    variance = 0.0
+    if values:
+        mean = sum(values) / len(values)
+        variance = sum((v - mean) ** 2 for v in values) / len(values)
+    kind_bonus = 0.2 if col and col.is_binary else 0.1 if col and col.dtype == "numeric" else 0.0
+    sparse_binary_bonus = 0.15 if col and col.is_binary and coverage < 0.95 else 0.0
+    constant_penalty = -1.0 if unique <= 1 else 0.0
+    return coverage + min(1.0, variance) + kind_bonus + sparse_binary_bonus + constant_penalty
+
+
+#: Agreement above which a BINARY signal is the label wearing another name.
+#: Not 1.0: one flipped case out of dozens is still a restatement, not a finding.
+_LABEL_RESTATEMENT_AGREEMENT = 0.95
+#: Below this many shared cases, perfect agreement is cheap and means nothing.
+_MIN_CASES_TO_JUDGE_RESTATEMENT = 10
+#: Names of fields an analyzer produces BY GRADING — i.e. by asking the same
+#: question the label answers.  Matched on the last dotted segment.
+_CORRECTNESS_NAME = re.compile(
+    r"(?:^|_)(?:correct|incorrect|label|labelled|passed|is_pass)(?:_|$)"
+    r"|_match$|^match$"
+)
+# These answer-audit fields are defined using the gold and/or the stored
+# PASS/FAIL label.  A sparse flag can have low whole-batch agreement while
+# still making ``P(FAIL | flag)=1`` by construction, so an association test is
+# circular.  They remain useful descriptive sanity signals, never candidate
+# discriminators for M2/M4.
+_LABEL_DERIVED_SUFFIXES = frozenset({
+    "answer_extraction_audit.extraction_suspect",
+    "answer_extraction_audit.extraction_point_miss",
+    "answer_extraction_audit.gold_in_output",
+    "answer_extraction_audit.gold_in_answer_region",
+    "answer_extraction_audit.strict_match",
+    "answer_extraction_audit.label_disagrees",
+    "answer_extraction_audit.labelled_fail",
+    # gold x pred conjunctions: P(FAIL | flag) = 1 by construction. Sparse
+    # enough to slip past label_leak_score, named nothing like "correct", so
+    # they sat in the VLM family as guaranteed BH survivors. The marginals
+    # (pope.answered_yes / pope.gold_yes) stay testable.
+    "pope.false_positive",
+    "pope.false_negative",
+})
+
+
+def restates_label(inp: Any, signal: str) -> bool:
+    """True when *signal* reproduces the PASS/FAIL label, or its complement.
+
+    Such a signal always rejects H0, always survives FDR, and carries no
+    information — "does the label predict the label". Left in the pool it also
+    CROWDS OUT the real ones, because the survivor list is what M4 draws on to
+    verify a hypothesis. Measured on qwen3.5-2b / bbh_causal_judgement: of 10
+    rejecting signals, 6 survived FDR and every one of those 6 was the label
+    (``labelled_fail``, ``strict_match``, ``calibration.correct``,
+    ``self_repair.baseline_correct`` / ``revised_correct``), while the two
+    genuinely informative extraction flags did not survive. M2's own narrative
+    called them out as "the label copied under another name"; M4 then verified a
+    hypothesis about label REPRODUCIBILITY using ``labelled_fail``, at
+    confidence 0.73.
+
+    Requires BOTH a correctness-shaped NAME and label agreement in the DATA,
+    because neither alone is safe and each fixes the other's failure:
+
+    * Data alone suppresses real findings. A composite discovered over
+      independent measurements — ``(obj_size < 40) and (attention < 0.3)`` —
+      can separate the labels perfectly, and that is the most valuable result
+      the loop can produce, not a tautology. It is statistically identical to
+      ``calibration.correct``; only its provenance differs.
+    * Name alone suppresses honest fields. A column called ``correct`` that
+      tracks something other than this batch's label is a legitimate signal.
+
+    Two further things this deliberately does NOT do:
+
+    * It does not judge arbitrary signals by the conditional rate. A sparse,
+      independently measured flag may legitimately have
+      ``P(FAIL | signal)=1``. The explicit answer-audit exceptions above are
+      removed by provenance because their definitions use gold/labels, not
+      because of their observed conditional rate.
+    * It does not touch continuous signals. Restatement is an identity claim,
+      and a continuous measure binarised at some threshold can drift into high
+      agreement without being the label at all.
+    """
+    labels = getattr(inp, "labels", {}) or {}
+    if str(signal).lower() in _LABEL_DERIVED_SUFFIXES:
+        return bool(labels)
+    if not _CORRECTNESS_NAME.search(str(signal).rsplit(".", 1)[-1].lower()):
+        return False
+    sigmap = (getattr(inp, "per_case", {}) or {}).get(signal, {})
+    shared = [key for key in sigmap if key in labels]
+    if len(shared) < _MIN_CASES_TO_JUDGE_RESTATEMENT:
+        return False
+    try:
+        values = {float(sigmap[key]) for key in shared}
+    except (TypeError, ValueError):
+        return False
+    if not values <= {0.0, 1.0} or len(values) < 2:
+        return False
+    agree = sum(1 for key in shared
+                if bool(float(sigmap[key])) == bool(labels[key]))
+    rate = agree / len(shared)
+    return (rate >= _LABEL_RESTATEMENT_AGREEMENT
+            or rate <= 1.0 - _LABEL_RESTATEMENT_AGREEMENT)
+
+
+def label_restating_signals(inp: Any) -> list[str]:
+    """The signals :func:`ranked_signal_names` drops, so callers can report them.
+
+    Dropped rather than tested-and-flagged on purpose: the harm is not that the
+    result is wrong, it is that the result is RIGHT and meaningless, and every
+    such test consumes a slot in the FDR family that a real signal needed.
+    """
+    return sorted(
+        signal for signal in (getattr(inp, "per_case", {}) or {})
+        if restates_label(inp, signal)
+    )
+
+
+def ranked_signal_names(inp: Any, *, max_signals: int | None = None) -> list[str]:
+    """Rank per-case signals by testability rather than insertion order."""
+    profile = profile_stats_input(inp)
+    scored = [
+        (_signal_priority(inp, signal, profile), signal)
+        for signal in (getattr(inp, "per_case", {}) or {})
+        if not restates_label(inp, signal)
+    ]
+    scored.sort(key=lambda item: (-item[0], item[1]))
+    names = [name for score, name in scored if score > 0]
+    if max_signals is not None:
+        names = names[:max(0, int(max_signals))]
+    return names
+
+
+def plan_stats_input(
+    inp: Any,
+    *,
+    max_signals: int | None = None,
+) -> list[AnalysisPlanItem]:
+    """Create a deterministic confirmatory plan for a ``StatsInput`` object."""
+    labels = getattr(inp, "labels", {}) or {}
+    groups = getattr(inp, "groups", None) or {}
+    vectors = getattr(inp, "per_case_vectors", {}) or {}
+    n_fail = sum(1 for value in labels.values() if value)
+    n_pass = len(labels) - n_fail
+    plan: list[AnalysisPlanItem] = []
+
+    if n_pass > 0 and n_fail > 0 and getattr(inp, "per_case", None):
+        # Dropped, not silently: a suppressed signal that turns out to be real
+        # has to leave a trace somewhere, and this is the only place that knows.
+        dropped = label_restating_signals(inp)
+        if dropped:
+            logger.info(
+                "planner: dropped %d signal(s) that restate the PASS/FAIL label "
+                "and would crowd the FDR family: %s",
+                len(dropped), ", ".join(dropped))
+        ranked = ranked_signal_names(inp, max_signals=max_signals)
+        for key in ranked:
+            plan.append(AnalysisPlanItem(
+                tool="signal_label_assoc",
+                config={"signal": key},
+                rationale=f"test whether per-case signal '{key}' predicts FAIL",
+                estimand=f"P(FAIL | {key}=present/high) - P(FAIL | {key}=absent/low)",
+                priority=1.0,
+            ))
+        continuous = [
+            key for key in ranked
+            if key in (getattr(inp, "per_case", {}) or {})
+            and len({float(v) for v in inp.per_case[key].values()}) > 2
+        ]
+        for key in continuous:
+            plan.append(AnalysisPlanItem(
+                tool="rank_corr",
+                config={"signal": key},
+                rationale=f"monotonic association between continuous '{key}' and FAIL",
+                estimand=f"Kendall tau({key}, FAIL)",
+                family="descriptive",
+                priority=0.5,
+            ))
+
+    if n_pass > 0 and n_fail > 0 and vectors:
+        vector_names = list(vectors)
+        if max_signals is not None:
+            vector_names = vector_names[:max(0, int(max_signals))]
+        for key in vector_names:
+            plan.append(AnalysisPlanItem(
+                tool="attention_decoding",
+                config={"signal": key},
+                rationale=f"omnibus: do FAIL/PASS per-case maps '{key}' differ?",
+                estimand=f"distributional difference in vector signal {key}",
+                priority=0.8,
+            ))
+
+    if labels:
+        plan.append(AnalysisPlanItem(
+            tool="single_rate_evalue",
+            config={},
+            rationale="describe the overall FAIL rate against an explicit baseline if provided",
+            estimand="P(FAIL) - p0",
+            family="descriptive",
+            priority=0.1,
+        ))
+
+    n_groups = len(groups)
+    if n_groups >= 3:
+        plan.append(AnalysisPlanItem(
+            tool="friedman_nemenyi",
+            config={},
+            rationale="rank 3+ strategies on shared cases",
+            estimand="strategy rank differences",
+            priority=0.9,
+        ))
+        names = list(groups)
+        base = names[0]
+        for variant in names[1:]:
+            plan.append(AnalysisPlanItem(
+                tool="mcnemar_evalue",
+                config={"strategies": [base, variant]},
+                rationale=f"paired contrast: does '{variant}' repair '{base}' failures?",
+                estimand=f"{variant} success - {base} success on paired cases",
+                priority=1.0,
+            ))
+    elif n_groups == 2:
+        plan.append(AnalysisPlanItem(
+            tool="mcnemar_evalue",
+            config={},
+            rationale="paired two-strategy comparison",
+            estimand="strategy success difference on paired cases",
+            priority=1.0,
+        ))
+
+    return plan

evalrx/analysis/probe_search.py

@@ -0,0 +1,203 @@
+"""Hierarchical MCTS probing search (ProbeLLM-style adaptive test-case discovery).
+
+Standalone tree-search skeleton: it knows nothing about models, judges, or
+datasets — those are injected as three callables (``generate_macro``,
+``generate_micro``, ``verify``). This mirrors the paper's two-level design
+(arXiv 2602.12966, "ProbeLLM"): a global choice between Macro (broad
+coverage — diversify topics) and Micro (local refinement — densify evidence
+around a seed) regimes, each running its own UCB-guided MCTS to pick which
+node to expand next, and a shared node/backup bookkeeping (Eqs. 5-10 in the
+paper).
+
+Failure cases discovered here are plain ``FailureCase`` objects — the result
+feeds directly into :func:`evalrx.analysis.failure_modes.cluster_failures`
+for failure-mode synthesis.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+from evalrx.core.case import CaseBatch, FailureCase, Label
+
+Regime = str  # "macro" | "micro"
+
+GenerateMacroFn = Callable[["ProbeNode", "list[ProbeNode]"], "FailureCase | None"]
+GenerateMicroFn = Callable[["ProbeNode"], "FailureCase | None"]
+VerifyFn = Callable[[FailureCase], FailureCase]
+
+
+@dataclass
+class ProbeNode:
+    """One evaluated test case in the search tree (paper Eq.5: (x, y*, y, fail)).
+
+    ``N``/``E`` are *this node's* expansion stats (paper Eq.6): how many times
+    a child was proposed+evaluated from here, and how many of those children
+    were failures. They are updated on nodes along the *selection path*
+    (root..this node), not on the freshly created child itself — a brand-new
+    node always starts at N=E=0 until it is later expanded.
+    """
+
+    case: FailureCase
+    regime: Regime
+    parent: "ProbeNode | None" = None
+    children: list["ProbeNode"] = field(default_factory=list)
+    N: int = 0
+    E: int = 0
+
+    def p_hat(self) -> float:
+        """Empirical failure rate of expansions from this node (Eq.6)."""
+        return self.E / max(1, self.N)
+
+    def is_fail(self) -> bool:
+        return self.case.label == Label.FAIL
+
+
+@dataclass
+class ProbeSearchResult:
+    """Output of one :meth:`ProbeSearch.run` call."""
+
+    macro_root: ProbeNode
+    micro_root: ProbeNode
+    n_simulations: int
+    n_macro: int
+    n_micro: int
+    failure_cases: CaseBatch
+    all_cases: CaseBatch
+
+    @property
+    def error_rate(self) -> float:
+        return len(self.failure_cases) / self.n_simulations if self.n_simulations else 0.0
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "n_simulations": self.n_simulations,
+            "n_macro": self.n_macro,
+            "n_micro": self.n_micro,
+            "error_rate": self.error_rate,
+            "n_failures": len(self.failure_cases),
+        }
+
+
+def _ucb(node: ProbeNode, parent_n: int, beta: float) -> float:
+    """UCB(v|u) = p_hat(v) + beta*sqrt(log(max(1,N(u))) / max(1,N(v)))  (Eq.7)."""
+    return node.p_hat() + beta * math.sqrt(math.log(max(1, parent_n)) / max(1, node.N))
+
+
+def _select_expandable(root: ProbeNode, *, beta: float, w_max: int) -> ProbeNode:
+    """Descend via UCB while the current node is already at max width,
+    stopping at the first node with room for a new child (paper Eq.7-8, the
+    "continue this selection step until reaching an expandable node")."""
+    u = root
+    while u.children and len(u.children) >= w_max:
+        u = max(u.children, key=lambda v: _ucb(v, u.N, beta))
+    return u
+
+
+def _choose_regime(macro_root: ProbeNode, micro_root: ProbeNode, beta: float) -> Regime:
+    """Top-level Macro/Micro choice, via the *same* UCB rule applied to the
+    two regime roots — the global root's two children in the paper's
+    hierarchy (Fig.2/Sec.3.2), rather than a fixed or alternating split."""
+    total_n = macro_root.N + micro_root.N
+    ucb_macro = _ucb(macro_root, total_n, beta)
+    ucb_micro = _ucb(micro_root, total_n, beta)
+    return "macro" if ucb_macro >= ucb_micro else "micro"
+
+
+class ProbeSearch:
+    """Hierarchical MCTS driver.
+
+    Args:
+        generate_macro:  ``(node, explored_macro_nodes) -> FailureCase | None``.
+                          Proposes one new, *unevaluated* candidate case aimed
+                          at broad topical coverage; receives every macro-tree
+                          node explored so far so it can pick its own
+                          diversity/frontier heuristic (e.g. embed + cluster
+                          medoids, paper Eq.11) — this module stays agnostic
+                          to how that is done.
+        generate_micro:  ``(node) -> FailureCase | None``. Proposes one local
+                          perturbation of ``node.case`` (paper's controlled
+                          entity/attribute substitutions).
+        verify:          ``(case) -> case``. Runs the target model + scorer,
+                          returning the same case with ``.observed``/``.label``
+                          populated (paper Eq.1's verifier V). A thin wrapper
+                          around ``CaseDiscoveryAgent.discover`` in practice.
+        seeds:           Initial seed cases (from an existing benchmark) used
+                          to anchor both regime roots.
+        budget:          Total number of simulations (T_max in Eq.4).
+        beta:            UCB exploration constant (Eq.7).
+        w_max:           Max children per node before progressive widening
+                          forces a deeper UCB descent instead of a new sibling.
+    """
+
+    def __init__(
+        self,
+        *,
+        generate_macro: GenerateMacroFn,
+        generate_micro: GenerateMicroFn,
+        verify: VerifyFn,
+        seeds: "list[FailureCase] | CaseBatch",
+        budget: int = 20,
+        beta: float = 1.0,
+        w_max: int = 3,
+    ) -> None:
+        seeds = list(seeds)
+        if not seeds:
+            raise ValueError("ProbeSearch requires at least one seed case")
+        self._generate_macro = generate_macro
+        self._generate_micro = generate_micro
+        self._verify = verify
+        self.budget = budget
+        self.beta = beta
+        self.w_max = max(1, w_max)
+
+        self.macro_root = ProbeNode(case=seeds[0], regime="macro")
+        self.micro_root = ProbeNode(case=seeds[0], regime="micro")
+        self._macro_nodes: list[ProbeNode] = [self.macro_root]
+        self._evaluated: list[FailureCase] = []
+        self._n_macro = 0
+        self._n_micro = 0
+
+    def run(self) -> ProbeSearchResult:
+        for _ in range(self.budget):
+            regime = _choose_regime(self.macro_root, self.micro_root, self.beta)
+            root = self.macro_root if regime == "macro" else self.micro_root
+            u = _select_expandable(root, beta=self.beta, w_max=self.w_max)
+
+            new_case = (
+                self._generate_macro(u, self._macro_nodes)
+                if regime == "macro"
+                else self._generate_micro(u)
+            )
+            if new_case is None:
+                continue  # generator exhausted for this node; budget step still spent
+
+            evaluated = self._verify(new_case)
+            child = ProbeNode(case=evaluated, regime=regime, parent=u)
+            u.children.append(child)
+            if regime == "macro":
+                self._macro_nodes.append(child)
+                self._n_macro += 1
+            else:
+                self._n_micro += 1
+            self._evaluated.append(evaluated)
+
+            fail = 1 if evaluated.label == Label.FAIL else 0
+            node: "ProbeNode | None" = u
+            while node is not None:
+                node.N += 1
+                node.E += fail
+                node = node.parent
+
+        failures = CaseBatch([c for c in self._evaluated if c.label == Label.FAIL])
+        return ProbeSearchResult(
+            macro_root=self.macro_root,
+            micro_root=self.micro_root,
+            n_simulations=len(self._evaluated),
+            n_macro=self._n_macro,
+            n_micro=self._n_micro,
+            failure_cases=failures,
+            all_cases=CaseBatch(self._evaluated),
+        )