evaluate.py

"""
Reproducible evaluation of the extraction system against a gold standard.

    python evaluate.py tables \\
        --reference testovani/reference_master.csv \\
        --export BERT=testy_BERT.csv --export CLAUDE=testy_CLAUDE.csv \\
        [--ner-match set|multiset] [--json results.json]

    python evaluate.py anonymize     # PII recall of the anonymizer on tiny samples

``tables`` joins every engine's exported records (``/records/export``, CSV or
JSON) to the gold labels by KEY - the ``input_id`` the ingest step assigns
(``<file stem>#<row>``, e.g. ``test_set1#17``) - never by row order, and
prints for each engine: coverage, sentiment / topic accuracy (+ macro-F1 and
the majority-class baseline), entity-level NER precision / recall / F1,
latency (median / p95 / mean over MEASURED rows only) and cost.

Requirements for a valid run (see evaluation/README.md for the full procedure):
  * the reference file is the gold standard (columns: set, row_id, text,
    gold_topic, gold_sentiment, gold_entities as JSON);
  * the input files were named so their stem ends in the set number
    (``test_set1.csv`` <-> ``reference_set1``), which ties ``input_id`` to a
    gold row;
  * for NER, the system ran with ``STORE_RAW_ENTITIES=true`` (otherwise the
    exported entity text is redacted and only entity TYPES can match) - use
    synthetic or consented data for such a run;
  * messages with no exported row (failed analyses are never stored) count as
    missing: reported as coverage and as ``accuracy_all`` (missing = wrong);
  * NER is scored only on entity types the reference set itself uses (e.g. a
    reference with just PER/ORG/LOC/MISC labels ignores predicted DATE/TIME
    entities - see ``ner.excluded_types`` in the JSON output for counts, and
    evaluation/README.md for why).

Metric definitions live in ``evaluation/metrics.py``.
"""

from __future__ import annotations

import argparse
import csv
import json
import re
import sys
from pathlib import Path
from typing import Any

from evaluation import metrics

_KEY = re.compile(r"(\d+)#(\d+)$")
_SENTIMENTS = ["negative", "neutral", "positive"]


# ---------------------------------------------------------------------------
# Loading
# ---------------------------------------------------------------------------
def _set_number(name: str) -> int | None:
    m = re.search(r"(\d+)$", (name or "").strip())
    return int(m.group(1)) if m else None


def load_reference(path: str | Path) -> dict[tuple[int, int], dict[str, Any]]:
    """Gold rows keyed by (set number, row id)."""
    gold: dict[tuple[int, int], dict[str, Any]] = {}
    with open(path, encoding="utf-8-sig", newline="") as f:
        for row in csv.DictReader(f):
            set_no = _set_number(row["set"])
            if set_no is None:
                raise ValueError(f"Cannot read a set number from '{row['set']}'.")
            gold[(set_no, int(row["row_id"]))] = {
                "topic": row["gold_topic"],
                "sentiment": metrics.map_sentiment(row["gold_sentiment"]),
                "entities": json.loads(row["gold_entities"] or "[]"),
            }
    return gold


def parse_entity_cell(cell: str) -> list[dict[str, str]]:
    """CSV export cell 'TYPE:text; TYPE:text' -> entity dicts."""
    cell = (cell or "").replace("&amp;", "&")
    out = []
    for chunk in cell.split(";"):
        chunk = chunk.strip()
        if ":" in chunk:
            etype, etext = chunk.split(":", 1)
            if etype.strip() and etext.strip():
                out.append({"type": etype.strip(), "text": etext.strip()})
    return out


