model_diff.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/server.py (_run_model_diff function).
Extracted into its own module so train_simulate.py and inspect tools can import it
without pulling in FastAPI or the full server.
"""
from __future__ import annotations
 
import torch
 
from aquin.compute.model_loader import get_config, load_model, resolve_model_id
from aquin.compute.causal_trace import DTYPE, DEVICE
from aquin.compute.device import empty_device_cache
 
 
def _run_model_diff_tl(
    base_model,
    ft_model,
    prompts: list[str],
    n_prompts: int = 5,
) -> dict:
    """Compare base vs fine-tuned HookedTransformer models (already loaded)."""
    _default_prompts = [
        "The capital of France is",
        "Water boils at",
        "The square root of 144 is",
        "Albert Einstein was born in",
        "The speed of light is approximately",
        "The chemical formula for water is",
        "DNA stands for",
        "The largest planet in our solar system is",
    ]
    eval_prompts = (prompts if prompts else _default_prompts)[:n_prompts]
 
    def _gen_output(model, prompt: str, max_new: int = 60) -> str:
        try:
            if hasattr(model, "hf_model"):
                from aquin.compute.causal_trace import run_chat
                return run_chat(
                    prompt,
                    max_new_tokens=max_new,
                    temperature=0.0,
                    model=model,
                )
            tokens = model.generate(prompt, max_new_tokens=max_new, do_sample=False)
            if isinstance(tokens, str):
                return tokens[len(prompt):].strip()
            input_len = model.to_tokens(prompt).shape[1]
            return model.tokenizer.decode(tokens[0][input_len:].tolist(), skip_special_tokens=True).strip()
        except Exception:
            return ""
 
    base_outs = [_gen_output(base_model, p) for p in eval_prompts]
    ft_outs   = [_gen_output(ft_model, p)   for p in eval_prompts]
 
    def _word_overlap(a: str, b: str) -> float:
        wa, wb = set(a.lower().split()), set(b.lower().split())
        if not wa and not wb:
            return 1.0
        return len(wa & wb) / max(len(wa | wb), 1)
 
    per_prompt_overlap = [_word_overlap(a, b) for a, b in zip(base_outs, ft_outs)]
    per_prompt_drift   = [1.0 - s for s in per_prompt_overlap]
    mean_drift         = sum(per_prompt_drift) / max(len(per_prompt_drift), 1)
    consistency_score  = round(max(0.0, 1.0 - mean_drift), 4)
 
    len_ratios = [
        min(max(len(ft.split()), 1) / max(len(base.split()), 1), 2.0)
        for base, ft in zip(base_outs, ft_outs)
    ]
    mean_len_ratio    = sum(len_ratios) / max(len(len_ratios), 1)
    suppression_score = round(min(mean_len_ratio, 1.0), 4)
 
    if len(per_prompt_drift) > 1:
        mean_d = sum(per_prompt_drift) / len(per_prompt_drift)
        variance = sum((d - mean_d) ** 2 for d in per_prompt_drift) / len(per_prompt_drift)
        robustness_score = round(max(0.0, 1.0 - variance * 4), 4)
    else:
        robustness_score = consistency_score
 
    _DOMAIN_KEYWORDS = {
        "factual":   ["capital", "born", "formula", "stands for", "invented", "wrote", "founded"],
        "science":   ["boils", "speed of light", "square root", "photosynthesis", "dna", "atom", "planet", "chemical"],
        "reasoning": ["if", "therefore", "implies", "deduce", "conclude", "given that", "factorial", "function"],
        "language":  ["write", "explain", "describe", "list", "name", "define", "summarise", "translate"],
    }
 
    def _classify_prompt(p: str) -> str:
        pl = p.lower()
        for domain, kws in _DOMAIN_KEYWORDS.items():
            if any(kw in pl for kw in kws):
                return domain
        return "general"
 
    from collections import defaultdict as _dd
    cat_data: dict = _dd(list)
    for i, (base_o, ft_o) in enumerate(zip(base_outs, ft_outs)):
        cat = _classify_prompt(eval_prompts[i])
        drift = per_prompt_drift[i]
        length_r = len_ratios[i]
        score = round((1.0 - drift) * 0.7 + min(length_r, 1.0) * 0.3, 4)
        cat_data[cat].append(score)
 
    BASELINE = 0.75
    category_deltas = []
    for cat, scores in cat_data.items():
        mean_score = round(sum(scores) / len(scores), 4)
        delta = round(mean_score - BASELINE, 4)
        direction = "improved" if delta > 0.05 else "degraded" if delta < -0.05 else "unchanged"
        category_deltas.append({"category": cat, "score": mean_score, "delta": delta, "direction": direction, "n": len(scores)})
    category_deltas.sort(key=lambda x: abs(x["delta"]), reverse=True)
 
    max_drift_idx = max(range(len(per_prompt_drift)), key=lambda i: per_prompt_drift[i]) if per_prompt_drift else 0
    max_drift_prompt = {
        "prompt": eval_prompts[max_drift_idx] if eval_prompts else "",
        "base":   base_outs[max_drift_idx]    if base_outs    else "",
        "ft":     ft_outs[max_drift_idx]      if ft_outs      else "",
        "drift":  round(per_prompt_drift[max_drift_idx], 4) if per_prompt_drift else 0.0,
    }
 
    return {
        "consistencyScore": consistency_score,
        "suppressionScore": suppression_score,
        "robustnessScore":  robustness_score,
        "categoryDeltas":   category_deltas,
        "maxDriftPrompt":   max_drift_prompt,
        "baseOutputs":      [{"prompt": p, "output": base_outs[i]} for i, p in enumerate(eval_prompts)],
        "ftOutputs":        [{"prompt": p, "output": ft_outs[i]}   for i, p in enumerate(eval_prompts)],
        "promptsUsed":      eval_prompts,
    }
 
 
def _run_model_diff(base_model_id: str, ft_ckpt_path: str, prompts: list[str], n_prompts: int = 5) -> dict:
    from transformer_lens import HookedTransformer
 
    _HF_TO_KEY = {
        "meta-llama/Llama-3.2-1B":          "llama-3.2-1b",
        "meta-llama/Llama-3.2-1B-Instruct": "llama-3.2-1b",
        "meta-llama/Llama-3.2-3B":          "llama-3.2-3b",
        "meta-llama/Llama-3.2-3B-Instruct": "llama-3.2-3b",
    }
    try:
        model_key = resolve_model_id(base_model_id)
    except ValueError:
        model_key = _HF_TO_KEY.get(base_model_id)
    if model_key is None:
        raise ValueError(
            f"[model-diff] unsupported base model '{base_model_id}'. "
            f"Supported: {list(_HF_TO_KEY.keys())}"
        )
 
    base_model = load_model(model_key)
    hf_name = get_config(model_key)["hf_name"]
    state = torch.load(ft_ckpt_path, map_location=DEVICE, weights_only=True)
    if isinstance(state, dict) and "model_state_dict" in state:
        state = state["model_state_dict"]
 
    try:
        from transformers import AutoModelForCausalLM
        hf_ft = AutoModelForCausalLM.from_pretrained(hf_name, torch_dtype=DTYPE, device_map=DEVICE)
        hf_ft.load_state_dict(state, strict=False)
        ft_model = HookedTransformer.from_pretrained(hf_name, hf_model=hf_ft, dtype=DTYPE, device=DEVICE)
        del hf_ft
    except Exception:
        ft_model = HookedTransformer.from_pretrained(hf_name, dtype=DTYPE)
        ft_model.eval()
        ft_model.to(DEVICE)
        missing, unexpected = ft_model.load_state_dict(state, strict=False)
        if unexpected:
            print(f"[model-diff] {len(unexpected)} unexpected keys: {unexpected[:3]}", flush=True)
        if missing:
            print(f"[model-diff] {len(missing)} missing keys — ft weights partially applied", flush=True)
 
    ft_model.eval()
    result = _run_model_diff_tl(base_model, ft_model, prompts, n_prompts=n_prompts)
    del ft_model
    empty_device_cache()
    return result