Coverage for transformer_lens/model_bridge/sources/inspect/profiles.py: 76%
70 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""Per-provider request/response codecs for the Inspect driver.
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.
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
16import warnings
17from typing import Any, Mapping
19import numpy as np
21from . import hooks, intervention, wire
24class TLBridgeProfile:
25 """Codec for our own ``tl_bridge`` provider: residual/attn/mlp hooks, full-seq logits.
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 """
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
35 def __init__(
36 self,
37 supported_kinds: Any = None,
38 provides_sequence_logits: bool = True,
39 supports_attention_mask: bool = True,
40 ) -> None:
41 self._kinds = supported_kinds
42 self.provides_sequence_logits = provides_sequence_logits
43 self.supports_attention_mask = supports_attention_mask
45 def supported_hooks(self, n_layers: int) -> frozenset[str]:
46 return hooks.supported_hook_points(n_layers, self._kinds)
48 def translate_interventions(self, intervene, supported):
49 return intervention.build_interventions(intervene, supported) # {wire_key: spec}
51 def build_request(self, ids, wire_keys, interventions, return_logits, tokenizer):
52 extra: dict[str, Any] = {
53 "input_ids": ids,
54 "capture": wire_keys,
55 "return_logits": return_logits,
56 }
57 if interventions:
58 extra["interventions"] = interventions
59 return "", extra # our provider reads input_ids from extra_args; prompt unused
61 def decode_logits(self, output, n_tokens, d_vocab, tokenizer):
62 entry = (getattr(output, "metadata", None) or {}).get("tl_logits")
63 if entry is not None: 63 ↛ 67line 63 didn't jump to line 67 because the condition on line 63 was always true
64 return wire.decode_array(entry)[np.newaxis, ...] # full (1, seq, d_vocab)
65 # Absent tl_logits despite a logits request means the forward produced none — warn
66 # rather than silently hand back an all -inf tensor (whose argmax is a bogus token 0).
67 warnings.warn(
68 "tl_bridge provider returned no 'tl_logits' in output metadata; emitting an "
69 "all -inf placeholder (argmax would be token 0). Check the provider/wire path "
70 "rather than trusting these logits.",
71 RuntimeWarning,
72 stacklevel=2,
73 )
74 return np.full((1, n_tokens, d_vocab), -np.inf, dtype=np.float32)
77class VLLMLensProfile:
78 """Codec for the third-party vllm-lens provider.
80 Residual-stream-only, additive-steering-only, and last-token logits synthesized
81 one-hot from the generated token (argmax-only — vllm-lens doesn't hand back full
82 logits through this path). Prompt is the detokenized text, so vllm-lens
83 re-tokenizes it: activations reflect that re-tokenization, which may differ from
84 the exact ids. UNVERIFIED against a live provider.
85 """
87 provides_sequence_logits = False
88 supports_attention_mask = False
90 def supported_hooks(self, n_layers: int) -> frozenset[str]:
91 return frozenset(f"blocks.{i}.hook_out" for i in range(n_layers))
93 def translate_interventions(self, intervene: Mapping[str, Any], supported) -> list:
94 """op='add' with a width-shaped vector → vllm-lens SteeringVector; others raise.
96 Validates before importing ``vllm_lens`` so non-additive ops (and the
97 no-intervention case) don't require the package installed.
98 """
99 if not intervene:
100 return []
101 steering_cls: Any = None # imported lazily once a valid additive spec is seen
102 vectors = []
103 for name, spec in intervene.items(): 103 ↛ 127line 103 didn't jump to line 127 because the loop on line 103 didn't complete
104 if callable(spec): 104 ↛ 105line 104 didn't jump to line 105 because the condition on line 104 was never true
105 raise NotImplementedError("vllm-lens requires intervention specs, not callables.")
106 if not isinstance(spec, Mapping) or spec.get("op") != "add": 106 ↛ 111line 106 didn't jump to line 111 because the condition on line 106 was always true
107 raise NotImplementedError(
108 f"vllm-lens supports only additive steering (op='add' with a width vector); "
109 f"got {spec!r} for {name!r}. Use boot_transformers() for suppress/scale/set."
110 )
111 if name not in supported:
112 raise ValueError(f"Cannot intervene on {name!r}: not in supported_hook_points.")
113 resolved = hooks.resolve(name)
114 assert resolved is not None # supported ⇒ resolvable
115 if steering_cls is None:
116 from vllm_lens import SteeringVector
118 steering_cls = SteeringVector
119 vectors.append(
120 steering_cls(
121 activations=np.asarray(spec["value"], dtype=np.float32),
122 layer_indices=[resolved[0]],
123 scale=float(spec.get("scale", 1.0)),
124 norm_match=bool(spec.get("norm_match", True)),
125 )
126 )
127 return vectors
129 def build_request(self, ids, wire_keys, interventions, return_logits, tokenizer):
130 layers = sorted({int(key.split(":")[0]) for key in wire_keys})
131 extra: dict[str, Any] = {"output_residual_stream": layers}
132 if interventions: 132 ↛ 133line 132 didn't jump to line 133 because the condition on line 132 was never true
133 extra["apply_steering_vectors"] = interventions
134 prompt = tokenizer.decode(ids) if tokenizer is not None else ""
135 return prompt, extra
137 def decode_logits(self, output, n_tokens, d_vocab, tokenizer):
138 # vllm-lens returns no full logits here; one-hot the generated token (argmax only).
139 logits = np.full((1, n_tokens, d_vocab), -np.inf, dtype=np.float32)
140 text = getattr(output, "completion", "") or ""
141 ids = tokenizer.encode(text) if (tokenizer is not None and text) else []
142 if len(ids): 142 ↛ 144line 142 didn't jump to line 144 because the condition on line 142 was always true
143 logits[0, -1, int(ids[0])] = 0.0
144 return logits
147def for_provider(provider: str) -> Any:
148 """Pick the codec for a provider name (the part before ``/`` in get_model)."""
149 if provider.startswith("vllm-lens"):
150 return VLLMLensProfile()
151 if provider in ("tl_bridge", "tl_bridge_vllm"):
152 return TLBridgeProfile(supports_attention_mask=False)
153 # An unknown provider would otherwise get full-capability codec and NaN downstream.
154 raise ValueError(
155 f"No Inspect codec for provider {provider!r}. Known providers: 'tl_bridge', "
156 "'tl_bridge_vllm', 'vllm-lens*'."
157 )