def _number(value: Any) -> float | None:
    if value in (None, ""):
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def load_export(path: str | Path) -> dict[tuple[int, int], dict[str, Any]]:
    """Exported records keyed by (set number, row) taken from ``input_id``."""
    path = Path(path)
    if path.suffix.lower() == ".json":
        raw_rows = json.loads(path.read_text(encoding="utf-8-sig"))
    else:
        with open(path, encoding="utf-8-sig", newline="") as f:
            raw_rows = list(csv.DictReader(f))

    rows: dict[tuple[int, int], dict[str, Any]] = {}
    for r in raw_rows:
        m = _KEY.search(str(r.get("input_id") or ""))
        if not m:
            continue  # exported before V9 (no input_id) or from another source
        entities = r.get("entities")
        if isinstance(entities, str):
            entities = parse_entity_cell(entities)
        rows[(int(m.group(1)), int(m.group(2)))] = {
            "topic": r.get("topic"),
            "sentiment": metrics.map_sentiment(r.get("sentiment")),
            "entities": entities or [],
            "latency_ms": _number(r.get("latency_ms")),
            "cost_usd": _number(r.get("cost_usd")),
            "engine": r.get("engine"),
        }
    return rows


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------
def evaluate_engine(
    gold: dict[tuple[int, int], dict[str, Any]],
    predictions: dict[tuple[int, int], dict[str, Any]],
    ner_match: str = "set",
) -> dict[str, Any]:
    keys = sorted(gold)
    covered = [k for k in keys if k in predictions]
    result: dict[str, Any] = {
        "gold_rows": len(keys),
        "covered": len(covered),
        "missing": len(keys) - len(covered),
        "coverage": len(covered) / len(keys) if keys else 0.0,
        "unmatched_export_rows": len([k for k in predictions if k not in gold]),
    }

    topic_labels = sorted({g["topic"] for g in gold.values() if g["topic"]})
    for name, labels in (("sentiment", _SENTIMENTS), ("topic", topic_labels)):
        g = [gold[k][name] for k in covered]
        p = [predictions[k][name] for k in covered]
        scores = metrics.classification_scores(g, p, labels)
        # accuracy_all: rows with no prediction count as wrong.
        scores["accuracy_all"] = (
            sum(1 for a, b in zip(g, p) if a == b) / len(keys) if keys else 0.0
        )
        result[name] = scores

    # A predicted entity of a type the gold standard never uses (e.g. DATE/TIME,
    # which processors/temporal.py adds but this reference set's annotation
    # scheme does not cover - see evaluation/README.md) can only ever be a false
    # positive: the gold side has no matching label space to score it against.
    # Scoring it anyway would penalize the system for extracting something the
    # reference set was never asked to judge, not for a real NER mistake.
    gold_entity_types = {e["type"] for row in gold.values() for e in row.get("entities", []) if e.get("type")}
    excluded_types: dict[str, int] = {}
    scored_predictions: list[list[dict[str, Any]]] = []
    for k in covered:
        kept, dropped = [], 0
        for e in predictions[k]["entities"]:
            if gold_entity_types and e.get("type") not in gold_entity_types:
                dropped += 1
                excluded_types[e["type"]] = excluded_types.get(e["type"], 0) + 1
            else:
                kept.append(e)
        scored_predictions.append(kept)

    result["ner"] = metrics.ner_scores(
        [gold[k]["entities"] for k in covered],
        scored_predictions,
        match=ner_match,
    )
    result["ner"]["excluded_types"] = excluded_types

    costs = [predictions[k]["cost_usd"] for k in covered]
    measured_costs = [c for c in costs if c is not None]
    is_local = all((predictions[k].get("engine") or "local") == "local" for k in covered) if covered else False
    result["performance"] = {
        "latency_ms": metrics.latency_stats(predictions[k]["latency_ms"] for k in covered),
        "cost_usd_total": sum(measured_costs) if measured_costs else None,
        "cost_rows_measured": len(measured_costs),
        "note": (
            "local engine: no external API fee (compute cost not measured)"
            if is_local
            else "cost summed over rows that report it"
        ),
    }
    return result


def _fmt(value: Any, digits: int = 3) -> str:
    return "n/a" if value is None else f"{value:.{digits}f}"


