Coverage for transformer_lens/model_bridge/sources/vllm/intervention_specs.py: 94%

1"""Intervention op vocabulary + spec validation shared by every producer path 

2(VLLMDriver, the worker RPCs — which also cover the Inspect vLLM provider). 

3 

4Adding a new op requires (a) listing it here, (b) allowing its keys in 

5``validate_spec``, (c) handling it in ``worker_extension._apply_intervention``. 

6""" 

7from __future__ import annotations 

8 

9from typing import Any, Mapping 

10 

11import torch 

12 

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

14 

15 

16def validate_spec(hook_name: str, spec: Any, *, width: int | None = None) -> dict: 

17 """Validate one declarative intervention spec; returns a plain-dict copy. 

18 

19 ``width`` enables the value-shape check where the caller knows the hook width. 

20 """ 

21 if callable(spec): 

22 raise NotImplementedError( 

23 "vLLM accepts intervention specs (dict), not callables. " 

24 "Supported ops: suppress, scale (factor: float), add (value: scalar or " 

25 "width-shaped), set (value: scalar or width-shaped)." 

26 ) 

27 if not isinstance(spec, Mapping) or "op" not in spec: 

28 raise ValueError( 

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

30 ) 

31 op = spec["op"] 

32 if op not in SUPPORTED_OPS: 

33 raise ValueError( 

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

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

36 ) 

37 if op == "scale" and "factor" not in spec: 

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

39 if op in ("add", "set") and "value" not in spec: 

40 raise ValueError( 

41 f"Intervention {hook_name!r}: op={op!r} requires 'value' " 

42 "(scalar or width-shaped tensor/list)." 

43 ) 

44 # Unknown keys must fail loud: a typo'd "position"/"positions" would otherwise 

45 # silently become a whole-sequence edit. 

46 allowed = {"op", "pos"} 

47 if op == "scale": 47 ↛ 48line 47 didn't jump to line 48 because the condition on line 47 was never true

48 allowed.add("factor") 

49 if op in ("add", "set"): 

50 allowed.add("value") 

51 extra = set(spec) - allowed 

52 if extra: 

53 raise ValueError( 

54 f"Intervention {hook_name!r}: unknown spec key(s) {sorted(extra)}; " 

55 f"allowed for op={op!r}: {sorted(allowed)}." 

56 ) 

57 value = spec.get("value") 

58 if value is not None and not isinstance(value, (int, float)): 

59 try: 

60 n_elements = int(torch.as_tensor(value).numel()) 

61 except (TypeError, ValueError, RuntimeError) as exc: 

62 raise ValueError( 

63 f"Intervention {hook_name!r}: 'value' must be a scalar or a " 

64 f"width-shaped tensor/list; got {type(value).__name__}." 

65 ) from exc 

66 if width is not None and n_elements != width: 

67 # A mis-shaped value would otherwise surface as a broadcast error 

68 # mid-forward — or broadcast along the wrong axis in square cases. 

69 raise ValueError( 

70 f"Intervention {hook_name!r}: 'value' has {n_elements} elements " 

71 f"but the hook width is {width}." 

72 ) 

73 pos = spec.get("pos") 

74 if pos is not None: 

75 valid_pos = isinstance(pos, int) and not isinstance(pos, bool) 

76 if not valid_pos and isinstance(pos, (list, tuple)): 

77 valid_pos = all(isinstance(p, int) and not isinstance(p, bool) for p in pos) 

78 if not valid_pos: 

79 raise ValueError( 

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

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

82 ) 

83 negative = [p for p in ([pos] if isinstance(pos, int) else pos) if p < 0] 

84 if negative: 

85 raise ValueError( 

86 f"Intervention {hook_name!r}: 'pos' must be non-negative; got {negative}." 

87 ) 

88 return dict(spec)