Coverage for transformer_lens/model_bridge/sources/inspect/_provider_base.py: 85%

1"""Shared base for ``tl_bridge``-style Inspect providers (HF, vLLM). 

2 

3Subclasses load their model + tokenizer, run their structural self-check, and implement 

4backend-specific ``_generate_capture`` / ``_generate_eval`` — this base owns the 

5``generate()`` dispatch, message-to-ids rendering (chat template + tools), logprob 

6construction, and validation of the per-turn ``capture=[...]`` config. 

7 

8Subclass contract: 

9- Set ``self._tokenizer`` (HF-style ``AutoTokenizer``), ``self._device`` (str), 

10 ``self._kinds`` (frozenset of served boundary kinds), ``self._capability_note`` (str 

11 explaining any gating), and ``self._eval_capture`` (dict ``{wire_key: hook_name}``, 

12 built via :meth:`_parse_eval_capture`) before any ``generate()`` call. 

13- Implement ``_generate_capture(input, extra_args, config)`` (TL-driven single forward) and 

14 ``_generate_eval(input, config, tools)`` (multi-token chat generation). 

15""" 

16from __future__ import annotations 

17 

18import json 

19import re 

20import uuid 

21import warnings 

22from typing import Any, Mapping 

23 

24import torch 

25from inspect_ai.model import GenerateConfig, Logprob, ModelAPI, TopLogprob 

26from inspect_ai.tool import ToolCall 

27 

28from . import hooks 

29 

30 

31def _message_text(message: Any) -> str: 

32 """Text of an Inspect chat message — its ``.text`` (handles multimodal content), 

33 falling back to string content, else ''.""" 

34 text = getattr(message, "text", None) 

35 if isinstance(text, str): 35 ↛ 37line 35 didn't jump to line 37 because the condition on line 35 was always true

36 return text 

37 content = getattr(message, "content", None) 

38 return content if isinstance(content, str) else "" 

39 

40 

41def _tool_schema(tool: Any) -> dict[str, Any]: 

42 """Inspect ``ToolInfo`` → OpenAI-style function schema for ``apply_chat_template``.""" 

43 params = tool.parameters 

44 params = params.model_dump() if hasattr(params, "model_dump") else dict(params) 

45 return { 

46 "type": "function", 

47 "function": {"name": tool.name, "description": tool.description, "parameters": params}, 

48 } 

49 

50 

51# Tool-call blocks emitted by common instruct templates (Qwen/Hermes-style); the bare-JSON 

52# fallback covers models that emit a single {"name", "arguments"} object. 

53_TOOL_CALL_BLOCK = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.S) 

54 

55 

56def _parse_tool_calls(text: str) -> list[ToolCall] | None: 

57 """Best-effort parse of tool calls from a completion. Model-specific formats vary; this 

58 handles ``<tool_call>{json}</tool_call>`` blocks and a single bare ``{name, arguments}``.""" 

59 blocks = _TOOL_CALL_BLOCK.findall(text) 

60 if not blocks: 

61 match = re.search(r"\{.*\}", text, re.S) 

62 blocks = [match.group(0)] if match else [] 

63 calls = [] 

64 for block in blocks: 

65 try: 

66 obj = json.loads(block) 

67 except (ValueError, TypeError): 

68 continue 

69 name = obj.get("name") if isinstance(obj, dict) else None 

70 if not isinstance(name, str): 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true

71 continue 

72 args = obj.get("arguments") or obj.get("parameters") or {} 

73 calls.append( 

74 ToolCall( 

75 id=uuid.uuid4().hex[:8], 

76 function=name, 

77 arguments=args if isinstance(args, dict) else {}, 

78 ) 

79 ) 

80 return calls or None 

81 

82 

83# Generation-semantics GenerateConfig fields neither provider maps; warn (once per field 

84# per process) instead of silently ignoring them. 

