"""Paired, source-stratified cluster bootstrap for token predictions.""" from __future__ import annotations from dataclasses import dataclass from typing import Iterable, Mapping, Sequence import numpy as np from baselines.structs import TokenPredictionRow from .metrics import average_precision, auroc @dataclass(frozen=True) class TokenObservation: support_key: tuple[int, str | None, int, int, int] score: float label: int source: str split_group: str row_index: int def _row_source_and_group(row: TokenPredictionRow) -> tuple[str, str]: source = row.example_metadata.source or str(row.metadata.get("source") or "unknown") extra = row.example_metadata.extra split_group = None if isinstance(extra, dict): for version in ("v5_3", "v5_2", "v5_1"): value = extra.get(version) if isinstance(value, dict) and value.get("split_group"): split_group = str(value["split_group"]) break split_group = split_group or extra.get("split_group") if not split_group: split_group = f"row:{row.index}" return str(source), str(split_group) def observations_from_predictions( rows: Iterable[TokenPredictionRow], ) -> list[TokenObservation]: observations = [] for row in rows: if ( row.labels is None or row.token_ids is None or row.token_indices is None or row.comparison_mask is None ): raise ValueError("schema-v2 labels/token IDs/indices/comparison mask required") source, group = _row_source_and_group(row) for token_id, token_index, score, label, keep in zip( row.token_ids, row.token_indices, row.scores, row.labels, row.comparison_mask, ): if not keep: continue observations.append( TokenObservation( (row.index, row.id, token_index, token_id, label), float(score), int(label), source, group, row.index, ) ) return observations def assert_common_support(runs: Mapping[str, Sequence[TokenObservation]]) -> list: if not runs: raise ValueError("at least one prediction run is required") supports = { name: sorted(observation.support_key for observation in observations) for name, observations in runs.items() } reference_name = next(iter(supports)) reference = supports[reference_name] for name, support in supports.items(): if support != reference: raise ValueError( f"prediction support mismatch: {name} differs from {reference_name}" ) return reference def weighted_average_precision(scores, labels, weights) -> float: scores = np.asarray(scores, dtype=float) labels = np.asarray(labels, dtype=int) weights = np.asarray(weights, dtype=float) if not (len(scores) == len(labels) == len(weights)) or len(scores) == 0: raise ValueError("weighted AP inputs must be non-empty and aligned") if len(set(labels.tolist())) != 2 or (weights <= 0).any(): raise ValueError("weighted AP requires both classes and positive weights") order = np.argsort(-scores, kind="mergesort") scores, labels, weights = scores[order], labels[order], weights[order] distinct = np.r_[np.flatnonzero(np.diff(scores)), len(scores) - 1] positive_weight = weights * (labels == 1) tp = np.cumsum(positive_weight)[distinct] total = np.cumsum(weights)[distinct] precision = tp / total recall = tp / positive_weight.sum() previous = np.r_[0.0, recall[:-1]] return float(np.sum((recall - previous) * precision)) def descriptive_metrics(observations: Sequence[TokenObservation]) -> dict: scores = [observation.score for observation in observations] labels = [observation.label for observation in observations] row_counts = {} for observation in observations: row_counts[observation.row_index] = row_counts.get(observation.row_index, 0) + 1 weights = [1.0 / row_counts[observation.row_index] for observation in observations] source_ap = {} for source in sorted({observation.source for observation in observations}): subset = [observation for observation in observations if observation.source == source] if len({observation.label for observation in subset}) == 2: source_ap[source] = average_precision( [observation.score for observation in subset], [observation.label for observation in subset], ) else: source_ap[source] = None finite_source_ap = [value for value in source_ap.values() if value is not None] return { "n_tokens": len(observations), "n_positive": labels.count(1), "prevalence": labels.count(1) / len(labels), "micro_average_precision": average_precision(scores, labels), "micro_auroc": auroc(scores, labels), "answer_balanced_average_precision": weighted_average_precision( scores, labels, weights ), "source_average_precision": source_ap, "source_macro_average_precision": float(np.mean(finite_source_ap)), } def paired_cluster_bootstrap( method_runs: Mapping[str, Sequence[Sequence[TokenObservation]]], *, primary_method: str, comparator_method: str, draws: int = 10_000, seed: int = 20260806, ) -> dict: """Bootstrap groups within source; average metrics across seeds per draw.""" if draws <= 0: raise ValueError("draws must be positive") flattened = { f"{method}:seed{index}": observations for method, runs in method_runs.items() for index, observations in enumerate(runs) } assert_common_support(flattened) reference = next(iter(flattened.values())) group_source = {} group_keys = {} for observation in reference: key = (observation.source, observation.split_group) group_source[key] = observation.source group_keys.setdefault(key, []).append(observation.support_key) by_source = {} for key, source in group_source.items(): by_source.setdefault(source, []).append(key) score_maps = { method: [ {observation.support_key: observation.score for observation in run} for run in runs ] for method, runs in method_runs.items() } label_map = {observation.support_key: observation.label for observation in reference} rng = np.random.default_rng(seed) deltas = [] skipped = 0 for _ in range(draws): sampled_support = [] for source in sorted(by_source): groups = sorted(by_source[source]) sampled = rng.choice(len(groups), size=len(groups), replace=True) for index in sampled: sampled_support.extend(group_keys[groups[int(index)]]) sampled_labels = [label_map[key] for key in sampled_support] if len(set(sampled_labels)) != 2: skipped += 1 continue metrics = {} for method, run_maps in score_maps.items(): per_seed = [ average_precision( [run_map[key] for key in sampled_support], sampled_labels ) for run_map in run_maps ] metrics[method] = float(np.mean(per_seed)) deltas.append(metrics[primary_method] - metrics[comparator_method]) if not deltas: raise ValueError("all bootstrap draws were degenerate") values = np.asarray(deltas, dtype=float) return { "estimand": ( f"mean-seed micro AP({primary_method}) - " f"mean-seed micro AP({comparator_method})" ), "draws_requested": draws, "draws_used": len(values), "draws_skipped_degenerate": skipped, "seed": seed, "delta_mean": float(values.mean()), "ci_95_percentile": [ float(np.quantile(values, 0.025)), float(np.quantile(values, 0.975)), ], "probability_delta_positive": float((values > 0).mean()), }