code/Q-Prefer/validation/metrics.py

"""Exact pairwise metrics used by the Q-Prefer D2 evaluation pipeline."""

from __future__ import annotations


def sufficient_statistics(
    labels: list[int], margins: list[float], epsilon: float
) -> tuple[int, int, int, int, int]:
    """Return C, D, human-tie, model-tie, and both-tie counts."""
    consistent = discordant = human_tie = model_tie = both_tie = 0
    for label, margin in zip(labels, margins, strict=True):
        if label == 0 and abs(margin) <= epsilon:
            both_tie += 1
        elif label == 0:
            human_tie += 1
        elif abs(margin) <= epsilon:
            model_tie += 1
        elif label * margin > 0:
            consistent += 1
        else:
            discordant += 1
    return consistent, discordant, human_tie, model_tie, both_tie


def accuracy_from_statistics(
    consistent: int,
    discordant: int,
    human_tie: int,
    model_tie: int,
    both_tie: int,
) -> float:
    total = consistent + discordant + human_tie + model_tie + both_tie
    return (consistent + both_tie) / total if total else 0.0


def calc_accuracy_with_ties(labels: list[int], margins: list[float]) -> float:
    """Search the scalar tie threshold that maximizes tie-aware accuracy."""
    statistics = list(sufficient_statistics(labels, margins, -1.0))
    best = float("-inf")
    current_epsilon = -1.0
    for label, margin in sorted(zip(labels, margins, strict=True), key=lambda item: abs(item[1])):
        if label == 0 and abs(margin) < current_epsilon:
            statistics[4] -= 1
        elif label == 0:
            statistics[2] -= 1
        elif abs(margin) < current_epsilon:
            statistics[3] -= 1
        elif label * margin > 0:
            statistics[0] -= 1
        else:
            statistics[1] -= 1

        current_epsilon = abs(margin)
        if label == 0 and abs(margin) <= current_epsilon:
            statistics[4] += 1
        elif label == 0:
            statistics[2] += 1
        elif abs(margin) <= current_epsilon:
            statistics[3] += 1
        elif label * margin > 0:
            statistics[0] += 1
        else:
            statistics[1] += 1
        best = max(best, accuracy_from_statistics(*statistics))
    return best if margins else 0.0


def calc_accuracy_without_ties(labels: list[int], margins: list[float]) -> float:
    """Exclude human ties and evaluate the sign on decisive A/B examples."""
    consistent, discordant, _, model_tie, _ = sufficient_statistics(labels, margins, -1.0)
    denominator = consistent + discordant + model_tie
    return consistent / denominator if denominator else 0.0


def calc_accuracy_with_ties_fixed(
    labels: list[int], margins: list[float], epsilon: float = 0.0
) -> float:
    """Evaluate three-way accuracy at a fixed scalar tie threshold."""
    return accuracy_from_statistics(*sufficient_statistics(labels, margins, epsilon))


def calc_stats_fixed(
    labels: list[int], margins: list[float], epsilon: float = 0.0
) -> dict[str, int]:
    """Return named sufficient statistics at a fixed threshold."""
    keys = ("C", "D", "Th", "Tm", "Thm")
    return dict(zip(keys, sufficient_statistics(labels, margins, epsilon), strict=True))