steer.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""
Copied from inspection-backend/server.py steer logic.
Synchronous generator — yields SSE event dicts.
"""
from __future__ import annotations
 
from typing import Generator
 
import torch
 
from aquin.compute.device import resolve_compute_device
 
DEVICE = resolve_compute_device()
 
 
def run_steer_stream(
    prompt: str,
    feature_idx: int,
    strength: float = 20.0,
    model_id: str = "llama-3.2-1b",
    max_new_tokens: int = 200,
    temperature: float = 0.7,
) -> Generator[dict, None, None]:
    from aquin.compute.model_loader import get_config, load_model, resolve_model_id
    from aquin.compute.feature_analysis import _load_sae, _load_norm
    from aquin.compute.causal_trace import _format_prompt
 
    model_id = resolve_model_id(model_id)
    m = load_model(model_id)
    cfg = get_config(model_id)
    layer = cfg.get("sae_layer", 8)
 
    sae = _load_sae(model_id, layer)
    norm = _load_norm(model_id, layer)
 
    feat_dir = sae.W_dec[feature_idx]
    if norm is not None and norm.get("std") is not None:
        feat_dir = feat_dir * norm["std"].to(device=feat_dir.device, dtype=feat_dir.dtype)
    steer_vec = feat_dir.to(device=m.W_E.device, dtype=m.W_E.dtype)
    hook_name = f"blocks.{layer}.hook_resid_post"
 
    formatted = _format_prompt(m, prompt)
    input_ids = m.tokenizer(formatted, return_tensors="pt").input_ids.to(m.W_E.device)
 
    # Baseline pass
    with torch.no_grad():
        cur = input_ids
        for _ in range(max_new_tokens):
            logits = m(cur)
            next_logits = logits[0, -1] / max(temperature, 1e-6)
            probs = torch.softmax(next_logits, dim=-1)
            next_id = int(torch.multinomial(probs, 1).item())
            if next_id == m.tokenizer.eos_token_id:
                break
            tok_str = m.tokenizer.decode([next_id], skip_special_tokens=True)
            yield {"type": "baseline", "token": tok_str}
            cur = torch.cat([cur, torch.tensor([[next_id]], device=cur.device)], dim=1)
    yield {"type": "baseline_done"}
 
    def steer_hook(value, hook):
        value = value.clone()
        vec = steer_vec.to(device=value.device, dtype=value.dtype)
        value[:, -1, :] = value[:, -1, :] + strength * vec
        return value
 
    # Steered pass
    with torch.no_grad():
        cur = input_ids
        for _ in range(max_new_tokens):
            logits = m.run_with_hooks(cur, fwd_hooks=[(hook_name, steer_hook)])
            next_logits = logits[0, -1] / max(temperature, 1e-6)
            probs = torch.softmax(next_logits, dim=-1)
            next_id = int(torch.multinomial(probs, 1).item())
            if next_id == m.tokenizer.eos_token_id:
                break
            tok_str = m.tokenizer.decode([next_id], skip_special_tokens=True)
            yield {"type": "steered", "token": tok_str}
            cur = torch.cat([cur, torch.tensor([[next_id]], device=cur.device)], dim=1)
 
    yield {"type": "done"}