85_UNSUPPORTED_GENERATE_FIELDS = ( 

86 "frequency_penalty", 

87 "presence_penalty", 

88 "logit_bias", 

89 "best_of", 

90 "num_choices", 

91) 

92_WARNED_UNSUPPORTED: set[str] = set() 

93 

94 

95def _warn_unsupported_config(config: GenerateConfig, provider: str) -> None: 

96 """Warn once per process per set-but-unsupported GenerateConfig field.""" 

97 new = [ 

98 field 

99 for field in _UNSUPPORTED_GENERATE_FIELDS 

100 if getattr(config, field, None) is not None and field not in _WARNED_UNSUPPORTED 

101 ] 

102 if new: 

103 _WARNED_UNSUPPORTED.update(new) 

104 warnings.warn( 

105 f"{provider}: ignoring unsupported GenerateConfig field(s) {new}.", 

106 UserWarning, 

107 stacklevel=3, 

108 ) 

109 

110 

111def _require_served(kind: str, served: frozenset[str], note: str, context: str) -> None: 

112 """Raise if ``kind`` was gated by the structural self-check — without this the eval path 

113 would silently return a derivation (e.g. ``resid_mid``) the driver path excludes.""" 

114 if kind not in served: 

115 raise ValueError( 

116 f"{context} requests kind {kind!r} which this model gated. {note} " 

117 f"Served kinds: {sorted(served)}." 

118 ) 

119 

120 

121def _require_interveneable(kind: str, served: frozenset[str], note: str, context: str) -> None: 

122 """``_require_served`` plus the capture-only gate — shared by every provider's 

123 intervention entry point so gated/capture-only kinds fail identically.""" 

124 from .hooks import INTERVENEABLE_KINDS 

125 

126 _require_served(kind, served, note, context) 

127 if kind not in INTERVENEABLE_KINDS: 

128 raise ValueError( 

129 f"{context}: kind {kind!r} is capture-only " 

130 f"(interveneable: {sorted(INTERVENEABLE_KINDS)})." 

131 ) 

132 

133 

134class _InspectModelAPIBase(ModelAPI): 

135 """Shared Inspect ModelAPI scaffolding for ``tl_bridge``-style providers. 

136 

137 Subclasses populate ``self._tokenizer``/``_device``/``_kinds``/``_capability_note``/ 

138 ``_eval_capture`` in ``__init__`` and implement the backend-specific generate paths. 

139 """ 

140 

141 # Attributes set by subclass __init__ before any generate() call (declared here so the 

142 # shared helpers' type-checking sees them; not initialized to avoid masking bugs). 

143 _tokenizer: Any 

144 _device: Any 

145 _kinds: frozenset 

146 _capability_note: str 

147 _eval_capture: dict[str, str] 

148 

149 # Class-level capability flag, read by source.py → TLBridgeProfile → InspectDriver → 

150 # RemoteBridge.forward to gate return_type ∈ {loss, both}. Must be set by every 

151 # subclass: True iff every position 0..n-1 of metadata['tl_logits'] holds real values 

152 # (not -inf padding). HF does a true forward → full logits; vLLM's sampler bypasses 

153 # lm_head and only the generated position has logprobs, so vLLM=False. 

154 provides_sequence_logits: bool 

155 

156 # --- subclass contract ------------------------------------------------------------- 

157 

158 def _generate_capture( 

159 self, input: Any, extra_args: Mapping[str, Any], config: GenerateConfig 

160 ) -> Any: 

161 """TL-driven single-forward capture (residual/attn/mlp boundaries + full logits).""" 

162 raise NotImplementedError 

163 

164 def _generate_eval(self, input: Any, config: GenerateConfig, tools: Any) -> Any: 

165 """Plain Inspect generation: chat input → multi-token completion + Logprobs + 

166 ModelUsage (+ parsed tool calls when ``tools`` is non-empty + per-turn capture).""" 

167 raise NotImplementedError 

168 

169 # --- shared API -------------------------------------------------------------------- 

170 

171 def supported_kinds(self) -> frozenset: 

