causal_trace.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.
 
"""
Ingested from inspection-backend/causal_trace.py.
Import adaptation only: model_config -> aquin.compute.model_loader.
load_model() delegates to aquin.compute.model_loader (same instance as aquin load).
"""
from __future__ import annotations
 
from collections import OrderedDict
from typing import Any
 
import torch
from transformer_lens import HookedTransformer
 
from aquin.compute.model_loader import get_config, load_model, resolve_model_id
 
from aquin.compute.device import (
    default_dtype_for_device,
    resolve_compute_device,
    synchronize_device,
)
 
DEVICE = resolve_compute_device()
DTYPE = default_dtype_for_device(DEVICE)
 
 
def _project_unembed(model, resid: torch.Tensor) -> torch.Tensor:
    """Map residual vector to vocab logits (handles tied vs untied lm_head layouts)."""
    from aquin.compute.hf_llm_shim import project_residual_to_logits
 
    return project_residual_to_logits(model, resid)
 
 
def _format_prompt(model: HookedTransformer, prompt: str) -> str:
    if hasattr(model.tokenizer, "chat_template") and model.tokenizer.chat_template:
        messages = [{"role": "user", "content": prompt}]
        return model.tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
    return prompt
 
 
def _get_embeds(tokens, model_id: str = "llama-3.2-1b"):
    m = load_model(model_id)
    with torch.no_grad():
        return m.W_E[tokens]
 
 
def _corrupt_all(tokens, noise_scale, model_id: str = "llama-3.2-1b"):
    embeds = _get_embeds(tokens, model_id)
    return embeds + torch.randn_like(embeds) * noise_scale
 
 
def _corrupt_single_pos(tokens, pos, noise_scale, model_id: str = "llama-3.2-1b"):
    embeds = _get_embeds(tokens, model_id)
    noise = torch.zeros_like(embeds)
    noise[0, pos] = torch.randn(embeds.shape[-1], device=embeds.device) * noise_scale
    return embeds + noise
 
 
def run_trace(
    prompt: str,
    target_token: str,
    model_id: str = "llama-3.2-1b",
    noise_scale: float = 3.0,
    n_noise_runs: int = 10,
):
    m = load_model(model_id)
    tokens = m.to_tokens(prompt)
    target_id = m.to_tokens(target_token, prepend_bos=False)[0, 0].item()
 
    with torch.no_grad():
        clean_logits, clean_cache = m.run_with_cache(tokens)
    baseline_prob = torch.softmax(clean_logits[0, -1], dim=-1)[target_id].item()
 
    corrupted_prob = 0.0
    for _ in range(n_noise_runs):
        ne = _corrupt_all(tokens, noise_scale, model_id)
        with torch.no_grad():
            cl = m.run_with_hooks(tokens, fwd_hooks=[("hook_embed", lambda v, hook=None: ne.to(v.device))])
        corrupted_prob += torch.softmax(cl[0, -1], dim=-1)[target_id].item()
    corrupted_prob /= n_noise_runs
    corruption_effect = baseline_prob - corrupted_prob
 
    results = []
    for layer in range(m.cfg.n_layers):
        patch_probs = []
        for _ in range(n_noise_runs):
            ne = _corrupt_all(tokens, noise_scale, model_id)
            def make_hooks(ne_, li):
                return [
                    ("hook_embed", lambda v, hook=None, n=ne_: n.to(v.device)),
                    (f"blocks.{li}.hook_resid_post", lambda v, hook=None, li=li: (
                        v.__setitem__(
                            (slice(None), -1, slice(None)),
                            clean_cache[f"blocks.{li}.hook_resid_post"][:, -1, :].to(v.device)
                        ) or v
                    )),
                ]
            with torch.no_grad():
                pl = m.run_with_hooks(tokens, fwd_hooks=make_hooks(ne, layer))
            patch_probs.append(torch.softmax(pl[0, -1], dim=-1)[target_id].item())
 
        restored = sum(patch_probs) / len(patch_probs)
        drop = round(max(restored - corrupted_prob, 0.0) / max(corruption_effect, 1e-6), 4)
        results.append({
            "layer":             layer,
            "baseline_prob":     round(baseline_prob, 4),
            "corrupted_prob":    round(corrupted_prob, 4),
            "corruption_effect": round(corruption_effect, 4),
            "drop":              drop,
            "attn_drop":         0.0,
            "mlp_drop":          0.0,
            "patched_prob":      round(restored, 4),
            "attn_patched_prob": 0.0,
            "mlp_patched_prob":  0.0,
        })
 
    synchronize_device()
    return sorted(results, key=lambda x: x["drop"], reverse=True)
 
 
