Coverage for transformer_lens/model_bridge/sources/inspect/intervention.py: 83%

1"""Validate intervention specs and key them by ``<layer>:<kind>`` for the provider. 

2 

3Our HF provider applies interventions as forward-hook affine ops at the residual/ 

4attn/mlp boundaries, so the full vocabulary works. ``resid_mid`` is derived 

5(capture-only), so intervening on it is rejected. 

6""" 

7from __future__ import annotations 

8 

9from typing import Any, Mapping 

10 

11from . import hooks 

12 

13# suppress (→0), scale (factor), add (value), set (value). 

14SUPPORTED_OPS = frozenset({"suppress", "scale", "add", "set"}) 

15 

16 

17def build_interventions( 

18 intervene: Mapping[str, Any], 

19 supported_hook_points: frozenset[str], 

20) -> dict[str, dict[str, Any]]: 

21 """Validate specs and return ``{wire_key: spec}`` for the provider to apply. 

22 

23 Rejects callables (remote drivers take specs, not callbacks), bad ops, unknown or 

24 unsupported hooks, and the capture-only ``resid_mid``. 

25 """ 

26 out: dict[str, dict[str, Any]] = {} 

27 for hook_name, spec in intervene.items(): 

28 if callable(spec): 

29 raise NotImplementedError( 

30 "InspectDriver requires intervention specs (dict), not callables. " 

31 "Supported ops: suppress, scale (factor: float), add/set (value: scalar or " 

32 "width-shaped)." 

33 ) 

34 if not isinstance(spec, Mapping) or "op" not in spec: 34 ↛ 35line 34 didn't jump to line 35 because the condition on line 34 was never true

35 raise ValueError( 

36 f"Intervention spec for {hook_name!r} must be a dict with 'op' key; got {spec!r}" 

37 ) 

38 op = spec["op"] 

39 if op not in SUPPORTED_OPS: 

40 raise ValueError( 

41 f"Unsupported intervention op {op!r} for {hook_name!r}. " 

42 f"Supported: {sorted(SUPPORTED_OPS)}." 

43 ) 

44 if hook_name not in supported_hook_points: 

45 raise ValueError(f"Cannot intervene on {hook_name!r}: not in supported_hook_points.") 

46 resolved = hooks.resolve(hook_name) 

47 if resolved is None or resolved[1] not in hooks.INTERVENEABLE_KINDS: 

48 raise ValueError( 

49 f"Cannot intervene on {hook_name!r}: capture-only " 

50 f"(intervene on resid_pre/attn_out/mlp_out/resid_post or attn q/k/v/z instead)." 

51 ) 

52 if op == "scale" and "factor" not in spec: 52 ↛ 53line 52 didn't jump to line 53 because the condition on line 52 was never true

53 raise ValueError(f"Intervention {hook_name!r}: op='scale' requires 'factor' (float).") 

54 if op in ("add", "set") and "value" not in spec: 54 ↛ 55line 54 didn't jump to line 55 because the condition on line 54 was never true

55 raise ValueError( 

56 f"Intervention {hook_name!r}: op={op!r} requires 'value' (scalar or width-shaped)." 

57 ) 

58 pos = spec.get("pos") 

59 if pos is not None and not ( 59 ↛ 63line 59 didn't jump to line 63 because the condition on line 59 was never true

60 isinstance(pos, int) 

61 or (isinstance(pos, (list, tuple)) and all(isinstance(p, int) for p in pos)) 

62 ): 

63 raise ValueError( 

64 f"Intervention {hook_name!r}: 'pos' must be an int or list of ints " 

65 f"(sequence positions to patch); got {pos!r}." 

66 ) 

67 layer, kind = resolved 

68 out[hooks.wire_key(layer, kind)] = dict(spec) 

69 return out