Coverage for transformer_lens/model_bridge/sources/inspect/eval.py: 93%

1"""Inspect solver for capturing TransformerLens activations during an eval. 

2 

3``capture_activations([...])`` is a solver you add to a Task's solver chain to harvest 

4activations alongside a behavioral eval (with a ``tl_bridge``-served model). Full 

5activations go to a per-sample side artifact (``output_dir/<sample_id>.npz``) for 

6probing/SAE; a compact ``reduce(...)`` summary lands in the sample store so 

7``inspect_ai.analysis.samples_df`` can correlate activation features with scores. 

8 

9Imports ``inspect_ai`` at module load (like ``provider.py``); the package ``__init__`` 

10exposes ``capture_activations`` lazily so importing the package stays inspect_ai-free. 

11""" 

12from __future__ import annotations 

13 

14import os 

15from typing import Any, Callable, Mapping, Optional, Sequence 

16 

17import numpy as np 

18from inspect_ai.model import GenerateConfig, get_model 

19from inspect_ai.solver import Generate, Solver, TaskState, solver 

20from inspect_ai.util import store 

21 

22from . import hooks, wire 

23 

24 

25def _default_reduce(activations: Mapping[str, np.ndarray]) -> dict[str, Any]: 

26 """Per-hook L2 norm + shape — small, JSON-able, queryable in samples_df.""" 

27 return { 

28 name: {"l2": float(np.linalg.norm(arr)), "shape": list(arr.shape)} 

29 for name, arr in activations.items() 

30 } 

31 

32 

33@solver 

34def capture_activations( 

35 capture: Sequence[str], 

36 output_dir: str = "tl_activations", 

37 reduce: Optional[Callable[[Mapping[str, np.ndarray]], dict[str, Any]]] = None, 

38 store_key: str = "tl_activations", 

39) -> Solver: 

40 """Capture ``capture`` hooks for each sample's current messages. 

41 

42 Writes full activations to ``output_dir/<sample_id>.npz`` and a ``reduce(...)`` summary 

43 (default: per-hook L2 + shape) to the sample store under ``store_key`` (+ ``_path``). 

44 Requires a ``tl_bridge``-served model; raises if the model returns no activations. 

45 Place before ``generate()`` to capture the prompt, after it to include the completion. 

46 """ 

47 # Resolve at construction so a multi-eval run from different CWDs doesn't scatter 

48 # artifacts (the solver runs later, possibly under a different working directory). 

49 output_dir = os.path.abspath(output_dir) 

50 reduce_fn = reduce or _default_reduce 

51 # TL hook names → provider wire keys (and back, to key the saved arrays by hook name). 

52 name_by_wire = {} 

53 for name in capture: 

54 resolved = hooks.resolve(name) 

55 if resolved is None: 

56 raise ValueError(f"capture_activations: {name!r} is not a fireable hook name.") 

57 name_by_wire[hooks.wire_key(*resolved)] = name 

58 wire_keys = list(name_by_wire) 

59 

60 async def solve(state: TaskState, generate: Generate) -> TaskState: 

61 # get_model() (no args) is the eval's active model; state.model is just its name. 

62 output = await get_model().generate( 

63 state.messages, 

64 config=GenerateConfig( 

65 extra_body={"extra_args": {"capture": wire_keys, "return_logits": False}} 

66 ), 

67 ) 

68 decoded = wire.decode_activations(getattr(output, "metadata", None), wire_keys) 

69 if not decoded: 69 ↛ 70line 69 didn't jump to line 70 because the condition on line 69 was never true

70 raise RuntimeError( 

71 "capture_activations got no activations back — the eval model must be a " 

72 "tl_bridge provider (e.g. model='tl_bridge/gpt2')." 

73 ) 

74 activations = {name_by_wire[wk]: arr for wk, arr in decoded.items()} 

75 os.makedirs(output_dir, exist_ok=True) 

76 # Epoch in the name — multi-epoch runs reuse sample_ids and would overwrite. 

77 epoch = getattr(state, "epoch", None) 

78 stem = f"{state.sample_id}_epoch{epoch}" if epoch is not None else str(state.sample_id) 

79 path = os.path.join(output_dir, f"{stem}.npz") 

80 # Explicit allow_pickle: plain float arrays never need pickle, and it keeps 

81 # the **activations unpack off numpy 2.3's typed keyword slot. 

82 np.savez_compressed(path, allow_pickle=False, **activations) 

83 store().set(store_key, reduce_fn(activations)) 

84 store().set(f"{store_key}_path", path) 

85 return state 

86 

87 return solve 

88 

89 

90def turn_activations(sample: Any) -> list[dict[str, np.ndarray]]: 

91 """Per-turn activations from an eval sample's model events, for a provider booted with 

92 ``capture=[...]`` (e.g. ``model_args={"capture": [...]}``). Returns one dict per model 

93 generation, in turn order — the activations of an agentic/multi-turn rollout. Every 

94 array gets a leading batch dim: boundaries are ``(1, seq, d_model)``, head-split 

95 q/k/v/z ``(1, seq, heads, d_head)``. 

96 """ 

97 turns = [] 

98 for event in getattr(sample, "events", []) or []: 

99 metadata = getattr(getattr(event, "output", None), "metadata", None) 

100 if not metadata or "activations" not in metadata: 

101 continue 

102 decoded = wire.decode_activations(metadata, list(metadata["activations"])) 

103 named = {} 

104 for wk, arr in decoded.items(): 

105 name = hooks.name_from_wire_key(wk) 

106 if name is None: 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true

107 continue 

108 # Rank-aware batch dim (mirrors driver._assemble_captures): boundary kinds 

109 # arrive rank-2, head-split kinds rank-3 — unsqueeze exactly once either way. 

110 batchless = hooks.WIRE_BATCHLESS_NDIM.get(wk.partition(":")[2], 2) 

111 named[name] = arr[np.newaxis, ...] if arr.ndim == batchless else arr 

112 if named: 112 ↛ 98line 112 didn't jump to line 98 because the condition on line 112 was always true

113 turns.append(named) 

114 return turns 

115 

116 

117def activations_column(store_key: str = "tl_activations", name: Optional[str] = None) -> Any: 

118 """A ``samples_df`` column surfacing :func:`capture_activations`'s reduction, for 

119 correlating activation features with scores:: 

120 

121 df = samples_df(logs, columns=[*SampleSummary, activations_column()]) 

122 """ 

123 from inspect_ai.analysis import SampleColumn 

124 

125 return SampleColumn(name or store_key, path=lambda s: s.store.get(store_key), full=True) 

126 

127 

128__all__ = ["activations_column", "capture_activations", "turn_activations"]