judgekit 0.2.0__py3-none-any.whl

judgekit/__init__.py

@@ -0,0 +1,21 @@
+"""judgekit: audit an LLM judge before you trust it."""
+
+from judgekit.audit import AuditReport, run_audit
+from judgekit.io import dump_verdicts, load_verdicts
+from judgekit.result import ProbeResult, SkippedProbe
+from judgekit.schema import BinaryVerdict, GradedVerdict, PairwiseVerdict
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "AuditReport",
+    "BinaryVerdict",
+    "GradedVerdict",
+    "PairwiseVerdict",
+    "ProbeResult",
+    "SkippedProbe",
+    "__version__",
+    "dump_verdicts",
+    "load_verdicts",
+    "run_audit",
+]

judgekit/agreement.py

@@ -0,0 +1,178 @@
+"""Agreement between a judge and a reference rater.
+
+Implemented from scratch on plain float arithmetic; the test suite cross-checks
+every function against an independent reference (scikit-learn, scipy, and the
+krippendorff package).  Nothing here calls into BLAS, so results are identical
+across platforms — reductions use numpy's fixed pairwise summation and the
+small confusion/coincidence matrices are accumulated in Python loops.
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+from collections.abc import Sequence
+from typing import Literal
+
+import numpy as np
+from numpy.typing import NDArray
+
+Weights = Literal["linear", "quadratic"] | None
+FloatArray = Sequence[float] | NDArray[np.float64]
+
+
+def _as_pair(
+    a: FloatArray, b: FloatArray
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    x = np.asarray(a, dtype=np.float64)
+    y = np.asarray(b, dtype=np.float64)
+    if x.shape != y.shape:
+        raise ValueError(f"length mismatch: {x.shape} vs {y.shape}")
+    if x.size == 0:
+        raise ValueError("no ratings")
+    return x, y
+
+
+def confusion(
+    a: FloatArray, b: FloatArray, labels: FloatArray
+) -> NDArray[np.float64]:
+    """Count matrix with rows indexed by ``a`` and columns by ``b``."""
+    x, y = _as_pair(a, b)
+    index = {label: i for i, label in enumerate(labels)}
+    matrix = np.zeros((len(labels), len(labels)), dtype=np.float64)
+    for u, v in zip(x.tolist(), y.tolist(), strict=True):
+        if u not in index or v not in index:
+            raise ValueError(f"rating pair ({u}, {v}) outside labels {list(labels)}")
+        matrix[index[u], index[v]] += 1.0
+    return matrix
+
+
+def cohen_kappa(
+    a: FloatArray,
+    b: FloatArray,
+    labels: FloatArray | None = None,
+    weights: Weights = None,
+) -> float:
+    """Cohen's kappa, optionally weighted for ordinal scales.
+
+    Disagreement weights follow scikit-learn: distances are between label
+    *positions* in ``labels``, so pass the full scale (for instance
+    ``[1, 2, 3, 4, 5]``) rather than letting a category that nobody used
+    silently shrink the scale.  ``quadratic`` is the usual choice for graded
+    scores.  Returns nan when raters use a single label, where chance-corrected
+    agreement is undefined.
+    """
+    x, y = _as_pair(a, b)
+    if labels is None:
+        labels = sorted(set(x.tolist()) | set(y.tolist()))
+    counts = confusion(x, y, labels)
+    n = counts.sum()
+    observed = counts / n
+    row = observed.sum(axis=1)
+    col = observed.sum(axis=0)
+    expected = np.outer(row, col)
+
+    k = len(labels)
+    positions = np.arange(k, dtype=np.float64)
+    delta = np.abs(positions[:, None] - positions[None, :])
+    if weights is None:
+        w = (delta > 0).astype(np.float64)
+    elif weights == "linear":
+        w = delta
+    elif weights == "quadratic":
+        w = delta**2
+    else:
+        raise ValueError(f"unknown weights {weights!r}")
+
+    expected_disagreement = float((w * expected).sum())
+    if expected_disagreement == 0.0:
+        return float("nan")
+    return 1.0 - float((w * observed).sum()) / expected_disagreement
+
+
+def krippendorff_alpha(
+    values_by_unit: Sequence[FloatArray],
+    level: Literal["nominal", "interval"] = "nominal",
+) -> float:
+    """Krippendorff's alpha over units rated by any number of raters.
+
+    ``values_by_unit`` holds, per unit, the ratings that were actually given;
+    missing ratings are simply absent, which is the point of alpha.  Units with
+    fewer than two ratings carry no agreement information and are skipped.
+    """
+
+    def distance(u: float, v: float) -> float:
+        if level == "interval":
+            return (u - v) ** 2
+        if level == "nominal":
+            return 0.0 if u == v else 1.0
+        raise ValueError(f"unknown level {level!r}")
+
+    units = [[float(v) for v in unit] for unit in values_by_unit if len(unit) >= 2]
+    if not units:
+        raise ValueError("no unit has two or more ratings")
+    n_pairable = float(sum(len(unit) for unit in units))
+
+    observed = 0.0
+    for unit in units:
+        m = len(unit)
+        pair_sum = 0.0
+        for i in range(m):
+            for j in range(m):
+                if i != j:
+                    pair_sum += distance(unit[i], unit[j])
+        observed += pair_sum / (m - 1)
+    observed /= n_pairable
+
+    pooled = Counter(value for unit in units for value in unit)
+    expected = 0.0
+    for u, n_u in pooled.items():
+        for v, n_v in pooled.items():
+            if u != v:
+                expected += n_u * n_v * distance(u, v)
+    expected /= n_pairable * (n_pairable - 1.0)
+
+    if expected == 0.0:
+        return float("nan")
+    return 1.0 - observed / expected
+
+
+def rank_average(values: FloatArray) -> NDArray[np.float64]:
+    """Ranks starting at 1, ties sharing the average rank."""
+    x = np.asarray(values, dtype=np.float64)
+    order = np.argsort(x, kind="stable")
+    ranks = np.empty(x.size, dtype=np.float64)
+    sorted_x = x[order]
+    i = 0
+    while i < x.size:
+        j = i
+        while j + 1 < x.size and sorted_x[j + 1] == sorted_x[i]:
+            j += 1
+        ranks[order[i : j + 1]] = (i + j) / 2.0 + 1.0
+        i = j + 1
+    return ranks
+
+
+def pearson(a: FloatArray, b: FloatArray) -> float:
+    x, y = _as_pair(a, b)
+    xc = x - x.mean()
+    yc = y - y.mean()
+    denominator = float(np.sqrt((xc**2).sum() * (yc**2).sum()))
+    if denominator == 0.0:
+        return float("nan")
+    return float((xc * yc).sum()) / denominator
+
+
+def spearman(a: FloatArray, b: FloatArray) -> float:
+    """Spearman rank correlation with average ranks for ties."""
+    return pearson(rank_average(a), rank_average(b))
+
+
+def exact_agreement(a: FloatArray, b: FloatArray) -> float:
+    x, y = _as_pair(a, b)
+    return float((x == y).mean())
+
+
+def within_one(a: FloatArray, b: FloatArray) -> float:
+    """Share of ratings at most one scale point apart, for integer scales."""
+    x, y = _as_pair(a, b)
+    return float((np.abs(x - y) <= 1.0 + 1e-12).mean())

judgekit/audit.py

@@ -0,0 +1,448 @@
+"""Run every applicable check on a judge's verdicts and collect the results.
+
+The audit takes two verdict files — the judge's and the reference (usually
+human) — and produces one ``AuditReport``: agreement, calibration, bias probes
+and consistency, each as a ``ProbeResult`` with a bootstrap interval.  Which
+checks run is decided by what the data contains; nothing is required beyond
+matching ``item_id`` values.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+
+from judgekit import bias, consistency
+from judgekit.agreement import cohen_kappa, exact_agreement, spearman, within_one
+from judgekit.bootstrap import bootstrap_ci
+from judgekit.calibration import brier, ece_equal_mass, murphy_decomposition, reliability_curve
+from judgekit.io import align_on_items, by_kind
+from judgekit.result import ProbeResult, SkippedProbe
+from judgekit.schema import BinaryVerdict, GradedVerdict, PairwiseVerdict, Verdict
+
+# The finite-sample floor of equal-mass ECE at the demo's scale is a few
+# hundredths even for a perfectly calibrated judge, so the innocent band
+# stops at 0.05 rather than zero.
+INNOCENT_ECE = (0.0, 0.05)
+
+
+@dataclass(frozen=True)
+class AuditReport:
+    judge_id: str
+    n_graded: int
+    n_binary: int
+    n_pairwise: int
+    agreement: list[ProbeResult] = field(default_factory=list)
+    calibration: list[ProbeResult] = field(default_factory=list)
+    probes: list[ProbeResult] = field(default_factory=list)
+    skipped: list[SkippedProbe] = field(default_factory=list)
+    # Small data slices the figures are drawn from.
+    reliability: list[tuple[float, float, int]] | None = None
+    graded_scores: tuple[list[float], list[float]] | None = None
+
+    @property
+    def results(self) -> list[ProbeResult]:
+        return [*self.agreement, *self.calibration, *self.probes]
+
+    @property
+    def flags(self) -> list[ProbeResult]:
+        return [result for result in self.results if result.triggered is True]
+
+    def to_dict(self) -> dict[str, object]:
+        """JSON-safe view of the whole audit, for pipelines that gate on it."""
+        return {
+            "judge_id": self.judge_id,
+            "n_graded": self.n_graded,
+            "n_binary": self.n_binary,
+            "n_pairwise": self.n_pairwise,
+            "agreement": [result.to_dict() for result in self.agreement],
+            "calibration": [result.to_dict() for result in self.calibration],
+            "probes": [result.to_dict() for result in self.probes],
+            "flags": [flag.name for flag in self.flags],
+            "skipped": [{"name": skip.name, "needs": skip.needs} for skip in self.skipped],
+        }
+
+
+def _grouped_scores(
+    pairs: Sequence[tuple[Verdict, Verdict]],
+    values: Sequence[tuple[float, float]],
+) -> list[list[tuple[float, float]]]:
+    units: dict[str, list[tuple[float, float]]] = {}
+    for (judge, _), row in zip(pairs, values, strict=True):
+        units.setdefault(judge.item_id, []).append(row)
+    return list(units.values())
+
+
+def _pooled(units: list[list[tuple[float, float]]]) -> tuple[list[float], list[float]]:
+    flat = [row for unit in units for row in unit]
+    return [row[0] for row in flat], [row[1] for row in flat]
+
+
+def _agreement_graded(
+    pairs: Sequence[tuple[GradedVerdict, GradedVerdict]], seed: int
+) -> list[ProbeResult]:
+    scale_min = pairs[0][0].scale_min
+    scale_max = pairs[0][0].scale_max
+    values = [(judge.score, human.score) for judge, human in pairs]
+    units = _grouped_scores(pairs, values)
+    results = []
+
+    integral = all(
+        float(judge).is_integer() and float(human).is_integer() for judge, human in values
+    )
+    if integral and scale_min.is_integer() and scale_max.is_integer():
+        labels = [float(v) for v in range(int(scale_min), int(scale_max) + 1)]
+
+        def kappa(sample: list[list[tuple[float, float]]]) -> float:
+            judge_scores, human_scores = _pooled(sample)
+            return cohen_kappa(judge_scores, human_scores, labels=labels, weights="quadratic")
+
+        value = kappa(units)
+        results.append(
+            ProbeResult(
+                name="quadratic-weighted kappa",
+                value=value,
+                ci=bootstrap_ci(kappa, units, seed=seed),
+                innocent=None,
+                n=len(values),
+                detail=(
+                    f"chance-corrected ordinal agreement with the reference "
+                    f"on {len(values)} verdicts"
+                ),
+            )
+        )
+
+    def rho(sample: list[list[tuple[float, float]]]) -> float:
+        judge_scores, human_scores = _pooled(sample)
+        return spearman(judge_scores, human_scores)
+
+    results.append(
+        ProbeResult(
+            name="spearman rho",
+            value=rho(units),
+            ci=bootstrap_ci(rho, units, seed=seed),
+            innocent=None,
+            n=len(values),
+            detail="rank correlation with the reference scores",
+        )
+    )
+
+    if integral:
+
+        def exact(sample: list[list[tuple[float, float]]]) -> float:
+            judge_scores, human_scores = _pooled(sample)
+            return exact_agreement(judge_scores, human_scores)
+
+        def close(sample: list[list[tuple[float, float]]]) -> float:
+            judge_scores, human_scores = _pooled(sample)
+            return within_one(judge_scores, human_scores)
+
+        results.append(
+            ProbeResult(
+                name="exact agreement",
+                value=exact(units),
+                ci=bootstrap_ci(exact, units, seed=seed),
+                innocent=None,
+                n=len(values),
+                detail="share of verdicts equal to the reference score",
+            )
+        )
+        results.append(
+            ProbeResult(
+                name="within one point",
+                value=close(units),
+                ci=bootstrap_ci(close, units, seed=seed),
+                innocent=None,
+                n=len(values),
+                detail="share of verdicts at most one scale point from the reference",
+            )
+        )
+    return results
+
+
+def _agreement_binary(
+    pairs: Sequence[tuple[BinaryVerdict, BinaryVerdict]], seed: int
+) -> list[ProbeResult]:
+    values = [(1.0 if judge.label else 0.0, 1.0 if human.label else 0.0) for judge, human in pairs]
+    units = _grouped_scores(pairs, values)
+
+    def accuracy(sample: list[list[tuple[float, float]]]) -> float:
+        judge_labels, human_labels = _pooled(sample)
+        return exact_agreement(judge_labels, human_labels)
+
+    def kappa(sample: list[list[tuple[float, float]]]) -> float:
+        judge_labels, human_labels = _pooled(sample)
+        return cohen_kappa(judge_labels, human_labels, labels=[0.0, 1.0])
+
+    return [
+        ProbeResult(
+            name="label accuracy",
+            value=accuracy(units),
+            ci=bootstrap_ci(accuracy, units, seed=seed),
+            innocent=None,
+            n=len(values),
+            detail=f"share of pass/fail labels matching the reference on {len(values)} verdicts",
+        ),
+        ProbeResult(
+            name="cohen kappa",
+            value=kappa(units),
+            ci=bootstrap_ci(kappa, units, seed=seed),
+            innocent=None,
+            n=len(values),
+            detail="chance-corrected label agreement",
+        ),
+    ]
+
+
+def _agreement_pairwise(
+    pairs: Sequence[tuple[PairwiseVerdict, PairwiseVerdict]], seed: int
+) -> list[ProbeResult]:
+    categories = {"a": 0.0, "b": 1.0, "tie": 2.0}
+    values = [(categories[judge.choice], categories[human.choice]) for judge, human in pairs]
+    units = _grouped_scores(pairs, values)
+
+    def match(sample: list[list[tuple[float, float]]]) -> float:
+        judge_choices, human_choices = _pooled(sample)
+        return exact_agreement(judge_choices, human_choices)
+
+    return [
+        ProbeResult(
+            name="choice agreement",
+            value=match(units),
+            ci=bootstrap_ci(match, units, seed=seed),
+            innocent=None,
+            n=len(values),
+            detail=(
+                f"share of presentations whose canonical choice matches the reference "
+                f"on {len(values)} verdicts"
+            ),
+        )
+    ]
+
+
+def _calibration(
+    pairs: Sequence[tuple[BinaryVerdict, BinaryVerdict]], seed: int
+) -> tuple[list[ProbeResult], list[tuple[float, float, int]] | None]:
+    confident = [
+        (judge.p_positive, 1.0 if human.label else 0.0)
+        for judge, human in pairs
+        if judge.p_positive is not None
+    ]
+    if len(confident) < 20:
+        return [], None
+    values = [(p, y) for p, y in confident if p is not None]
+    units = _grouped_scores(
+        [pair for pair in pairs if pair[0].p_positive is not None], values
+    )
+
+    def ece(sample: list[list[tuple[float, float]]]) -> float:
+        probabilities, outcomes = _pooled(sample)
+        return ece_equal_mass(probabilities, outcomes)
+
+    probabilities, outcomes = _pooled(units)
+    reliability_probe = ProbeResult(
+        name="expected calibration error",
+        value=ece(units),
+        ci=bootstrap_ci(ece, units, seed=seed),
+        innocent=INNOCENT_ECE,
+        n=len(values),
+        detail="mean gap between claimed confidence and observed frequency, equal-mass bins",
+    )
+    rel, res, unc = murphy_decomposition(probabilities, outcomes)
+    decomposition = ProbeResult(
+        name="brier score",
+        value=brier(probabilities, outcomes),
+        ci=(float("nan"), float("nan")),
+        innocent=None,
+        n=len(values),
+        detail=(
+            f"decomposes exactly into reliability {rel:.3f} - resolution {res:.3f} "
+            f"+ uncertainty {unc:.3f}"
+        ),
+    )
+    return [reliability_probe, decomposition], reliability_curve(probabilities, outcomes)
+
+
+def run_audit(
+    judge: Sequence[Verdict],
+    human: Sequence[Verdict],
+    *,
+    judge_model: str | None = None,
+    seed: int = 0,
+) -> AuditReport:
+    """Audit one judge against one reference; see the module docstring."""
+    judge_ids = {verdict.judge_id for verdict in judge}
+    if len(judge_ids) > 1:
+        raise ValueError(
+            f"one audit reads one judge, got judge_ids {sorted(judge_ids)}; "
+            "split the verdicts and audit each judge separately"
+        )
+    judge_graded, judge_binary, judge_pairwise = by_kind(judge)
+    human_graded, human_binary, human_pairwise = by_kind(human)
+    judge_id = judge[0].judge_id if judge else "judge"
+
+    agreement: list[ProbeResult] = []
+    calibration: list[ProbeResult] = []
+    probes: list[ProbeResult] = []
+    skipped: list[SkippedProbe] = []
+    reliability = None
+    graded_scores = None
+
+    graded_pairs = align_on_items(judge_graded, human_graded) if human_graded else []
+    if judge_graded and not graded_pairs:
+        skipped.append(
+            SkippedProbe(
+                "graded agreement and bias probes",
+                "reference graded verdicts sharing item_id values with the judge's",
+            )
+        )
+    if graded_pairs:
+        scales = {(v.scale_min, v.scale_max) for v, _ in graded_pairs} | {
+            (v.scale_min, v.scale_max) for _, v in graded_pairs
+        }
+        if len(scales) > 1:
+            raise ValueError(f"mixed score scales in one audit: {sorted(scales)}")
+        agreement.extend(_agreement_graded(graded_pairs, seed))
+        graded_scores = (
+            [judge_verdict.score for judge_verdict, _ in graded_pairs],
+            [human_verdict.score for _, human_verdict in graded_pairs],
+        )
+        verbosity_probe = bias.verbosity_graded(graded_pairs, seed=seed)
+        if verbosity_probe is not None:
+            probes.append(verbosity_probe)
+        else:
+            skipped.append(
+                SkippedProbe("verbosity bias", "candidate_len on at least 3 graded verdicts")
+            )
+
+        if judge_model is None:
+            skipped.append(
+                SkippedProbe(
+                    "self-preference",
+                    "the judge's model name (judge_model / --judge-model) "
+                    "plus candidate_model on graded verdicts",
+                )
+            )
+        else:
+            self_probe = bias.self_preference(graded_pairs, judge_model, seed=seed)
+            if self_probe is not None:
+                probes.append(self_probe)
+            else:
+                skipped.append(
+                    SkippedProbe(
+                        "self-preference",
+                        "candidate_model on graded verdicts, with both own-model "
+                        "and other-model candidates present",
+                    )
+                )
+
+        spread_probe = bias.central_tendency(graded_pairs, seed=seed)
+        if spread_probe is not None:
+            probes.append(spread_probe)
+        else:
+            skipped.append(
+                SkippedProbe(
+                    "central tendency",
+                    "at least 3 aligned graded verdicts with varying reference scores",
+                )
+            )
+
+    binary_pairs = align_on_items(judge_binary, human_binary) if human_binary else []
+    if judge_binary and not binary_pairs:
+        skipped.append(
+            SkippedProbe(
+                "binary agreement and calibration",
+                "reference binary verdicts sharing item_id values with the judge's",
+            )
+        )
+    if binary_pairs:
+        agreement.extend(_agreement_binary(binary_pairs, seed))
+        calibration, reliability = _calibration(binary_pairs, seed)
+        if not calibration:
+            confident = sum(1 for j, _ in binary_pairs if j.p_positive is not None)
+            skipped.append(
+                SkippedProbe(
+                    "calibration (ECE, Brier)",
+                    f"confidence on at least 20 aligned binary verdicts (found {confident})",
+                )
+            )
+
+    if judge_pairwise:
+        pairwise_pairs = (
+            align_on_items(judge_pairwise, human_pairwise) if human_pairwise else []
+        )
+        if not pairwise_pairs:
+            skipped.append(
+                SkippedProbe(
+                    "choice agreement and verbosity bias (pairwise)",
+                    "reference pairwise verdicts sharing item_id values with the judge's",
+                )
+            )
+        else:
+            agreement.extend(_agreement_pairwise(pairwise_pairs, seed))
+            verbosity = bias.verbosity_pairwise(pairwise_pairs, seed=seed)
+            if verbosity is not None:
+                probes.append(verbosity)
+            else:
+                skipped.append(
+                    SkippedProbe(
+                        "verbosity bias (pairwise)",
+                        "a_len and b_len on pairwise verdicts, with unequal lengths "
+                        "and decisive verdicts on both sides",
+                    )
+                )
+        for probe, skip in (
+            (
+                bias.position_preference(judge_pairwise, seed=seed),
+                SkippedProbe(
+                    "position preference", "at least one decisive (non-tie) pairwise verdict"
+                ),
+            ),
+            (
+                bias.swap_flip_rate(judge_pairwise, seed=seed),
+                SkippedProbe(
+                    "swap flip rate",
+                    "each pair judged in both presentation orders (swapped false and true)",
+                ),
+            ),
+            (
+                bias.identical_pair_decisiveness(judge_pairwise, seed=seed),
+                SkippedProbe(
+                    "identical-pair decisiveness",
+                    'pairs of identical candidates marked meta {"identical": true}',
+                ),
+            ),
+        ):
+            if probe is not None:
+                probes.append(probe)
+            else:
+                skipped.append(skip)
+
+    stability_ran = False
+    for kind_verdicts in (judge_graded, judge_binary, judge_pairwise):
+        if kind_verdicts:
+            stability = consistency.resample_consistency(kind_verdicts, seed=seed)
+            if stability is not None:
+                probes.append(stability)
+                stability_ran = True
+    if judge and not stability_ran:
+        skipped.append(
+            SkippedProbe(
+                "re-judgment unanimity",
+                "the same item judged more than once (sample_index 0..k-1, "
+                "per presentation order for pairwise)",
+            )
+        )
+
+    return AuditReport(
+        judge_id=judge_id,
+        n_graded=len(judge_graded),
+        n_binary=len(judge_binary),
+        n_pairwise=len(judge_pairwise),
+        agreement=agreement,
+        calibration=calibration,
+        probes=probes,
+        skipped=skipped,
+        reliability=reliability,
+        graded_scores=graded_scores,
+    )