Coverage for transformer_lens/model_bridge/sources/inspect/source.py: 93%
60 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""``boot_inspect`` — wrap an ``inspect_ai`` provider in a RemoteBridge via InspectDriver."""
2from __future__ import annotations
4import logging
5import warnings
6from typing import Any, Optional
8import torch
10# Module import (not a name import): architecture_adapter_factory imports model_bridge,
11# which imports this module, so its names may not be defined yet at this point.
12from transformer_lens.factories import architecture_adapter_factory
13from transformer_lens.model_bridge.remote_bridge import RemoteBridge
14from transformer_lens.model_bridge.sources._bridge_builder import (
15 build_bridge_config_from_hf,
16 configure_tokenizer,
17 skip_tokenizer_for_modality,
18)
19from transformer_lens.model_bridge.sources._hf_format import (
20 determine_architecture_from_hf_config,
21)
22from transformer_lens.utilities.hf_utils import get_hf_token
24from . import profiles
25from .driver import InspectDriver
27# Providers that expose the structural self-check + capture wire format the InspectDriver
28# consumes. boot_inspect queries supported_kinds on these; others route via for_provider.
29_TL_BRIDGE_PROVIDERS = {"tl_bridge", "tl_bridge_vllm"}
32def boot_inspect(
33 model_name: str,
34 tokenizer: Optional[Any] = None,
35 dtype: Optional[torch.dtype] = None,
36 provider: str = "tl_bridge",
37 **inspect_kwargs: Any,
38) -> RemoteBridge:
39 """Boot a model via an ``inspect_ai`` provider and wrap it in a :class:`RemoteBridge`.
41 The driver is provider-agnostic: ``provider`` defaults to our own HF-backed
42 ``tl_bridge`` provider (residual/attn/mlp capture + full affine interventions +
43 full-sequence logits); ``"vllm-lens"`` targets a running vllm-lens vLLM provider
44 (residual-only, additive-steering-only) — wire-aligned with its documented format,
45 but not yet verified against a live provider.
47 Fireable hooks (``tl_bridge``, TransformerBridge-native names): ``blocks.{i}.hook_in``
48 (resid_pre) / ``ln2.hook_in`` (resid_mid) /
49 ``hook_out`` (resid_post) / ``attn.hook_out`` / ``mlp.hook_out``, plus the head-split
50 attention hooks where the structural probe finds them: ``attn.hook_q/k/v`` (pre-RoPE
51 projection outputs; separate-projection archs only — fused qkv gates them),
52 ``attn.hook_z`` (out-projection input), and ``attn.hook_pattern`` (post-softmax,
53 capture-only, eager attention required). The provider runs a structural self-check per
54 model and gates any boundary it can't serve faithfully: ``resid_mid`` for
55 parallel-residual or norm-variant blocks, ``attn_out``/``mlp_out`` when their submodule
56 isn't locatable (it warns when it gates one). ``embed``, ``ln_final``, and
57 ``attn.hook_attn_scores`` are always non-fireable — use ``boot_transformers()`` for
58 those.
60 For parity with ``boot_transformers`` the provider loads with the same dtype (fp32 by
61 default) and eager attention. Full-sequence logits ride on ``return_logits=True`` (the
62 default); pass ``return_logits=False`` to skip the (seq × d_vocab) payload for pure
63 activation capture (``run_with_cache`` keeps them since it returns logits).
65 ``tl_bridge`` accepts a binary, single-row ``attention_mask`` matching the input
66 length. Masked forwards use padding-aware positions where the model accepts 2-D
67 position IDs; model-owned position derivation is left intact. The vLLM Inspect
68 providers reject masks rather than silently ignore them.
69 """
70 from inspect_ai.model import get_model
71 from transformers import AutoConfig, AutoTokenizer
73 from . import ( # noqa: F401 — import registers @modelapi
74 transformers_provider as _provider,
75 )
77 hf_token = get_hf_token()
78 hf_config = AutoConfig.from_pretrained(model_name, token=hf_token)
79 # Shared resolution (not architectures[0]) handles architectures=None configs via
80 # model_type and rejects unsupported archs before the provider loads weights.
81 architecture = determine_architecture_from_hf_config(hf_config)
82 # Default fp32 to match boot_transformers (which loads/casts fp32 regardless of the
83 # config's native dtype); an explicit dtype still wins.
84 resolved_dtype = dtype if dtype is not None else torch.float32
86 bridge_config = build_bridge_config_from_hf(hf_config, architecture, model_name, resolved_dtype)
87 adapter = architecture_adapter_factory.ArchitectureAdapterFactory.select_architecture_adapter(
88 bridge_config
89 )
90 if tokenizer is None and not skip_tokenizer_for_modality(adapter.cfg):
91 tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
92 if tokenizer is not None:
93 # Match boot_transformers' tokenizer setup so to_tokens(str) is token-identical.
94 tokenizer = configure_tokenizer(tokenizer, adapter.cfg)
96 if provider == "tl_bridge":
97 # The provider's raw HF forward must match boot_transformers' load: same dtype,
98 # eager attention (TL forces eager — SDPA/flash diverge and accumulate with depth),
99 # and auth/remote-code so gated/custom models load at all.
100 inspect_kwargs["model_kwargs"] = _provider_model_kwargs(
101 dict(inspect_kwargs.get("model_kwargs", {})), adapter, resolved_dtype, hf_token
102 )
103 elif provider == "tl_bridge_vllm": 103 ↛ 111line 103 didn't jump to line 111 because the condition on line 103 was always true
104 # Otherwise the provider defaults to the HF-config dtype and bridge_config.dtype
105 # lies about what the engine actually loaded.
106 inspect_kwargs["dtype"] = resolved_dtype
108 # memoize=False: inspect_ai caches get_model by name, which would (a) return a stale
109 # model ignoring a changed dtype/kwargs on re-boot and (b) keep weights resident past
110 # close(). Each boot must honor its own args and own its model's lifecycle.
111 model = get_model(f"{provider}/{model_name}", memoize=False, **inspect_kwargs)
112 # Both TL-bridge providers (HF + vLLM) restrict their profile to the boundaries the
113 # provider's structural self-check / overlay found this model can serve; warn only
114 # if it gated something. Third-party providers (e.g. vllm-lens) route via for_provider.
115 if provider in _TL_BRIDGE_PROVIDERS: 115 ↛ 134line 115 didn't jump to line 134 because the condition on line 115 was always true
116 api = getattr(model, "api", None)
117 kinds = None
118 note = ""
119 # Default True for back-compat; vLLM provider sets False so RemoteBridge.forward
120 # rejects loss/both (otherwise loss over -inf earlier positions silently NaNs).
121 psl = True
122 if api is not None: 122 ↛ 126line 122 didn't jump to line 126 because the condition on line 122 was always true
123 kinds = api.supported_kinds() if hasattr(api, "supported_kinds") else None
124 note = api.capability_note() if hasattr(api, "capability_note") else ""
125 psl = bool(getattr(api, "provides_sequence_logits", True))
126 profile = profiles.TLBridgeProfile(
127 supported_kinds=kinds,
128 provides_sequence_logits=psl,
129 supports_attention_mask=bool(getattr(api, "supports_attention_mask", False)),
130 )
131 if note:
132 warnings.warn(note, UserWarning, stacklevel=2)
133 else:
134 profile = profiles.for_provider(provider)
135 driver = InspectDriver(model=model, adapter=adapter, tokenizer=tokenizer, profile=profile)
136 bridge = RemoteBridge(adapter=adapter, tokenizer=tokenizer, driver=driver)
137 _log_hook_summary(model_name, architecture, provider, driver)
138 return bridge
141def _provider_model_kwargs(
142 model_kwargs: dict[str, Any], adapter: Any, dtype: torch.dtype, hf_token: Optional[str]
143) -> dict[str, Any]:
144 """Load kwargs for the HF provider that mirror boot_transformers, so the provider's
145 raw forward matches the bridge. Caller-supplied keys win (setdefault)."""
146 model_kwargs.setdefault("torch_dtype", dtype)
147 # boot_transformers forces eager unless the adapter pins an implementation.
148 model_kwargs.setdefault(
149 "attn_implementation", getattr(adapter.cfg, "attn_implementation", None) or "eager"
150 )
151 if hf_token: 151 ↛ 153line 151 didn't jump to line 153 because the condition on line 151 was always true
152 model_kwargs.setdefault("token", hf_token)
153 return model_kwargs
156def _log_hook_summary(
157 model_name: str, architecture: str, provider: str, driver: InspectDriver
158) -> None:
159 log = logging.getLogger("transformer_lens.inspect")
160 fireable = sorted(driver.supported_hook_points)
161 log.info(
162 "Inspect source on %s (%s) via provider %r serves %d fireable hook(s).",
163 model_name,
164 architecture,
165 provider,
166 len(fireable),
167 )
170__all__ = ["boot_inspect"]