Coverage for transformer_lens/model_bridge/sources/inspect/profiles.py: 75%

1"""Per-provider request/response codecs for the Inspect driver. 

2 

3The driver speaks one internal shape; each provider speaks its own. A Profile 

4encapsulates everything provider-specific: which hooks it serves, whether it 

5returns full-sequence logits, how to phrase the ``generate`` request (prompt + 

6``extra_args``), how to translate interventions, and how to read logits back. 

7Torch-free (numpy only) so the driver stays torch-free. 

8 

9``tl_bridge`` is our own HF provider (``provider.py``). ``vllm-lens`` is the 

10third-party provider. **The vllm-lens codec is written from its documented API and 

11is NOT verified against a live provider** (CI has none — it needs their GPU-served 

12provider). It's isolated here so this is the single place to fix once validated. 

13""" 

14from __future__ import annotations 

15 

16import warnings 

17from typing import Any, Mapping 

18 

19import numpy as np 

20 

21from . import hooks, intervention, wire 

22 

23 

24class TLBridgeProfile: 

25 """Codec for our own ``tl_bridge`` provider: residual/attn/mlp hooks, full-seq logits. 

26 

27 ``supported_kinds`` is the provider's structurally-detected boundary set (e.g. 

28 ``resid_mid`` dropped for parallel/norm-variant archs); ``None`` exposes all boundaries. 

29 """ 

30 

31 # Class-level default — overridden per-instance via __init__ for backends (vLLM) that 

32 # only populate the gen position; lets RemoteBridge.forward reject loss/both there. 

33 provides_sequence_logits = True 

34 

35 def __init__(self, supported_kinds: Any = None, provides_sequence_logits: bool = True) -> None: 

36 self._kinds = supported_kinds 

37 self.provides_sequence_logits = provides_sequence_logits 

38 

39 def supported_hooks(self, n_layers: int) -> frozenset[str]: 

40 return hooks.supported_hook_points(n_layers, self._kinds) 

41 

42 def translate_interventions(self, intervene, supported): 

43 return intervention.build_interventions(intervene, supported) # {wire_key: spec} 

44 

45 def build_request(self, ids, wire_keys, interventions, return_logits, tokenizer): 

46 extra: dict[str, Any] = { 

47 "input_ids": ids, 

48 "capture": wire_keys, 

49 "return_logits": return_logits, 

50 } 

51 if interventions: 

52 extra["interventions"] = interventions 

53 return "", extra # our provider reads input_ids from extra_args; prompt unused 

54 

55 def decode_logits(self, output, n_tokens, d_vocab, tokenizer): 

56 entry = (getattr(output, "metadata", None) or {}).get("tl_logits") 

57 if entry is not None: 57 ↛ 61line 57 didn't jump to line 61 because the condition on line 57 was always true

58 return wire.decode_array(entry)[np.newaxis, ...] # full (1, seq, d_vocab) 

59 # Absent tl_logits despite a logits request means the forward produced none — warn 

60 # rather than silently hand back an all -inf tensor (whose argmax is a bogus token 0). 

61 warnings.warn( 

62 "tl_bridge provider returned no 'tl_logits' in output metadata; emitting an " 

63 "all -inf placeholder (argmax would be token 0). Check the provider/wire path " 

64 "rather than trusting these logits.", 

65 RuntimeWarning, 

66 stacklevel=2, 

67 ) 

68 return np.full((1, n_tokens, d_vocab), -np.inf, dtype=np.float32) 

69 

70 

71class VLLMLensProfile: 

72 """Codec for the third-party vllm-lens provider. 

73 

74 Residual-stream-only, additive-steering-only, and last-token logits synthesized 

75 one-hot from the generated token (argmax-only — vllm-lens doesn't hand back full 

76 logits through this path). Prompt is the detokenized text, so vllm-lens 

77 re-tokenizes it: activations reflect that re-tokenization, which may differ from 

78 the exact ids. UNVERIFIED against a live provider. 

79 """ 