172 """Boundary kinds this model is structurally able to serve.""" 

173 return self._kinds 

174 

175 def capability_note(self) -> str: 

176 """Human-readable reason for any gated boundary, or '' if all are served.""" 

177 return self._capability_note 

178 

179 async def generate(self, input, tools, tool_choice, config): # type: ignore[override] 

180 # Two callers: the TL driver (extra_args carries input_ids/capture/interventions — 

181 # single-forward activation capture) and a plain Inspect eval (chat messages, real 

182 # multi-token generation). Branch on whether a TL request is present. 

183 extra_args: Mapping[str, Any] = (config.extra_body or {}).get("extra_args", {}) 

184 if extra_args.get("input_ids") is not None or extra_args.get("capture"): 

185 return self._generate_capture(input, extra_args, config) 

186 return self._generate_eval(input, config, tools or []) 

187 

188 # --- shared helpers ---------------------------------------------------------------- 

189 

190 def _parse_eval_capture(self, model_args: dict[str, Any]) -> dict[str, str]: 

191 """Validate ``model_args["capture"]`` against the structural self-check (same 

192 protection as the driver path) and key by wire key. Returns ``{wire_key: name}``.""" 

193 eval_capture: dict[str, str] = {} 

194 for name in model_args.pop("capture", None) or []: 

195 resolved = hooks.resolve(name) 

196 if resolved is None: 

197 raise ValueError(f"capture={name!r} is not a fireable hook name.") 

198 _require_served(resolved[1], self._kinds, self._capability_note, f"capture={name!r}") 

199 eval_capture[hooks.wire_key(*resolved)] = name 

200 return eval_capture 

201 

202 def _messages_to_ids(self, input: Any, tools: Any = ()) -> Any: 

203 """Render Inspect chat messages (+ any ``tools``) to input ids — chat template 

204 when the tokenizer has one, else newline-joined message text (e.g. gpt2).""" 

205 if isinstance(input, str): 

206 messages = [{"role": "user", "content": input}] 

207 else: 

208 messages = [ 

209 {"role": getattr(m, "role", "user"), "content": _message_text(m)} for m in input 

210 ] 

211 template = getattr(self._tokenizer, "chat_template", None) 

212 if len(tools) and not template: 

213 raise NotImplementedError( 

214 f"tl_bridge: tool use needs a tool-aware chat template; {self.model_name} has " 

215 "none. Serve a tool-capable instruct model for agentic evals." 

216 ) 

217 if template: 217 ↛ 218line 217 didn't jump to line 218 because the condition on line 217 was never true

218 kwargs: dict[str, Any] = {"add_generation_prompt": True} 

219 if len(tools): 

220 kwargs["tools"] = [_tool_schema(t) for t in tools] 

221 token_ids = self._tokenizer.apply_chat_template(messages, **kwargs) 

222 else: 

223 token_ids = self._tokenizer("\n".join(m["content"] for m in messages)).input_ids 

224 return torch.tensor([list(token_ids)], device=self._device) 

225 

226 def _logprob_entry(self, token_id: int, step_logits: Any, top_n: Any) -> Logprob: 

227 """One token's log-prob (+ top-k alternatives) from a position's logits.""" 

228 lp = torch.log_softmax(step_logits.float(), dim=-1) 

229 top = [] 

230 if top_n: 230 ↛ 238line 230 didn't jump to line 238 because the condition on line 230 was always true

231 vals, idx = lp.topk(int(top_n)) 

232 top = [ 

233 TopLogprob( 

234 token=str(self._tokenizer.decode([int(i)])), logprob=float(v), bytes=None 

235 ) 

236 for v, i in zip(vals.tolist(), idx.tolist()) 

237 ] 

238 return Logprob( 

239 token=str(self._tokenizer.decode([int(token_id)])), 

240 logprob=float(lp[int(token_id)]), 

241 bytes=None, 

242 top_logprobs=top, 

243 )