| """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()), |
| } |
|
|
