transformer_lens.model_bridge.remote_bridge module

Non-torch bridge: vLLM workers, Inspect remote providers.

class transformer_lens.model_bridge.remote_bridge.RemoteBridge(adapter: Any, tokenizer: Any, driver: Any)

Bases: BridgeCore, HookIntrospectionMixin

Bridge for backends with no local nn.Module (vLLM, Inspect).

No nn.Module parentage strips the torch-only surface; driver pre-declares supported_hook_points (no model to walk).

static boot_inspect(*args: Any, **kwargs: Any) → RemoteBridge

Boot a model via an inspect_ai provider. Returns a RemoteBridge wrapping an InspectDriver. Lazy import keeps remote_bridge inspect-agnostic. See sources.inspect.boot_inspect() for kwargs.

static boot_vllm(*args: Any, **kwargs: Any) → RemoteBridge

Boot a model via vLLM. Returns a RemoteBridge wrapping a VLLMDriver.

Mirrors TransformerBridge.boot_transformers. Lazy import so remote_bridge itself stays vLLM-agnostic — only callers of this method need vLLM installed. See sources.vllm.boot_vllm() for kwargs.

forward(input: Any = None, *, return_type: str | None = 'logits', loss_per_token: bool = False, labels: Any = None, **kwargs: Any) → Any

Tokenize → driver.forward → replay captures → finalize per return_type.

Explicit labels use shifted causal loss; remote encoder-decoder loss is unsupported.

run_with_cache(*args: Any, **kwargs: Any) → Any

Cache via driver captures; stop/start_at_layer and incl_bwd rejected.

run_with_hooks(input: Any, fwd_hooks: List[Tuple[str | Callable, Callable]] = [], bwd_hooks: List[Tuple[str | Callable, Callable]] = [], reset_hooks_end: bool = True, clear_contexts: bool = False, return_type: str | None = 'logits', stop_at_layer: int | None = None, start_at_layer: int | None = None, remove_batch_dim: bool = False, **kwargs: Any) → Any

Run with hooks. Remote fwd_hooks fire post-forward on captured activations (read-only) — they can’t alter the computation, so warn; use intervene= specs to mutate. bwd_hooks are unsupported (no backward).

to_tokens(input: Any, prepend_bos: bool | None = None, truncate: bool = True) → Any

Tokenize a string with the same BOS handling as TransformerBridge.

Mirrors cfg.default_prepend_bos / tokenizer_prepends_bos so boot_inspect(m).run_with_cache("text") matches boot_transformers(m) on the same string — a bare encode (no BOS) would silently diverge.