Eval Metrics — Python source

"""Eval Metrics — the standard LLM-eval metrics with exact, testable formulas.

Language: Python (3.10+, standard library only)
Port of src/lib/evalMetrics.ts (the canonical TypeScript implementation).
Tool page: https://dev.cosmolabs.org/tools/eval-metrics

pass@k is the unbiased estimator from the Codex paper (Chen et al. 2021),
identical to the HumanEval implementation's combinatorial form. Precision,
recall, and F1 run over confusion counts and over label sets; exact-match
rate runs over paired strings (case-sensitive).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, Sequence


@dataclass(frozen=True)
class PRF1:
    """Precision / recall / F1 triple."""

    precision: float
    recall: float
    f1: float


@dataclass(frozen=True)
class ConfusionCounts:
    """Confusion counts; tn is accepted but unused by P/R/F1."""

    tp: int
    fp: int
    fn: int
    tn: int | None = None


def pass_at_k(n: int, c: int, k: int) -> float:
    """Unbiased pass@k: probability that at least one of k samples drawn
    without replacement from n (of which c are correct) passes.

        1                                       when n - c < k  (a wrong draw is impossible)
        1 - Π_{i=0..k-1} (n - c - i) / (n - i)  otherwise

    Raises ValueError on impossible inputs (the TS RangeError contract).
    """
    if n <= 0:
        raise ValueError("n must be > 0")
    if c < 0 or c > n:
        raise ValueError("c must be in [0, n]")
    if k <= 0 or k > n:
        raise ValueError("k must be in [1, n]")
    if n - c < k:
        return 1.0
    product = 1.0
    for i in range(k):
        product *= (n - c - i) / (n - i)
    return 1.0 - product


def precision_recall(counts: ConfusionCounts) -> PRF1:
    """Precision/recall/F1 over confusion counts. Zero denominators score 0.

    Raises ValueError on negative counts.
    """
    tp, fp, fn = counts.tp, counts.fp, counts.fn
    if any(v < 0 for v in (tp, fp, fn)):
        raise ValueError("counts must be >= 0")
    precision = tp / (tp + fp) if tp + fp > 0 else 0.0
    recall = tp / (tp + fn) if tp + fn > 0 else 0.0
    f1 = (2 * precision * recall) / (precision + recall) if precision + recall > 0 else 0.0
    return PRF1(precision, recall, f1)


def micro_average(per_class: Sequence[ConfusionCounts]) -> PRF1:
    """Micro-averaged P/R/F1 across per-class confusion counts."""
    sums = ConfusionCounts(
        tp=sum(c.tp for c in per_class),
        fp=sum(c.fp for c in per_class),
        fn=sum(c.fn for c in per_class),
    )
    return precision_recall(sums)


def exact_match_rate(predictions: Sequence[str], references: Sequence[str]) -> float:
    """Exact-match rate over paired predictions/references (case-sensitive).

    Empty input scores 0; length mismatch raises ValueError.
    """
    if len(predictions) != len(references):
        raise ValueError("predictions and references must have the same length")
    if not predictions:
        return 0.0
    hits = sum(1 for p, r in zip(predictions, references) if p == r)
    return hits / len(predictions)


def set_match(prediction: Iterable[str], reference: Iterable[str]) -> PRF1:
    """P/R/F1 over label SETS — the standard multi-label / extraction metric."""
    p = set(prediction)
    r = set(reference)
    tp = sum(1 for label in r if label in p)
    fp = len(p - r)
    fn = len(r - p)
    return precision_recall(ConfusionCounts(tp=tp, fp=fp, fn=fn))