eval/knowledge/score.py

"""Precision / recall / F1 per branch of the knowledge-extraction pipeline.

Migrated verbatim (bar the import bootstrap) from the `kex` prototype, 2026-08-19.
The prototype is not being ported; this scorer is, because it is the only thing
that can show extraction v2 matches or beats the measured baseline in
`results/baseline_prototype_2026-08-13_145132.json`.

Pipeline-independent by design: it scores plain lists of surfaces / entry dicts,
so it works against the prototype's artifacts and against v2 alike.

The one thing this module refuses to do is conflate **term-filter recall** with
**extraction precision** (spec §5). They are different failure modes with
different fixes: recall is fixed at stage 2 (GLiNER labels), precision is fixed
at stage 3 (model tier / prompt). E1 is the recall number specifically.

Gold sets are treated as partial by design — Mas Beta is the labelling
bottleneck, so scoring reports coverage rather than blocking on a complete file.
"""

from __future__ import annotations

import re
import unicodedata
from dataclasses import asdict, dataclass
from pathlib import Path

import yaml

GOLD_PATH = Path(__file__).resolve().parent / "knowledge_gold.yaml"


def norm(s: str) -> str:
    s = unicodedata.normalize("NFKC", s).casefold()
    s = re.sub(r"[^\w\s]", " ", s)
    return re.sub(r"\s+", " ", s).strip()


@dataclass
class Score:
    label: str
    n_gold: int
    n_pred: int
    true_positives: int
    precision: float
    recall: float
    f1: float
    misses: list[str]

    def as_dict(self) -> dict:
        d = asdict(self)
        if hasattr(self, "coverage"):
            d["coverage"] = self.coverage
        return d


def load_gold(path: Path) -> dict:
    with open(path, encoding="utf-8") as fh:
        return yaml.safe_load(fh)


def _prf(tp: int, n_pred: int, n_gold: int) -> tuple[float, float, float]:
    precision = tp / n_pred if n_pred else 0.0
    recall = tp / n_gold if n_gold else 0.0
    f1 = (
        2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
    )
    return round(precision, 4), round(recall, 4), round(f1, 4)


def score_term_filter(gold: dict, surfaces: list[str]) -> Score:
    """E1. A gold term counts as recalled if ANY of its variants appears among
    the filter's mention surfaces (substring match both ways, so 'PA' inside
    'Physical Availability (PA)' counts).

    Precision is reported but is NOT E1's criterion: the filter is deliberately
    over-inclusive, and clustering plus evidence ranking absorb the noise.
    """
    normed = {norm(s) for s in surfaces if norm(s)}
    blob = " | ".join(sorted(normed))
    tp, misses = 0, []
    gold_terms = gold.get("terms", [])
    for entry in gold_terms:
        variants = [entry["term"]] + list(entry.get("variants", []))
        if entry.get("full_name"):
            variants.append(entry["full_name"])
        hit = False
        for v in variants:
            nv = norm(v)
            if not nv:
                continue
            if nv in normed or re.search(rf"(?<![\w]){re.escape(nv)}(?![\w])", blob):
                hit = True
                break
        if hit:
            tp += 1
        else:
            misses.append(entry["term"])
    precision, recall, f1 = _prf(tp, len(normed), len(gold_terms))
    return Score(
        label="term_filter_recall(E1)",
        n_gold=len(gold_terms),
        n_pred=len(normed),
        true_positives=tp,
        precision=precision,
        recall=recall,
        f1=f1,
        misses=misses,
    )


def score_glossary(gold: dict, entries: list[dict]) -> Score:
    """E3: when nano fills the schema, is it right?

    Scoring is restricted to the SCOREABLE subset: entries that produced a
    definition AND whose term is in the gold set AND whose gold record carries
    `definition_contains` to check against.

    Why not simply tp/len(entries): the term filter is deliberately
    over-inclusive and the gold set is deliberately partial, so most entries are
    for terms gold says nothing about. Counting those as errors would measure
    gold coverage while claiming to measure nano's accuracy — precisely the
    conflation spec §5 forbids. Coverage is reported separately in as_dict().

    Substring matching, not exact — exact match would under-report
    correct-but-differently-worded extractions (spec §5).
    """
    gold_by_term: dict[str, dict] = {}
    for entry in gold.get("terms", []):
        for v in [entry["term"], *entry.get("variants", [])]:
            gold_by_term.setdefault(norm(v), entry)

    checkable_gold = [
        g for g in gold.get("terms", []) if g.get("definition_contains")
    ]
    n_checkable_gold = len(checkable_gold)

    correct, incorrect = 0, 0
    unscoreable_no_gold, unscoreable_no_criteria = 0, 0
    matched_gold, wrong, failures = set(), [], []

    for pred in entries:
        if not (pred.get("definition") or "").strip():
            continue  # abstention is scored separately, not as an error
        g = gold_by_term.get(norm(pred.get("term", "")))
        if not g:
            unscoreable_no_gold += 1
            continue
        required = [norm(x) for x in g.get("definition_contains", [])]
        if not required:
            unscoreable_no_criteria += 1
            continue
        definition = norm(pred.get("definition") or "")
        if all(r in definition for r in required):
            correct += 1
            matched_gold.add(g["term"])
        else:
            incorrect += 1
            wrong.append(f"{pred.get('term')} (definition did not match gold)")
            failures.append(
                {
                    "term": pred.get("term"),
                    "gold_requires": g.get("definition_contains"),
                    "extracted": (pred.get("definition") or "")[:240],
                }
            )

    misses = [g["term"] for g in checkable_gold if g["term"] not in matched_gold]
    n_scoreable = correct + incorrect
    precision, recall, f1 = _prf(correct, n_scoreable, n_checkable_gold)

    score = Score(
        label="glossary_schema_fill(E3)",
        n_gold=n_checkable_gold,
        n_pred=n_scoreable,
        true_positives=correct,
        precision=precision,
        recall=recall,
        f1=f1,
        misses=misses + wrong,
    )
    score.coverage = {  # type: ignore[attr-defined]
        "entries_total": len(entries),
        "entries_with_definition": sum(
            1 for e in entries if (e.get("definition") or "").strip()
        ),
        "abstained_null_definition": sum(
            1 for e in entries if not (e.get("definition") or "").strip()
        ),
        "scoreable": n_scoreable,
        "unscoreable_term_not_in_gold": unscoreable_no_gold,
        "unscoreable_gold_has_no_criteria": unscoreable_no_criteria,
        "gold_terms_with_criteria": n_checkable_gold,
        "failures": failures,
    }
    return score


def score_rules(gold: dict, entries: list[dict]) -> Score:
    gold_rules = gold.get("rules", [])
    pred_blobs = [
        norm(
            " ".join(
                str(v)
                for v in (e.get("statement"), e.get("condition"), e.get("consequence"))
                if v
            )
        )
        for e in entries
    ]
    tp, misses = 0, []
    for rule in gold_rules:
        required = [norm(x) for x in rule.get("statement_contains", [])]
        if any(all(r in blob for r in required) for blob in pred_blobs):
            tp += 1
        else:
            misses.append(rule["rule_id"])
    precision, recall, f1 = _prf(tp, len(entries), len(gold_rules))
    return Score(
        label="rule",
        n_gold=len(gold_rules),
        n_pred=len(entries),
        true_positives=tp,
        precision=precision,
        recall=recall,
        f1=f1,
        misses=misses,
    )