def print_tables(results: dict[str, dict[str, Any]]) -> None:
    print("=== Coverage and classification ===")
    header = f"{'Engine':10s} {'covered':>9s} {'sent.acc':>9s} {'sent.F1':>8s} {'topic.acc':>10s} {'topic.F1':>9s} {'base(sent)':>11s}"
    print(header)
    for label, r in results.items():
        print(
            f"{label:10s} {r['covered']:4d}/{r['gold_rows']:<4d} "
            f"{r['sentiment']['accuracy']:9.3f} {r['sentiment']['macro_f1']:8.3f} "
            f"{r['topic']['accuracy']:10.3f} {r['topic']['macro_f1']:9.3f} "
            f"{r['sentiment']['majority_baseline']:11.3f}"
        )
    print(f"\n=== NER (entity level, match={next(iter(results.values()))['ner']['match']}) ===")
    print(f"{'Engine':10s} {'P':>7s} {'R':>7s} {'F1':>7s} {'TP':>6s} {'FP':>6s} {'FN':>6s}")
    for label, r in results.items():
        n = r["ner"]
        print(f"{label:10s} {n['precision']:7.3f} {n['recall']:7.3f} {n['f1']:7.3f} {n['tp']:6d} {n['fp']:6d} {n['fn']:6d}")
        if n["excluded_types"]:
            print(f"           (ignored, not in reference's label space: {n['excluded_types']})")
    print("\n=== Performance (measured rows only) ===")
    print(f"{'Engine':10s} {'n':>5s} {'median ms':>10s} {'p95 ms':>9s} {'mean ms':>9s} {'cost USD':>10s}  note")
    for label, r in results.items():
        lat, perf = r["performance"]["latency_ms"], r["performance"]
        print(
            f"{label:10s} {lat['n']:5d} {_fmt(lat['median'], 0):>10s} {_fmt(lat['p95'], 0):>9s} "
            f"{_fmt(lat['mean'], 0):>9s} {_fmt(perf['cost_usd_total'], 4):>10s}  {perf['note']}"
        )


def run_tables(args: argparse.Namespace) -> int:
    gold = load_reference(args.reference)
    results: dict[str, dict[str, Any]] = {}
    for spec in args.export:
        if "=" not in spec:
            print(f"--export expects LABEL=path, got '{spec}'", file=sys.stderr)
            return 2
        label, path = spec.split("=", 1)
        predictions = load_export(path)
        if not predictions:
            print(f"{label}: no rows with an input_id in {path} - was it exported after V9?", file=sys.stderr)
            return 2
        results[label] = evaluate_engine(gold, predictions, args.ner_match)
    print_tables(results)
    if args.json:
        Path(args.json).write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
        print(f"\nWritten: {args.json}")
    return 0


# ---------------------------------------------------------------------------
# Anonymization (tiny illustrative check, see README for a real evaluation)
# ---------------------------------------------------------------------------
def evaluate_anonymize() -> None:
    from processors import anonymizer

    samples = [
        ("My name is John Smith, email john@example.com", ["John Smith", "john@example.com"]),
        ("Call Sarah at +1 202 555 0143", ["Sarah", "+1 202 555 0143"]),
        ("I live in Berlin and work at Google", ["Berlin", "Google"]),
    ]

    def recall(fn) -> float:
        found = total = 0
        for text, pii_list in samples:
            anon = fn(text)
            for pii in pii_list:
                total += 1
                found += pii.lower() not in anon.lower()
        return found / total if total else 0.0

    print("=== Anonymization (PII recall, 3 illustrative samples) ===")
    print(f"Active backend:       {anonymizer.backend_name()}")
    print(f"Recall (regex only):  {recall(anonymizer._regex_anonymize):.2%}")
    try:
        print(f"Recall (full system): {recall(anonymizer.anonymize_text):.2%}")
    except anonymizer.AnonymizationUnavailable as exc:
        print(f"Full system unavailable: {exc}")
    print("Sample is far too small for a claim - use a labelled PII set (see evaluation/README.md).")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="command", required=True)

    tables = sub.add_parser("tables", help="recompute the result tables from exports + gold standard")
    tables.add_argument("--reference", required=True)
    tables.add_argument("--export", action="append", required=True, metavar="LABEL=path")
    tables.add_argument("--ner-match", choices=("set", "multiset"), default="set")
    tables.add_argument("--json", help="also write the full results to this file")
    tables.set_defaults(func=run_tables)

    anon = sub.add_parser("anonymize", help="illustrative PII recall check")
    anon.set_defaults(func=lambda _a: evaluate_anonymize() or 0)

    args = parser.parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())