def run_prompt_attribution(
    prompt, response, prompt_tokens, response_tokens,
    sig_prompt_tis, sig_response_tis,
    model_id: str = "llama-3.2-1b",
    noise_scale: float = 3.0,
    n_noise_runs: int = 5,
):
    m = load_model(model_id)
    full_ctx = f"{prompt}\n{response}"
    tokens = m.to_tokens(full_ctx)
    n_tokens = tokens.shape[1]
 
    with torch.no_grad():
        clean_logits = m(tokens)
 
    target_ids, baseline_probs = {}, {}
    for rti in sig_response_tis:
        tok_str = f" {response_tokens[rti].strip()}"
        try:
            tid = m.to_tokens(tok_str, prepend_bos=False)[0, 0].item()
        except Exception:
            continue
        target_ids[rti] = tid
        baseline_probs[rti] = torch.softmax(clean_logits[0, -1], dim=-1)[tid].item()
 
    if not target_ids:
        return {"attributions": []}
 
    prefix_len = m.to_tokens(prompt + "\n", prepend_bos=True).shape[1]
 
    def get_token_pos(word_idx):
        words = prompt.split()
        if word_idx >= len(words):
            return None
        char_pos = sum(len(w) + 1 for w in words[:word_idx])
        try:
            return m.to_tokens(prompt[:char_pos], prepend_bos=True).shape[1] - 1
        except Exception:
            return None
 
    scores = {rti: {} for rti in target_ids}
    for pti in sig_prompt_tis:
        tok_pos = get_token_pos(pti)
        if tok_pos is None or tok_pos >= n_tokens:
            continue
 
        corrupted_probs = {rti: 0.0 for rti in target_ids}
        for _ in range(n_noise_runs):
            ne = _corrupt_single_pos(tokens, tok_pos, noise_scale, model_id)
            with torch.no_grad():
                logits = m.run_with_hooks(
                    tokens,
                    fwd_hooks=[("hook_embed", lambda v, hook=None, n=ne: n.to(v.device))]
                )
            probs = torch.softmax(logits[0, -1], dim=-1)
            for rti, tid in target_ids.items():
                corrupted_probs[rti] += probs[tid].item()
 
        for rti in target_ids:
            corrupted_probs[rti] /= n_noise_runs
            scores[rti][pti] = max(round(baseline_probs[rti] - corrupted_probs[rti], 4), 0.0)
 
    attributions = []
    for rti, tid in target_ids.items():
        raw = scores[rti]
        mx = max(raw.values(), default=1e-6) or 1e-6
        prompt_scores = sorted([
            {
                "prompt_ti":    pti,
                "prompt_token": prompt_tokens[pti],
                "score":        round(raw.get(pti, 0.0) / mx, 4),
                "raw_score":    raw.get(pti, 0.0),
            }
            for pti in sig_prompt_tis
        ], key=lambda x: x["score"], reverse=True)
        attributions.append({
            "response_ti":    rti,
            "response_token": response_tokens[rti],
            "baseline_prob":  round(baseline_probs[rti], 4),
            "prompt_scores":  prompt_scores,
        })
 
    synchronize_device()
    return {"attributions": attributions}
 
 