80 

81 provides_sequence_logits = False 

82 

83 def supported_hooks(self, n_layers: int) -> frozenset[str]: 

84 return frozenset(f"blocks.{i}.hook_out" for i in range(n_layers)) 

85 

86 def translate_interventions(self, intervene: Mapping[str, Any], supported) -> list: 

87 """op='add' with a width-shaped vector → vllm-lens SteeringVector; others raise. 

88 

89 Validates before importing ``vllm_lens`` so non-additive ops (and the 

90 no-intervention case) don't require the package installed. 

91 """ 

92 if not intervene: 

93 return [] 

94 steering_cls: Any = None # imported lazily once a valid additive spec is seen 

95 vectors = [] 

96 for name, spec in intervene.items(): 96 ↛ 120line 96 didn't jump to line 120 because the loop on line 96 didn't complete

97 if callable(spec): 97 ↛ 98line 97 didn't jump to line 98 because the condition on line 97 was never true

98 raise NotImplementedError("vllm-lens requires intervention specs, not callables.") 

99 if not isinstance(spec, Mapping) or spec.get("op") != "add": 99 ↛ 104line 99 didn't jump to line 104 because the condition on line 99 was always true

100 raise NotImplementedError( 

101 f"vllm-lens supports only additive steering (op='add' with a width vector); " 

102 f"got {spec!r} for {name!r}. Use boot_transformers() for suppress/scale/set." 

103 ) 

104 if name not in supported: 

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

106 resolved = hooks.resolve(name) 

107 assert resolved is not None # supported ⇒ resolvable 

108 if steering_cls is None: 

109 from vllm_lens import SteeringVector 

110 

111 steering_cls = SteeringVector 

112 vectors.append( 

113 steering_cls( 

114 activations=np.asarray(spec["value"], dtype=np.float32), 

115 layer_indices=[resolved[0]], 

116 scale=float(spec.get("scale", 1.0)), 

117 norm_match=bool(spec.get("norm_match", True)), 

118 ) 

119 ) 

120 return vectors 

121 

122 def build_request(self, ids, wire_keys, interventions, return_logits, tokenizer): 

123 layers = sorted({int(key.split(":")[0]) for key in wire_keys}) 

124 extra: dict[str, Any] = {"output_residual_stream": layers} 

125 if interventions: 125 ↛ 126line 125 didn't jump to line 126 because the condition on line 125 was never true

126 extra["apply_steering_vectors"] = interventions 

127 prompt = tokenizer.decode(ids) if tokenizer is not None else "" 

128 return prompt, extra 

129 

130 def decode_logits(self, output, n_tokens, d_vocab, tokenizer): 

131 # vllm-lens returns no full logits here; one-hot the generated token (argmax only). 

132 logits = np.full((1, n_tokens, d_vocab), -np.inf, dtype=np.float32) 

133 text = getattr(output, "completion", "") or "" 

134 ids = tokenizer.encode(text) if (tokenizer is not None and text) else [] 

135 if len(ids): 135 ↛ 137line 135 didn't jump to line 137 because the condition on line 135 was always true

136 logits[0, -1, int(ids[0])] = 0.0 

137 return logits 

138 

139 

140def for_provider(provider: str) -> Any: 

141 """Pick the codec for a provider name (the part before ``/`` in get_model).""" 

142 if provider.startswith("vllm-lens"): 

143 return VLLMLensProfile() 

144 if provider in ("tl_bridge", "tl_bridge_vllm"): 

145 return TLBridgeProfile() 

146 # An unknown provider would otherwise get full-capability codec and NaN downstream. 

147 raise ValueError( 

148 f"No Inspect codec for provider {provider!r}. Known providers: 'tl_bridge', " 

149 "'tl_bridge_vllm', 'vllm-lens*'." 

150 )