Coverage for transformer_lens/tools/analysis/jacobian_lens_causal_swap_benchmark.py: 97%
237 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""Causal coordinate-swap benchmark for ``JacobianLens.coordinate_patch_hooks``.
3Measures whether an anchored J-space coordinate edit installed live inside a forward pass
4via ``coordinate_patch_hooks`` causes a directional change in model output, under three
5controls: baseline-capability filtering (only intervene on prompts the model already
6answers correctly), displacement-matched random-atom controls (isolate "this concept
7mattered" from "any edit of similar magnitude would have mattered"), and exact
8Clopper-Pearson uncertainty on every reported rate. Each trial draws its control arm under
9several seeds and records every draw, so the control arm's own spread is visible rather
10than folded into a single number.
12This module is layered bottom-up: the model-free prompt corpus and rank/margin metric,
13baseline-capability filtering, displacement-matched control-token selection, a per-trial
14runner that wires the first three together with real ``coordinate_patch_hooks`` calls
15against a live model, and exact Clopper-Pearson confidence intervals plus a versioned JSON
16artifact schema that fingerprints both the protocol and the result blocks and embeds the
17corpus definition. Generating the frozen artifact consumed by
18``demos/Jacobian_Lens_Coordinate_Patch_Benchmark_Demo.ipynb`` is a separate entry point:
19``transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark_cli``.
20"""
22from __future__ import annotations
24import hashlib
25import itertools
26import json
27from collections import Counter
28from dataclasses import asdict, dataclass
29from pathlib import Path
30from typing import (
31 Any,
32 Container,
33 Dict,
34 Iterator,
35 List,
36 Literal,
37 Mapping,
38 MutableMapping,
39 Optional,
40 Sequence,
41 Tuple,
42)
44import torch
46from transformer_lens.tools.analysis.jacobian_lens import DEFAULT_K, JacobianLens
47from transformer_lens.tools.analysis.jacobian_lens_decomposition import (
48 JSpaceDecomposition,
49)
52@dataclass(frozen=True)
53class FunctionSpec:
54 """A templated prompt function evaluated over a shared set of concepts.
56 ``template`` takes a single ``{arg}`` placeholder, e.g. ``"The capital of {arg} is"``.
57 ``answers`` maps each concept to its answer word under this function, e.g.
58 ``{"France": "Paris"}``.
59 """
61 name: str
62 template: str
63 answers: Dict[str, str]
66@dataclass(frozen=True)
67class BenchmarkCorpus:
68 """A named set of concepts and the prompt functions evaluated over them."""
70 name: str
71 concepts: Sequence[str]
72 functions: Sequence[FunctionSpec]
75@dataclass(frozen=True)
76class PromptTrialSpec:
77 """One (function, ordered source/target concept pair) prompt instance."""
79 function: str
80 source: str
81 target: str
82 prompt: str
83 source_answer: str
84 target_answer: str
87def iter_prompt_trials(corpus: BenchmarkCorpus) -> Iterator[PromptTrialSpec]:
88 """Yields one spec per (function, ordered source/target concept pair).
90 Ordered pairs are every element of ``itertools.permutations(corpus.concepts, 2)``, the
91 same cross product Jacobian_Lens_Demo.ipynb's country benchmark already uses, now as
92 tested library code instead of a notebook cell.
93 """
94 for function in corpus.functions:
95 for source, target in itertools.permutations(corpus.concepts, 2):
96 yield PromptTrialSpec(
97 function=function.name,
98 source=source,
99 target=target,
100 prompt=function.template.format(arg=source),
101 source_answer=function.answers[source],
102 target_answer=function.answers[target],
103 )
106def corpus_definition(corpus: BenchmarkCorpus) -> Dict[str, Any]:
107 """Serializes a corpus into the JSON shape embedded in the protocol manifest.
109 The prompts and answers decide which prompts are baseline-capable and how every
110 condition is scored, so they belong inside the fingerprinted manifest rather than
111 travelling as a bare name. Key order is normalized so the same corpus always
112 serializes identically.
113 """
114 return {
115 "name": corpus.name,
116 "concepts": list(corpus.concepts),
117 "functions": [
118 {
119 "name": function.name,
120 "template": function.template,
121 "answers": dict(function.answers),
122 }
123 for function in corpus.functions
124 ],
125 }
128@dataclass(frozen=True)
129class AnswerMetrics:
130 """Rank/margin/tie metrics for one target token against one next-token logit vector."""
132 top1_token_id: int
133 target_rank: int
134 target_is_top1: bool
135 target_tied_for_top: bool
136 target_logit_margin: float
139def compute_answer_metrics(logits: torch.Tensor, target_token_id: int) -> AnswerMetrics:
140 """Computes rank/margin/tie metrics for ``target_token_id`` against ``logits``.
142 Ports Jacobian_Lens_Demo.ipynb's ``_target_metrics`` cell verbatim (arithmetic
143 unchanged, only renamed and restructured into a dataclass), so results stay directly
144 comparable with that notebook's already-reviewed success/rank definitions.
145 """
146 if logits.ndim != 1 or logits.numel() < 2:
147 raise ValueError("expected one-dimensional next-token logits")
148 if not 0 <= target_token_id < logits.shape[0]:
149 raise ValueError("target token id is outside the vocabulary")
150 if not torch.isfinite(logits).all():
151 raise ValueError("logits must be finite")
153 target_logit = logits[target_token_id]
154 top_logit = logits.max()
155 top1_token_id = int(logits.argmax().item())
156 top_logit_tie_count = int((logits == top_logit).sum().item())
157 competitors = torch.cat((logits[:target_token_id], logits[target_token_id + 1 :]))
158 return AnswerMetrics(
159 top1_token_id=top1_token_id,
160 target_rank=int((logits > target_logit).sum().item()) + 1,
161 target_is_top1=top1_token_id == target_token_id,
162 target_tied_for_top=bool(target_logit == top_logit) and top_logit_tie_count > 1,
163 target_logit_margin=float((target_logit - competitors.max()).item()),
164 )
167@dataclass(frozen=True)
168class BaselineRecord:
169 """A source prompt's own-answer metrics under the unperturbed baseline forward pass."""
171 function: str
172 source: str
173 prompt: str
174 metrics: AnswerMetrics
177def filter_baseline_capable(
178 baselines: Sequence[BaselineRecord],
179) -> Tuple[List[BaselineRecord], List[BaselineRecord]]:
180 """Splits baseline records into (capable, excluded) prompts.
182 A prompt is baseline-capable when the model's own deterministic argmax already matches
183 the source's answer (``metrics.target_is_top1``) and that answer is not tied for the
184 maximum (``metrics.target_tied_for_top``); only such prompts are eligible for later
185 intervention trials, so an edit's effect is never measured against a prompt the
186 unperturbed model already gets wrong. A tied-for-top baseline is treated as not capable:
187 ``argmax`` returns the lowest index on an exact tie, so admitting ties would let token-id
188 order decide whether a prompt enters the trial set. Order-preserving in both outputs;
189 never mutates ``baselines``.
190 """
191 capable = [
192 record
193 for record in baselines
194 if record.metrics.target_is_top1 and not record.metrics.target_tied_for_top
195 ]
196 excluded = [
197 record
198 for record in baselines
199 if not record.metrics.target_is_top1 or record.metrics.target_tied_for_top
200 ]
201 return capable, excluded
204def _control_generator(seed: int) -> torch.Generator:
205 """Builds the CPU generator used to pick a control token.
207 The generator is pinned to CPU rather than the dictionary's device. ``torch.randint``
208 infers its output device from the generator, so a device-local generator would make the
209 same seed select a different control token per device and the frozen artifact would stop
210 being reproducible. A CPU generator also avoids the device mismatch that raises on
211 accelerators, where ``torch.randint`` allocates on CPU while a device-local generator
212 expects its own device.
213 """
214 return torch.Generator(device="cpu").manual_seed(seed)
217def select_displacement_matched_control_token(
218 dictionary: torch.Tensor,
219 source_token_id: int,
220 target_token_id: int,
221 excluded_ids: Container[int],
222 active_support: Container[int],
223 *,
224 tolerance: float = 0.1,
225 seed: int = 0,
226) -> Optional[int]:
227 """Deterministically selects a displacement-matched control token id.
229 The real condition perturbs the activation by ``c_src * (a_target - a_source)``, so the
230 edit's size is set by ``||a_target - a_source||``. Matching atom norms leaves that size
231 free: two atoms of equal norm can sit at very different distances from ``a_source``, and
232 the control edit then runs at a different magnitude from the real one. A candidate token
233 id ``t`` therefore qualifies when
234 ``abs(||a_t - a_source|| - ||a_target - a_source||) <= tolerance * ||a_target - a_source||``
235 and ``t`` is not ``source_token_id``, not ``target_token_id``, not a member of
236 ``excluded_ids``, and not a member of ``active_support`` -- an atom already carrying the
237 source coordinate is not a clean control.
239 One qualifying candidate is picked with a seeded CPU ``torch.Generator`` so the same
240 ``seed`` always yields the same control token on every device. Returns ``None`` when no
241 candidate qualifies; the tolerance is never silently widened and selection never falls
242 back to the globally nearest atom. Returning ``None`` rather than raising lets the caller
243 record an empty pool as its own skip status instead of aborting the sweep.
244 """
245 if dictionary.ndim != 2:
246 raise ValueError(
247 f"dictionary must be 2-D [num_atoms, d_model], got shape {tuple(dictionary.shape)}"
248 )
249 atoms = dictionary.float()
250 source_atom = atoms[source_token_id]
251 target_displacement = float((atoms[target_token_id] - source_atom).norm())
252 displacements = (atoms - source_atom).norm(dim=1)
253 within_tolerance = (displacements - target_displacement).abs() <= (
254 tolerance * target_displacement
255 )
256 candidates = [
257 token_id
258 for token_id in range(dictionary.shape[0])
259 if token_id != source_token_id
260 and token_id != target_token_id
261 and token_id not in excluded_ids
262 and token_id not in active_support
263 and bool(within_tolerance[token_id])
264 ]
265 if not candidates:
266 return None
267 generator = _control_generator(seed)
268 pick = int(torch.randint(len(candidates), (1,), generator=generator).item())
269 return candidates[pick]
272def _resolve_answer_token_id(model: Any, word: str) -> int:
273 """Resolves a concept or answer word to the token id it maps to as a continuation.
275 Prepends a leading space so the id matches how the word tokenizes when it follows
276 other prompt text (e.g. " Paris", not "Paris"), the same convention the corpus
277 templates themselves rely on.
278 """
279 return int(model.to_single_token(f" {word}"))
282@dataclass(frozen=True)
283class ControlDraw:
284 """One seeded control-arm draw: the token it selected and how that condition scored."""
286 seed: int
287 token_id: int
288 metrics: AnswerMetrics
291@dataclass(frozen=True)
292class TrialResult:
293 """One ``(function, source, target, layer)`` causal-swap trial's full record.
295 ``control_token_id`` and ``control_target_metrics`` mirror the first entry of
296 ``control_draws`` so a single-draw reader sees the same fields it always did. The full
297 list carries the control arm's own spread across seeds, which one draw cannot measure.
298 """
300 function: str
301 source: str
302 target: str
303 layer: int
304 status: Literal["ok", "skipped_source_inactive", "skipped_no_control_token"]
305 baseline: AnswerMetrics
306 real_target_metrics: Optional[AnswerMetrics]
307 control_token_id: Optional[int]
308 control_target_metrics: Optional[AnswerMetrics]
309 control_draws: List[ControlDraw]
310 error: Optional[str]
313def run_causal_swap_trial(
314 lens: JacobianLens,
315 model: Any,
316 trial_spec: PromptTrialSpec,
317 layer: int,
318 *,
319 decomposition_cache: Optional[MutableMapping[Tuple[int, int, int], JSpaceDecomposition]] = None,
320 control_tolerance: float = 0.1,
321 control_seeds: Sequence[int] = (0,),
322 alpha: float = 1.0,
323 k: int = DEFAULT_K,
324) -> TrialResult:
325 """Runs one causal-swap trial: baseline, then real and control coordinate-patch conditions.
327 A baseline forward pass scores the prompt's own (unperturbed) source answer. A
328 displacement-matched control token is then selected from ``layer``'s lens-vector
329 dictionary, excluding the source token, the real target token, both prompts' answer
330 tokens, and every atom in the layer's active support. The real and control conditions
331 each install ``coordinate_patch_hooks`` at ``layer`` and position ``-1``, sharing one
332 ``decomposition_cache`` so only the first of the two performs the vocabulary-scale
333 decomposition; both are scored against the target's own answer token, so a real-vs-control
334 gap isolates "swapping toward this concept mattered" from "any edit of this magnitude
335 would have mattered."
337 The control arm is drawn once per entry in ``control_seeds``. A single draw leaves the
338 control arm's own variance unmeasured, so a reader cannot tell a real-vs-control gap from
339 draw variation; the per-seed results are carried on ``control_draws``. Every draw reuses
340 the same ``decomposition_cache``, so only the first pays the vocabulary-scale solve.
342 Two conditions produce a skip rather than a result. When the source is not in the layer's
343 active support, the trial is recorded with ``status="skipped_source_inactive"``. When no
344 displacement-matched control token survives the exclusions, the trial is recorded with
345 ``status="skipped_no_control_token"``. Both carry a diagnostic message in ``error``.
347 The active-support precondition is checked up front, against the same decomposition the
348 hooks consume, rather than inferred from a caught exception. ``coordinate_patch_hooks``
349 raises plain ``ValueError`` for several unrelated protocol faults (a same-id source and
350 target, an out-of-vocabulary answer, non-finite logits) and defines no narrower subclass
351 to discriminate on, so catching ``ValueError`` would file those as skips and drop them
352 from both denominators. Every other error therefore propagates unchanged.
353 ``coordinate_patch_hooks``'s own ``UserWarning``s (both the per-call install notice and
354 any solver-side conditioning warning) are not suppressed here and propagate to the caller
355 unchanged.
357 Args:
358 lens: The fitted lens.
359 model: The model to run trials against.
360 trial_spec: The prompt, source/target concepts, and their answer words.
361 layer: The single layer to patch at.
362 decomposition_cache: Shared cache passed to both the real and control
363 ``coordinate_patch_hooks`` calls. A fresh cache is used if omitted.
364 control_tolerance: Relative tolerance for the displacement-matched control token.
365 control_seeds: Seeds to draw the control arm under. Each seed selects its own
366 control token and contributes one :class:`ControlDraw`; the first seed's draw
367 is mirrored onto ``control_token_id`` and ``control_target_metrics``.
368 alpha: Interpolation strength forwarded to ``coordinate_patch_hooks``.
369 k: Sparse-solver upper bound forwarded to ``coordinate_patch_hooks``.
371 Returns:
372 The trial's :class:`TrialResult`.
374 Raises:
375 ValueError: If the source and target concepts resolve to the same token id, which
376 would make the coordinate patch a silent no-op, or if ``control_seeds`` is
377 empty.
378 """
379 if decomposition_cache is None:
380 decomposition_cache = {}
381 if not control_seeds:
382 raise ValueError("control_seeds must contain at least one seed")
383 tokens = model.to_tokens(trial_spec.prompt)
384 source_id = _resolve_answer_token_id(model, trial_spec.source)
385 target_id = _resolve_answer_token_id(model, trial_spec.target)
386 source_answer_id = _resolve_answer_token_id(model, trial_spec.source_answer)
387 target_answer_id = _resolve_answer_token_id(model, trial_spec.target_answer)
388 if source_id == target_id:
389 raise ValueError(
390 f"source and target resolve to the same token id {source_id}; a coordinate "
391 "patch would be a silent no-op"
392 )
394 with torch.no_grad():
395 baseline_logits = model(tokens)[0, -1].float()
396 baseline_metrics = compute_answer_metrics(baseline_logits, source_answer_id)
398 dictionary = lens.lens_vector_dictionary(model, layer)
399 decomposition = lens.decompose(model, trial_spec.prompt, layer=layer, position=-1, k=k)
400 # Seed the shared cache under the hook's own key so the first firing is a cache hit and
401 # the up-front solve is not repeated. The hook normalizes -1 against the activation's
402 # sequence length, which equals the tokenized prompt length for a single-example pass.
403 decomposition_cache[(layer, 0, tokens.shape[1] - 1)] = decomposition
404 if source_id not in decomposition.support.tolist():
405 return TrialResult(
406 function=trial_spec.function,
407 source=trial_spec.source,
408 target=trial_spec.target,
409 layer=layer,
410 status="skipped_source_inactive",
411 baseline=baseline_metrics,
412 real_target_metrics=None,
413 control_token_id=None,
414 control_target_metrics=None,
415 control_draws=[],
416 error=(
417 f"source token id {source_id} is not in layer {layer}'s active support "
418 f"for this prompt"
419 ),
420 )
421 control_token_ids: List[int] = []
422 for seed in control_seeds: 422 ↛ 451line 422 didn't jump to line 451 because the loop on line 422 didn't complete
423 token_id = select_displacement_matched_control_token(
424 dictionary,
425 source_id,
426 target_id,
427 excluded_ids={source_answer_id, target_answer_id},
428 active_support=decomposition.support.tolist(),
429 tolerance=control_tolerance,
430 seed=seed,
431 )
432 if token_id is None:
433 return TrialResult(
434 function=trial_spec.function,
435 source=trial_spec.source,
436 target=trial_spec.target,
437 layer=layer,
438 status="skipped_no_control_token",
439 baseline=baseline_metrics,
440 real_target_metrics=None,
441 control_token_id=None,
442 control_target_metrics=None,
443 control_draws=[],
444 error=(
445 "no displacement-matched control token within relative tolerance "
446 f"{control_tolerance} of ||a_target - a_source|| for seed {seed}"
447 ),
448 )
449 control_token_ids.append(token_id)
451 def _condition_metrics(condition_target_id: int) -> AnswerMetrics:
452 hooks = lens.coordinate_patch_hooks(
453 model,
454 source_id,
455 condition_target_id,
456 layers=[layer],
457 positions=[-1],
458 decomposition_cache=decomposition_cache,
459 k=k,
460 alpha=alpha,
461 )
462 with model.hooks(fwd_hooks=hooks), torch.no_grad():
463 condition_logits = model(tokens)[0, -1].float()
464 return compute_answer_metrics(condition_logits, target_answer_id)
466 real_metrics = _condition_metrics(target_id)
467 control_draws = [
468 ControlDraw(
469 seed=seed,
470 token_id=token_id,
471 metrics=_condition_metrics(token_id),
472 )
473 for seed, token_id in zip(control_seeds, control_token_ids)
474 ]
476 return TrialResult(
477 function=trial_spec.function,
478 source=trial_spec.source,
479 target=trial_spec.target,
480 layer=layer,
481 status="ok",
482 baseline=baseline_metrics,
483 real_target_metrics=real_metrics,
484 control_token_id=control_draws[0].token_id,
485 control_target_metrics=control_draws[0].metrics,
486 control_draws=control_draws,
487 error=None,
488 )
491def run_causal_swap_benchmark(
492 lens: JacobianLens,
493 model: Any,
494 corpus: BenchmarkCorpus,
495 layers: Sequence[int],
496 **trial_kwargs: Any,
497) -> Tuple[List[TrialResult], List[BaselineRecord]]:
498 """Runs the full causal-swap sweep: baseline filtering, then every surviving trial.
500 Each ``(function, source)`` prompt's baseline is computed once -- it does not depend on
501 ``layer`` -- and only prompts that survive :func:`filter_baseline_capable` proceed to
502 :func:`run_causal_swap_trial`, once per remaining ``layer``. Each trial gets its own fresh
503 ``decomposition_cache``: the cache key is ``(layer, batch_idx, position)``, which collides
504 across different prompts run as independent single-example forward passes, so a cache may
505 only be reused within one trial's real/control pair, never across trials.
507 Args:
508 lens: The fitted lens.
509 model: The model to run trials against.
510 corpus: The prompt corpus to sweep.
511 layers: Layers to sweep as an independent trial dimension.
512 **trial_kwargs: Forwarded to :func:`run_causal_swap_trial` (``control_tolerance``,
513 ``control_seeds``, ``alpha``, ``k``; ``decomposition_cache`` is not accepted here
514 since each trial always uses its own).
516 Returns:
517 ``(trials, excluded_baselines)``.
518 """
519 all_specs = list(iter_prompt_trials(corpus))
520 specs_by_prompt: Dict[Tuple[str, str], List[PromptTrialSpec]] = {}
521 baselines: List[BaselineRecord] = []
522 for spec in all_specs:
523 prompt_key = (spec.function, spec.source)
524 if prompt_key not in specs_by_prompt: 524 ↛ 538line 524 didn't jump to line 538 because the condition on line 524 was always true
525 specs_by_prompt[prompt_key] = []
526 tokens = model.to_tokens(spec.prompt)
527 with torch.no_grad():
528 baseline_logits = model(tokens)[0, -1].float()
529 source_answer_id = _resolve_answer_token_id(model, spec.source_answer)
530 baselines.append(
531 BaselineRecord(
532 function=spec.function,
533 source=spec.source,
534 prompt=spec.prompt,
535 metrics=compute_answer_metrics(baseline_logits, source_answer_id),
536 )
537 )
538 specs_by_prompt[prompt_key].append(spec)
540 capable, excluded = filter_baseline_capable(baselines)
541 capable_prompt_keys = {(record.function, record.source) for record in capable}
543 trials: List[TrialResult] = []
544 for layer in layers:
545 for prompt_key, specs in specs_by_prompt.items():
546 if prompt_key not in capable_prompt_keys:
547 continue
548 for spec in specs:
549 trials.append(
550 run_causal_swap_trial(
551 lens,
552 model,
553 spec,
554 layer,
555 **trial_kwargs,
556 )
557 )
558 return trials, excluded
561@dataclass(frozen=True)
562class SuccessRateInterval:
563 """An exact Clopper-Pearson confidence interval around a success rate.
565 ``n_trials`` and ``n_successes`` are the counts the interval was computed from, so a
566 reader can see the denominator behind the bounds rather than inferring it.
567 """
569 point_estimate: float
570 ci_low: float
571 ci_high: float
572 n_trials: int
573 n_successes: int
574 confidence: float
577# Bisection depth for the Clopper-Pearson bounds. Each step halves the bracket, so 200 steps
578# resolve the bound far below the precision any reported rate needs.
579_CLOPPER_PEARSON_STEPS = 200
582def _binomial_cdf(k: int, n: int, p: float) -> float:
583 """Exact ``P(X <= k)`` for ``X ~ Binomial(n, p)``.
585 Accumulates the pmf by recurrence (``term_{i} = term_{i-1} * (n - i + 1) / i * p / (1 - p)``)
586 rather than forming ``comb(n, i)`` and ``p ** i`` separately, which overflow for the
587 trial counts this benchmark reports.
588 """
589 if k < 0: 589 ↛ 590line 589 didn't jump to line 590 because the condition on line 589 was never true
590 return 0.0
591 if k >= n: 591 ↛ 592line 591 didn't jump to line 592 because the condition on line 591 was never true
592 return 1.0
593 if p <= 0.0: 593 ↛ 594line 593 didn't jump to line 594 because the condition on line 593 was never true
594 return 1.0
595 if p >= 1.0: 595 ↛ 596line 595 didn't jump to line 596 because the condition on line 595 was never true
596 return 0.0
597 term = (1.0 - p) ** n
598 total = term
599 odds = p / (1.0 - p)
600 for i in range(1, k + 1):
601 term *= (n - i + 1) / i * odds
602 total += term
603 return total
606def _clopper_pearson_bound(n_successes: int, n_trials: int, tail: float, *, upper: bool) -> float:
607 """Solves one Clopper-Pearson bound by bisecting the exact binomial CDF.
609 The lower bound solves ``P(X >= k) = tail`` and the upper bound solves ``P(X <= k) =
610 tail``. ``P(X >= k)`` rises with ``p`` while ``P(X <= k)`` falls, so the two bounds
611 bracket from opposite sides; reversing that collapses every interval to ``[0, 1]``.
612 """
613 if upper and n_successes == n_trials:
614 return 1.0
615 if not upper and n_successes == 0:
616 return 0.0
617 low, high = 0.0, 1.0
618 for _ in range(_CLOPPER_PEARSON_STEPS):
619 midpoint = (low + high) / 2.0
620 if upper:
621 # P(X <= k) decreases in p, so a probability above the tail puts the bound below
622 # the midpoint.
623 if _binomial_cdf(n_successes, n_trials, midpoint) > tail:
624 low = midpoint
625 else:
626 high = midpoint
627 else:
628 # P(X >= k) increases in p, so a probability above the tail puts the bound above
629 # the midpoint.
630 if 1.0 - _binomial_cdf(n_successes - 1, n_trials, midpoint) > tail:
631 high = midpoint
632 else:
633 low = midpoint
634 return (low + high) / 2.0
637def success_rate_ci(
638 successes: Sequence[bool],
639 *,
640 confidence: float = 0.95,
641) -> SuccessRateInterval:
642 """Computes an exact Clopper-Pearson confidence interval for a success rate.
644 A percentile bootstrap cannot express uncertainty about an all-failure or all-success
645 sample: every resample of an all-zero vector is all zeros, so the interval collapses to
646 ``[0, 0]`` and reports the same certainty for one trial as for a thousand. The exact
647 interval instead widens as the trial count shrinks, giving ``[0.0, 0.459]`` for zero
648 successes out of six.
650 Deterministic by construction: there is no resampling and no seed.
652 Raises:
653 ValueError: If ``successes`` is empty.
654 """
655 if len(successes) == 0:
656 raise ValueError("successes must be a non-empty sequence")
657 n_trials = len(successes)
658 n_successes = sum(1 for success in successes if success)
659 tail = (1.0 - confidence) / 2.0
660 return SuccessRateInterval(
661 point_estimate=n_successes / n_trials,
662 ci_low=_clopper_pearson_bound(n_successes, n_trials, tail, upper=False),
663 ci_high=_clopper_pearson_bound(n_successes, n_trials, tail, upper=True),
664 n_trials=n_trials,
665 n_successes=n_successes,
666 confidence=confidence,
667 )
670_REQUIRED_MANIFEST_FIELDS = (
671 "model_id",
672 "model_revision",
673 "lens_repo",
674 "lens_file",
675 "lens_revision",
676 "corpus_name",
677 "corpus",
678 "corpus_repo",
679 "corpus_path",
680 "corpus_revision",
681 "layers_swept",
682 "alpha",
683 "k",
684 "control_tolerance",
685 "control_seeds",
686 "success_definition",
687 "baseline_definition",
688 "rank_definition",
689)
692def build_protocol_manifest(**fields: Any) -> Dict[str, Any]:
693 """Assembles a protocol manifest, requiring the benchmark's fixed field set.
695 Requires at least :data:`_REQUIRED_MANIFEST_FIELDS`, the same key set (and, where they
696 overlap, the same string values) as ``Jacobian_Lens_Demo.ipynb``'s existing
697 ``protocol_manifest`` cell. Beyond that cell's fields the manifest carries the full
698 corpus definition (``corpus``) and the repo/path/revision triple it was taken from
699 (``corpus_repo``, ``corpus_path``, ``corpus_revision``), so the prompts and answers
700 that decide capability and scoring enter the fingerprint instead of a bare name.
702 Raises:
703 ValueError: If any required field is missing, naming the missing field(s).
704 """
705 missing = [name for name in _REQUIRED_MANIFEST_FIELDS if name not in fields]
706 if missing:
707 raise ValueError(f"protocol manifest is missing required field(s): {', '.join(missing)}")
708 return dict(fields)
711def fingerprint_manifest(manifest: Dict[str, Any]) -> str:
712 """Fingerprints a protocol manifest with the recipe ``Jacobian_Lens_Demo.ipynb`` uses."""
713 return hashlib.sha256(
714 json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode()
715 ).hexdigest()
718def results_fingerprint(artifact: Mapping[str, Any]) -> str:
719 """Fingerprints an artifact's result blocks with the manifest recipe.
721 ``protocol_fingerprint`` covers only the manifest, so an artifact whose trials,
722 excluded baselines, or confidence-interval blocks were rewritten would still load
723 clean and print the same provenance hash. Hashing those blocks separately catches
724 that tampering while leaving the manifest recipe and its stored value untouched.
725 """
726 return hashlib.sha256(
727 json.dumps(
728 {
729 "trials": artifact["trials"],
730 "excluded_baselines": artifact["excluded_baselines"],
731 "real_success_ci": artifact["real_success_ci"],
732 "control_success_ci": artifact["control_success_ci"],
733 },
734 sort_keys=True,
735 separators=(",", ":"),
736 ).encode()
737 ).hexdigest()
740SCHEMA_VERSION = 2
743def serialize_artifact(
744 manifest: Dict[str, Any],
745 trials: Sequence[TrialResult],
746 excluded_baselines: Sequence[BaselineRecord],
747 real_ci: SuccessRateInterval,
748 control_ci: SuccessRateInterval,
749) -> Dict[str, Any]:
750 """Assembles a JSON-serializable artifact dict from a benchmark run's results.
752 ``n_independent_prompts`` records how many distinct ``(function, source)`` prompts the
753 pooled rate rests on. Pooling trials as independent draws overstates the evidence when
754 several trials share one prompt, so the count travels with the rate.
756 ``layers_executed`` lists the sorted distinct layers that produced an ``ok`` trial, and
757 ``trial_status_counts`` counts trials per status, so a reader can see how many of the
758 swept layers contributed nothing and how many trials were skipped rather than executed.
759 """
760 artifact: Dict[str, Any] = {
761 "schema_version": SCHEMA_VERSION,
762 "protocol_manifest": manifest,
763 "protocol_fingerprint": fingerprint_manifest(manifest),
764 "trials": [asdict(trial) for trial in trials],
765 "excluded_baselines": [asdict(record) for record in excluded_baselines],
766 "real_success_ci": asdict(real_ci),
767 "control_success_ci": asdict(control_ci),
768 "n_independent_prompts": len(
769 {(trial.function, trial.source) for trial in trials if trial.status == "ok"}
770 ),
771 "layers_executed": sorted({trial.layer for trial in trials if trial.status == "ok"}),
772 "trial_status_counts": dict(Counter(trial.status for trial in trials)),
773 }
774 artifact["results_fingerprint"] = results_fingerprint(artifact)
775 return artifact
778_REQUIRED_ARTIFACT_FIELDS = (
779 "schema_version",
780 "protocol_manifest",
781 "protocol_fingerprint",
782 "trials",
783 "excluded_baselines",
784 "real_success_ci",
785 "control_success_ci",
786 "n_independent_prompts",
787 "layers_executed",
788 "trial_status_counts",
789 "results_fingerprint",
790)
793def load_artifact(path: Path) -> Dict[str, Any]:
794 """Reads and validates a frozen artifact produced by :func:`serialize_artifact`.
796 Validates that every required top-level key is present, that ``schema_version``
797 matches :data:`SCHEMA_VERSION`, that ``protocol_fingerprint`` matches a fresh
798 :func:`fingerprint_manifest` of the loaded ``protocol_manifest``, and that
799 ``results_fingerprint`` matches a fresh :func:`results_fingerprint` of the loaded
800 result blocks -- catching a hand-edited or corrupted artifact rather than trusting
801 the stored fingerprints blindly.
803 Raises:
804 ValueError: Naming the missing or mismatched field.
805 """
806 artifact = json.loads(Path(path).read_text())
807 missing = [name for name in _REQUIRED_ARTIFACT_FIELDS if name not in artifact]
808 if missing:
809 raise ValueError(f"artifact is missing required field(s): {', '.join(missing)}")
810 if artifact["schema_version"] != SCHEMA_VERSION:
811 raise ValueError(
812 f"artifact schema_version {artifact['schema_version']!r} does not match the "
813 f"expected {SCHEMA_VERSION!r}"
814 )
815 expected_fingerprint = fingerprint_manifest(artifact["protocol_manifest"])
816 if artifact["protocol_fingerprint"] != expected_fingerprint:
817 raise ValueError(
818 "artifact protocol_fingerprint does not match a freshly computed fingerprint of "
819 "its protocol_manifest (the manifest may have been hand-edited or corrupted)"
820 )
821 expected_results_fingerprint = results_fingerprint(artifact)
822 if artifact["results_fingerprint"] != expected_results_fingerprint:
823 raise ValueError(
824 "artifact results_fingerprint does not match a freshly computed fingerprint of "
825 "its trials, excluded_baselines, and confidence-interval blocks (they may have "
826 "been hand-edited or corrupted)"
827 )
828 return artifact