def run_logit_lens(
    prompt: str,
    model_id: str = "llama-3.2-1b",
    top_k: int = 5,
):
    m = load_model(model_id)
    tokens = m.to_tokens(prompt)
    with torch.no_grad():
        _, cache = m.run_with_cache(tokens)
 
    results = []
    for layer in range(m.cfg.n_layers):
        resid = cache[f"blocks.{layer}.hook_resid_post"][0, -1, :]
        resid_normed = m.ln_final(resid.unsqueeze(0).unsqueeze(0))[0, 0, :]
        logits = _project_unembed(m, resid_normed)
        probs = torch.softmax(logits, dim=-1)
        topk = probs.topk(top_k)
        results.append({
            "layer": layer,
            "top_tokens": [
                {"token": m.to_string([tid.item()]), "prob": round(p.item(), 4)}
                for tid, p in zip(topk.indices, topk.values)
            ],
        })
 
    synchronize_device()
    return results
 
 
def _resolve_chat_model(model_id: str, model: Any | None = None):
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id
 
    mid = resolve_model_id(model_id)
    m = model if model is not None else get_loaded_model()
    if m is None:
        m = load_model(mid)
    return m
 
 
def _generation_device(model: Any) -> torch.device:
    if hasattr(model, "hf_model"):
        embed = model.hf_model.get_input_embeddings()
        if embed is not None and hasattr(embed, "weight"):
            return embed.weight.device
    dev = getattr(model, "device", None)
    if dev is not None:
        return dev if isinstance(dev, torch.device) else torch.device(str(dev))
    return torch.device(DEVICE)
 
 
def _decode_chat_completion(model: Any, input_ids: torch.Tensor, output_ids: torch.Tensor) -> str:
    prompt_len = int(input_ids.shape[-1])
    new_tokens = output_ids[0, prompt_len:]
    return model.tokenizer.decode(new_tokens.tolist(), skip_special_tokens=True).strip()
 
 
def run_chat(
    prompt: str,
    model_id: str = "llama-3.2-1b",
    max_new_tokens: int = 200,
    temperature: float = 0.7,
    model: Any | None = None,
) -> str:
    m = _resolve_chat_model(model_id, model)
    formatted = _format_prompt(m, prompt)
    device = _generation_device(m)
    input_ids = m.tokenizer(formatted, return_tensors="pt").input_ids.to(device)
 
    if hasattr(m, "hf_model"):
        tok = m.tokenizer
        if tok.pad_token_id is None and tok.eos_token_id is not None:
            tok.pad_token_id = tok.eos_token_id
        pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
        gen_kwargs: dict[str, Any] = {
            "max_new_tokens": max_new_tokens,
            "do_sample": temperature > 0,
        }
        if pad_id is not None:
            gen_kwargs["pad_token_id"] = pad_id
        if temperature > 0:
            gen_kwargs["temperature"] = temperature
        with torch.no_grad():
            output_ids = m.hf_model.generate(input_ids, **gen_kwargs)
        return _decode_chat_completion(m, input_ids, output_ids)
 
    eos_id = getattr(m.tokenizer, "eos_token_id", None)
    generated: list[int] = []
 
    with torch.no_grad():
        cur = input_ids
        for _ in range(max_new_tokens):
            logits = m(cur)
            row = logits[0, -1]
            if temperature <= 0:
                next_id = int(row.argmax().item())
            else:
                probs = torch.softmax(row / temperature, dim=-1)
                next_id = int(torch.multinomial(probs, 1).item())
            if eos_id is not None and next_id == eos_id:
                break
            generated.append(next_id)
            cur = torch.cat(
                [cur, torch.tensor([[next_id]], device=device, dtype=cur.dtype)],
                dim=1,
            )
 
    return m.tokenizer.decode(generated, skip_special_tokens=True).strip()
 
 
def stream_chat(
    prompt: str,
    model_id: str = "llama-3.2-1b",
    max_new_tokens: int = 200,
    temperature: float = 0.7,
):
    m = load_model(model_id)
    formatted = _format_prompt(m, prompt)
    input_ids = m.tokenizer(formatted, return_tensors="pt").input_ids.to(DEVICE)
 
    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 tok_str
            cur = torch.cat([cur, torch.tensor([[next_id]], device=DEVICE)], dim=1)