quantdiff 0.1.0rc1__py3-none-any.whl

quantdiff/metrics/logit.py

@@ -0,0 +1,214 @@
+"""Teacher-forced logit metrics: top-1 agreement and a lower bound on KL divergence.
+
+Servers return only the k most likely tokens at each position, so the full-vocabulary
+KL(P_ref || Q_cand) cannot be computed. Both distributions are instead collapsed onto a
+coarser partition: every token listed in both top-k lists keeps its own cell, the
+reference's tokens missing from the candidate's list share a cell A, and everything else
+shares a cell B. Merging outcomes can never increase KL divergence (data processing
+inequality), so KL on this partition is a lower bound on the true KL.
+
+The candidate's mass on A is not reported, but it is bounded twice over: an unlisted token
+cannot be more likely than the least likely listed one, so Q(A) <= |A| * min listed
+probability, and A lies outside the candidate's list, so Q(A) <= 1 - the listed mass.
+We take the Q(A) within those bounds that minimizes the partition KL, which keeps the
+result a lower bound while still penalizing a candidate whose top-k misses tokens the
+reference considers likely.
+"""
+
+from __future__ import annotations
+
+import math
+import sys
+from collections.abc import Collection, Hashable, Sequence
+from dataclasses import dataclass
+from typing import Final
+
+from quantdiff.types import (
+    LogitMetrics,
+    PromptLogit,
+    ReferenceTrace,
+    TokenProb,
+    TokenStep,
+    TopK,
+)
+
+MASS_FLOOR: Final = 1e-12
+"""Smallest mass a cell may hold, so a cell one side assigns ~0 cannot produce log(0)."""
+
+PERCENTILE: Final = 0.99
+
+
+def partition_kld(ref_top: TopK, cand_top: TopK, *, by_id: bool) -> float:
+    """Return a lower bound, in nats, on the full-vocabulary KL(P_ref || Q_cand).
+
+    Tokens are matched by id when `by_id` is True and every token carries an id,
+    otherwise by string. `cand_top` must not be empty.
+    """
+    if not cand_top:
+        raise ValueError("candidate top-k is empty; the KL bound needs at least one token")
+    use_ids = by_id and _all_have_ids(ref_top, cand_top)
+    ref_mass = _mass_by_token(ref_top, use_ids=use_ids)
+    cand_mass = _mass_by_token(cand_top, use_ids=use_ids)
+    shared = ref_mass.keys() & cand_mass.keys()
+    ref_only = ref_mass.keys() - shared
+
+    p_shared = math.fsum(ref_mass[key] for key in shared)
+    p_missing = math.fsum(ref_mass[key] for key in ref_only)
+    p_rest = max(1.0 - p_shared - p_missing, 0.0)
+    q_unlisted = _rest_mass([cand_mass[key] for key in shared])
+    # Tokens in A are absent from the candidate's whole list, so they share what that
+    # list leaves over, not merely what the shared tokens leave over.
+    q_free = _rest_mass(cand_mass.values())
+
+    # Q(A) that minimizes the A and B terms is proportional to P; clamp it to its bounds.
+    proportional = q_unlisted * p_missing / max(p_missing + p_rest, MASS_FLOOR)
+    q_missing = min(proportional, len(ref_only) * min(cand_mass.values()), q_free, q_unlisted)
+    q_rest = q_unlisted - q_missing
+
+    cells = [(ref_mass[key], cand_mass[key]) for key in shared]
+    cells += [(p_missing, q_missing), (p_rest, q_rest)]
+    divergence = math.fsum(_kl_term(p, q) for p, q in cells)
+    # Floors and float rounding can leave a tiny negative residue on near-equal inputs.
+    return max(divergence, 0.0)
+
+
+def top1_match(ref_step: TokenStep, cand_top: TopK, *, by_id: bool) -> bool:
+    """Return True if the candidate's most likely token is the reference's greedy choice."""
+    if not cand_top:
+        return False
+    best = max(cand_top, key=lambda token: token.logprob)
+    chosen = ref_step.chosen
+    if by_id and chosen.token_id is not None and best.token_id is not None:
+        return best.token_id == chosen.token_id
+    return best.token == chosen.token
+
+
+def logit_metrics(
+    traces: Sequence[ReferenceTrace],
+    candidate: Sequence[Sequence[TopK]],
+    *,
+    exact_token_ids: bool,
+) -> LogitMetrics:
+    """Aggregate top-1 agreement and partition KL over every teacher-forced position.
+
+    `candidate[j][i]` is the candidate's top-k at position i of `traces[j]`. An empty
+    candidate list means the server stopped instead of predicting a token (it put
+    end-of-sequence first), so it counts as a top-1 miss and is left out of the KL
+    statistics, which need a distribution. Positions where the reference list is empty
+    are skipped. `per_prompt` repeats the counts for each trace so two candidates can be
+    compared prompt by prompt. Raises ValueError if the candidate shape does not match the
+    traces.
+    """
+    _check_shapes(traces, candidate)
+    divergences: list[float] = []
+    per_prompt: list[PromptLogit] = []
+    for trace, cand_steps in zip(traces, candidate, strict=True):
+        prompt = _score_prompt(trace, cand_steps, by_id=exact_token_ids)
+        per_prompt.append(prompt.summary)
+        divergences.extend(prompt.divergences)
+
+    positions = sum(prompt.positions for prompt in per_prompt)
+    if positions == 0:
+        return LogitMetrics(
+            prompts=0,
+            positions=0,
+            top1_agreement=0.0,
+            kld_mean=0.0,
+            kld_p99=0.0,
+            kld_max=0.0,
+            exact_token_ids=exact_token_ids,
+            per_prompt=tuple(per_prompt),
+        )
+    ordered = sorted(divergences) or [0.0]
+    return LogitMetrics(
+        prompts=sum(prompt.positions > 0 for prompt in per_prompt),
+        positions=positions,
+        top1_agreement=sum(prompt.top1_matches for prompt in per_prompt) / positions,
+        kld_mean=math.fsum(ordered) / len(ordered),
+        kld_p99=_nearest_rank(ordered, PERCENTILE),
+        kld_max=ordered[-1],
+        exact_token_ids=exact_token_ids,
+        per_prompt=tuple(per_prompt),
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class _ScoredPrompt:
+    summary: PromptLogit
+    divergences: list[float]
+
+
+def _score_prompt(
+    trace: ReferenceTrace, cand_steps: Sequence[TopK], *, by_id: bool
+) -> _ScoredPrompt:
+    divergences: list[float] = []
+    matches = 0
+    positions = 0
+    for ref_step, cand_top in zip(trace.steps, cand_steps, strict=True):
+        if not ref_step.top:
+            continue
+        positions += 1
+        matches += top1_match(ref_step, cand_top, by_id=by_id)
+        if cand_top:
+            divergences.append(partition_kld(ref_step.top, cand_top, by_id=by_id))
+    kld_mean = math.fsum(divergences) / len(divergences) if divergences else None
+    summary = PromptLogit(trace.prompt_id, positions, matches, kld_mean)
+    return _ScoredPrompt(summary, divergences)
+
+
+def _check_shapes(traces: Sequence[ReferenceTrace], candidate: Sequence[Sequence[TopK]]) -> None:
+    if len(traces) != len(candidate):
+        raise ValueError(
+            f"got candidate scores for {len(candidate)} traces, expected {len(traces)}"
+        )
+    for trace, cand_steps in zip(traces, candidate, strict=True):
+        if len(trace.steps) != len(cand_steps):
+            raise ValueError(
+                f"trace {trace.prompt_id!r} has {len(trace.steps)} steps but the candidate "
+                f"scored {len(cand_steps)} positions"
+            )
+
+
+def _all_have_ids(*tops: TopK) -> bool:
+    return all(token.token_id is not None for top in tops for token in top)
+
+
+def _mass_by_token(top: TopK, *, use_ids: bool) -> dict[Hashable, float]:
+    """Map each token key to its probability, summing duplicates of the same key."""
+    masses: dict[Hashable, float] = {}
+    for token in top:
+        key = _token_key(token, use_ids=use_ids)
+        masses[key] = masses.get(key, 0.0) + _probability(token.logprob)
+    return masses
+
+
+def _token_key(token: TokenProb, *, use_ids: bool) -> Hashable:
+    return token.token_id if use_ids else token.token
+
+
+def _probability(logprob: float) -> float:
+    if math.isnan(logprob):
+        raise ValueError("logprob is NaN")
+    # Servers occasionally report logprobs a hair above 0 for near-certain tokens.
+    return math.exp(min(logprob, 0.0))
+
+
+def _rest_mass(listed: Collection[float]) -> float:
+    """Return the candidate mass outside `listed`, padded by the rounding error of the sum.
+
+    When a cap binds at a tiny true mass, even an error of a few ulps in 1 - sum is a large
+    relative error. Padding errs toward more candidate mass, which can only lower the KL.
+    """
+    rounding = (len(listed) + 1) * sys.float_info.epsilon
+    return min(max(1.0 - math.fsum(listed) + rounding, MASS_FLOOR), 1.0)
+
+
+def _kl_term(p: float, q: float) -> float:
+    if p <= 0.0:
+        return 0.0
+    return p * math.log(p / max(q, MASS_FLOOR))
+
+
+def _nearest_rank(ordered: Sequence[float], fraction: float) -> float:
+    rank = max(math.ceil(fraction * len(ordered)), 1)
+    return ordered[rank - 1]

quantdiff/metrics/tasks.py

@@ -0,0 +1,114 @@
+"""Per-case scoring for task suites and the summaries built from it."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from typing import Final
+
+from quantdiff.errors import SuiteError
+from quantdiff.metrics.codeexec import DEFAULT_TIMEOUT_SECONDS, run_code_case
+from quantdiff.metrics.jsonschema import extract_json, validate
+from quantdiff.metrics.toolcheck import check_tool_case
+from quantdiff.types import CaseOutcome, ChatResult, PerfMetrics, TaskCase, TaskKind, TaskMetrics
+
+MAX_LISTED_FAILURES: Final = 25
+SCORED_KINDS: Final[tuple[TaskKind, ...]] = ("json", "tools", "code")
+CODE_EXEC_DISABLED: Final = "code execution disabled (pass --allow-code-exec)"
+CHAT_SCORED_BY_AGREEMENT: Final = "scored by agreement"
+
+
+def evaluate_case(
+    case: TaskCase,
+    result: ChatResult,
+    *,
+    allow_code_exec: bool,
+    timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
+) -> CaseOutcome:
+    """Score one chat result against its case. Code cases are skipped unless allowed."""
+    if case.kind == "json":
+        return _check_json_case(case, result)
+    if case.kind == "tools":
+        return check_tool_case(case, result)
+    if case.kind == "code":
+        if not allow_code_exec:
+            return CaseOutcome(case.id, case.kind, passed=None, reason=CODE_EXEC_DISABLED)
+        return run_code_case(case, result.text, timeout_seconds=timeout_seconds)
+    return CaseOutcome(case.id, case.kind, passed=None, reason=CHAT_SCORED_BY_AGREEMENT)
+
+
+def summarize_tasks(outcomes: Sequence[CaseOutcome]) -> tuple[TaskMetrics, ...]:
+    """Return one TaskMetrics per scored kind present, in json, tools, code order."""
+    summaries = []
+    for kind in SCORED_KINDS:
+        group = [outcome for outcome in outcomes if outcome.kind == kind]
+        if not group:
+            continue
+        failures = [outcome for outcome in group if outcome.passed is False]
+        summaries.append(
+            TaskMetrics(
+                kind=kind,
+                total=len(group),
+                passed=sum(outcome.passed is True for outcome in group),
+                skipped=sum(outcome.passed is None for outcome in group),
+                failures=tuple(failures[:MAX_LISTED_FAILURES]),
+            )
+        )
+    return tuple(summaries)
+
+
+def perf_metrics(results: Sequence[ChatResult]) -> PerfMetrics:
+    """Aggregate generation speed and mean request latency over chat results.
+
+    When the server timed its own decoding for every result that generated more than one
+    token, `tokens_per_second` is total completion tokens over total decode time (the
+    token-weighted harmonic mean of the per-request rates) with source "server". That
+    excludes model load and prompt processing, so it is comparable across models. A
+    one-token answer has no decode step to time and is left out of that figure.
+
+    Otherwise it falls back to completion tokens over wall-clock request time with source
+    "wall_clock", which includes network overhead and prompt evaluation and is therefore
+    only comparable between candidates measured on the same machine with the same cases.
+    """
+    mean_latency = sum(result.seconds for result in results) / len(results) if results else None
+    server_rate = _server_decode_rate(results)
+    if server_rate is not None:
+        return PerfMetrics(server_rate, mean_latency, source="server")
+    timed = [
+        result for result in results if result.completion_tokens is not None and result.seconds > 0
+    ]
+    total_seconds = sum(result.seconds for result in timed)
+    tokens_per_second = (
+        sum(result.completion_tokens or 0 for result in timed) / total_seconds
+        if total_seconds > 0
+        else None
+    )
+    return PerfMetrics(tokens_per_second, mean_latency, source="wall_clock")
+
+
+def _server_decode_rate(results: Sequence[ChatResult]) -> float | None:
+    tokens = 0
+    decode_seconds: list[float] = []
+    for result in results:
+        count = result.completion_tokens
+        if count is None or count < 2:
+            continue
+        rate = result.decode_tokens_per_second
+        if not rate:
+            return None
+        tokens += count
+        decode_seconds.append(count / rate)
+    return tokens / math.fsum(decode_seconds) if decode_seconds else None
+
+
+def _check_json_case(case: TaskCase, result: ChatResult) -> CaseOutcome:
+    if case.json_schema is None:
+        raise SuiteError(f"json case {case.id!r} has no json_schema")
+    value, problem = extract_json(result.text)
+    if problem is not None:
+        return CaseOutcome(case.id, case.kind, passed=False, reason=problem)
+    errors = validate(value, case.json_schema)
+    if errors:
+        extra = f" (and {len(errors) - 1} more)" if len(errors) > 1 else ""
+        return CaseOutcome(case.id, case.kind, passed=False, reason=errors[0] + extra)
+    return CaseOutcome(case.id, case.kind, passed=True)

quantdiff/metrics/textsim.py

@@ -0,0 +1,66 @@
+"""Text normalization and similarity for comparing free-form chat answers."""
+
+from __future__ import annotations
+
+import difflib
+import math
+import re
+from collections.abc import Sequence
+from typing import Final
+
+from quantdiff.types import AgreementMetrics
+
+MAX_COMPARED_CHARS: Final = 4000
+"""SequenceMatcher is quadratic in the worst case, so long answers are truncated first."""
+
+_OPEN_TAG: Final = "<think>"
+_CLOSE_TAG: Final = "</think>"
+_LEADING_THINK: Final = re.compile(r"\A\s*<think>.*?</think>", re.DOTALL)
+
+
+def strip_reasoning(text: str) -> str:
+    """Remove a leading <think>...</think> block emitted by reasoning models.
+
+    Some chat templates open the block inside the prompt, so the answer starts with the
+    reasoning and only contains the closing tag; that prefix is removed too. An opening
+    tag that is never closed means generation stopped mid-reasoning, so nothing is left.
+    """
+    match = _LEADING_THINK.match(text)
+    if match:
+        return text[match.end() :]
+    if text.lstrip().startswith(_OPEN_TAG):
+        return ""
+    if _CLOSE_TAG in text and _OPEN_TAG not in text:
+        return text.split(_CLOSE_TAG, 1)[1]
+    return text
+
+
+def normalize(text: str) -> str:
+    """Strip reasoning, trim, and collapse every run of whitespace to one space."""
+    return " ".join(strip_reasoning(text).split())
+
+
+def similarity(a: str, b: str) -> float:
+    """Return a similarity ratio in [0, 1] between two answers after normalization."""
+    left = normalize(a)[:MAX_COMPARED_CHARS]
+    right = normalize(b)[:MAX_COMPARED_CHARS]
+    if left == right:
+        return 1.0
+    return difflib.SequenceMatcher(None, left, right, autojunk=False).ratio()
+
+
+def agreement_metrics(answers: Sequence[tuple[str, str, str]]) -> AgreementMetrics:
+    """Score (case id, reference answer, candidate answer) triples by exact match and mean
+    similarity, keeping each case's similarity for paired comparisons."""
+    if not answers:
+        return AgreementMetrics(cases=0, exact_match_rate=0.0, mean_similarity=0.0)
+    exact = sum(normalize(reference) == normalize(candidate) for _, reference, candidate in answers)
+    per_case = tuple(
+        (case_id, similarity(reference, candidate)) for case_id, reference, candidate in answers
+    )
+    return AgreementMetrics(
+        cases=len(answers),
+        exact_match_rate=exact / len(answers),
+        mean_similarity=math.fsum(score for _, score in per_case) / len(answers),
+        per_case=per_case,
+    )

quantdiff/metrics/toolcheck.py

@@ -0,0 +1,99 @@
+"""Scoring for tool-calling cases: right tool, valid arguments, expected values."""
+
+from __future__ import annotations
+
+from typing import Final
+
+from quantdiff.errors import SuiteError
+from quantdiff.metrics.jsonschema import validate
+from quantdiff.types import CaseOutcome, ChatResult, JSONValue, TaskCase, ToolCall
+
+_PREVIEW_CHARS: Final = 60
+
+
+def check_tool_case(case: TaskCase, result: ChatResult) -> CaseOutcome:
+    """Pass when the first tool call matches the case's expectations.
+
+    A case with no `expected_tool` passes only when the model answers without calling
+    any tool. Raises SuiteError if the expected tool is not among the case's tools.
+    """
+    reason = _failure_reason(case, result)
+    return CaseOutcome(case_id=case.id, kind=case.kind, passed=reason is None, reason=reason or "")
+
+
+def arguments_match(expected: JSONValue, actual: JSONValue) -> bool:
+    """Lenient equality for tool arguments.
+
+    Numbers compare numerically (1 equals 1.0, but never equals true), strings compare
+    after strip() and casefold(), lists compare element by element, and objects match
+    when every expected key is present with a matching value.
+    """
+    if isinstance(expected, bool) or isinstance(actual, bool):
+        return type(expected) is type(actual) and expected == actual
+    if isinstance(expected, (int, float)):
+        return isinstance(actual, (int, float)) and expected == actual
+    if isinstance(expected, str):
+        return isinstance(actual, str) and _fold(expected) == _fold(actual)
+    if isinstance(expected, list):
+        return (
+            isinstance(actual, list)
+            and len(expected) == len(actual)
+            and all(map(arguments_match, expected, actual))
+        )
+    if isinstance(expected, dict):
+        return isinstance(actual, dict) and all(
+            key in actual and arguments_match(value, actual[key]) for key, value in expected.items()
+        )
+    return expected is None and actual is None
+
+
+def _failure_reason(case: TaskCase, result: ChatResult) -> str | None:
+    expected_tool = case.expected_tool
+    if expected_tool is None:
+        return _direct_answer_failure(result)
+    if not result.tool_calls:
+        return f"made no tool call, expected {expected_tool}"
+
+    call = result.tool_calls[0]
+    if call.name != expected_tool:
+        return f"called {call.name}, expected {expected_tool}"
+    if call.arguments is None:
+        return f"{call.name} arguments are not a JSON object: {_preview(call.raw_arguments)}"
+    schema_errors = validate(call.arguments, _parameters_schema(case, expected_tool))
+    if schema_errors:
+        extra = f" (and {len(schema_errors) - 1} more)" if len(schema_errors) > 1 else ""
+        return f"{call.name} arguments violate the schema: {schema_errors[0]}{extra}"
+    return _argument_mismatch(call, case.expected_arguments or {})
+
+
+def _direct_answer_failure(result: ChatResult) -> str | None:
+    if result.tool_calls:
+        return f"called {result.tool_calls[0].name}, expected a direct answer"
+    return None
+
+
+def _parameters_schema(case: TaskCase, tool_name: str) -> dict[str, JSONValue]:
+    for tool in case.tools:
+        if tool.name == tool_name:
+            return tool.parameters
+    raise SuiteError(f"case {case.id!r} expects tool {tool_name!r} but does not define it")
+
+
+def _argument_mismatch(call: ToolCall, expected: dict[str, JSONValue]) -> str | None:
+    arguments = call.arguments or {}
+    for key, value in expected.items():
+        if key not in arguments:
+            return f"{call.name} is missing argument {key!r}"
+        if not arguments_match(value, arguments[key]):
+            actual = _preview(repr(arguments[key]))
+            return f"{call.name} argument {key!r} is {actual}, expected {_preview(repr(value))}"
+    return None
+
+
+def _fold(text: str) -> str:
+    return text.strip().casefold()
+
+
+def _preview(text: str) -> str:
+    flat = " ".join(text.split())
+    return flat if len(flat) <= _PREVIEW_CHARS else flat[: _PREVIEW_CHARS - 3] + "..."