1706_scope/main.py

"""
SAE-based topic classification during generation
=====================================================

Pipeline:
  1. Load Qwen3 base model + per-layer TopK SAEs
  2. Load topic data, split train/test
  3. Generate from all prompts, storing hidden states at every layer & step
  4. Run SAEs on stored hidden states → per-sample firing rates per layer
  5. Discover best layer + top-K features via cross-topic variance
  6. Build binary feature vectors, train MLP, evaluate

Install (torch must be installed separately for your CUDA version):
    pip install transformers datasets huggingface_hub numpy sklearn pandas
    pip install flash-attn --no-build-isolation
"""

import argparse
import json
import os
from datetime import datetime
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from sklearn.metrics import precision_recall_fscore_support
from transformers import AutoModelForCausalLM, AutoTokenizer

# Local imports here
from models import TopKSAE, TopicClassifier
from utils import (
    build_dataset_df,
    build_labeled_set,
    filter_dead_samples,
    load_topic_data,
    print_results,
    split_data,
)


# ==========================================================================
# CONFIG
# ==========================================================================

DEFAULT_MODEL_SIZES = ["1.7B-Instruct"]

CONFIGS = {
    "0.6B": dict(
        model_name="Qwen/Qwen3-0.6B-Base",
        sae_repo="Qwen/SAE-Res-Qwen3-0.6B-Base-W32K-L0_100",
        top_k_sae=100,
        layer_stride=1,
    ),
    "1.7B": dict(
        model_name="Qwen/Qwen3-1.7B-Base",
        sae_repo="Qwen/SAE-Res-Qwen3-1.7B-Base-W32K-L0_100",
        top_k_sae=100,
        layer_stride=1,
    ),
    "1.7B-Instruct": dict(
        model_name="Qwen/Qwen3-1.7B",
        sae_repo="Qwen/SAE-Res-Qwen3-1.7B-Base-W32K-L0_100",
        top_k_sae=100,
        layer_stride=1,
    ),
    "8B": dict(
        model_name="Qwen/Qwen3-8B-Base",
        sae_repo="Qwen/SAE-Res-Qwen3-8B-Base-W64K-L0_100",
        top_k_sae=100,
        layer_stride=2,  # skip every other layer to save memory
    ),
}

N_POS = 5000  # max samples per topic
SEED = 42  # global RNG seed
BATCH_SIZE = 64  # generation batch size
N_TOP_FEATURES = 100  # SAE features selected per layer
TRAIN_RATIO = 0.8  # fraction of data used for training
MLP_HIDDEN = 64  # hidden units in the classifier
MLP_EPOCHS = 30  # training epochs
MLP_LR = 1e-3  # AdamW learning rate
MLP_BATCH_SIZE = 256  # mini-batch size during MLP training
MAX_GEN_TOKENS = 256  # max new tokens per generation
MAX_INPUT_TOKENS = 2048  # prompt truncation length
CHECKPOINT_DIR = "checkpoints_new"

# Path to prompts.jsonl from topic-data-gen; falls back to HF datasets if missing.
JSONL_DATA_PATH = Path("prompts.jsonl")

# HuggingFace dataset specs used when JSONL_DATA_PATH is absent.
DATASETS = {
    "medical": {
        "hf_id": "starmpcc/Asclepius-Synthetic-Clinical-Notes",
        "config": None,
        "split": "train",
        "columns": ["question"],
        "label": "Asclepius clinical questions",
    },
    "insurance": {
        "hf_id": "bitext/Bitext-insurance-llm-chatbot-training-dataset",
        "config": None,
        "split": "train",
        "columns": ["instruction"],
        "label": "Bitext Insurance questions",
    },
}


# ==========================================================================
# MODEL LOADING
# ==========================================================================


def configure_runtime_speed():
    """Enable TF32 and high-precision matmul for faster CUDA computation."""
    if torch.cuda.is_available():
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
    try:
        torch.set_float32_matmul_precision("high")
    except Exception:
        pass


def load_model(name, device, dtype):
    """Load a causal LM with the best available attention backend, return (model, tokenizer)."""
    tok = AutoTokenizer.from_pretrained(name, trust_remote_code=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    # left-padding keeps the real tokens right-aligned so generation starts at the correct position
    tok.padding_side = "left"

    model = None
    last_error = None
    for attn_impl in ("flash_attention_2", "sdpa", None):
        try:
            kwargs = {"torch_dtype": dtype, "trust_remote_code": True}
            if attn_impl is not None:
                kwargs["attn_implementation"] = attn_impl
            model = (
                AutoModelForCausalLM.from_pretrained(name, **kwargs).to(device).eval()
            )
            print(f"  attention: {attn_impl or 'default'}")
            break
        except Exception as exc:
            last_error = exc
            if attn_impl is not None:
                print(f"  attention {attn_impl} unavailable; falling back")

    if model is None:
        raise RuntimeError(f"failed to load model {name}") from last_error

    model.config.use_cache = True
    print(
        f"  model: {name} | {model.config.num_hidden_layers} layers, "
        f"d={model.config.hidden_size}"
    )
    return model, tok


def load_all_saes(repo, layers, expected_d_model, k, device):
    """Download and instantiate one TopKSAE per layer from a HuggingFace repo."""
    saes = {}
    d_sae = 0  # set inside loop; defined here so it's always bound after
    for layer in layers:
        print(f"  loading SAE layer {layer}/{max(layers)} ...", end="\r")
        path = hf_hub_download(repo_id=repo, filename=f"layer{layer}.sae.pt")
        # weights_only=True prevents arbitrary code execution from pickle payloads
        params = torch.load(path, map_location="cpu", weights_only=True)
        d_sae, d_model = params["W_enc"].shape
        if d_model != expected_d_model:
            raise ValueError(
                f"SAE layer {layer} in {repo} has d_model={d_model}, "
                f"but model hidden_size={expected_d_model}"
            )

        sae = TopKSAE(d_model, d_sae, k)
        with torch.no_grad():
            # cast to float32 — SAE checkpoints are often bfloat16/float16
            sae.W_enc.copy_(params["W_enc"].float())
            sae.b_enc.copy_(params["b_enc"].float())
        saes[layer] = sae.to(device).eval()

    print(f"  loaded {len(saes)} SAEs (d_sae={d_sae}, top_k={k})" + " " * 30)
    return saes


# ==========================================================================
# GENERATION + INLINE SAE FIRING
# ==========================================================================


@torch.inference_mode()
def generate_and_fire(texts, model, tok, saes, layers, device, batch_size=BATCH_SIZE):
    """Generate and accumulate SAE firing rates inline — no hidden state storage.

    Returns:
        fired_matrices: dict[layer] -> (n_samples, d_sae) float32 binary numpy array
        generations:    list of decoded generation strings, one per sample
        n_gen_steps:    int array of per-sample decode step counts (0 = dead sample)
    """
    n = len(texts)
    d_sae = next(iter(saes.values())).d_sae
    fired_matrices = {layer: np.zeros((n, d_sae), dtype=np.float32) for layer in layers}
    generations = []
    n_gen_steps = np.zeros(n, dtype=np.int32)
    n_batches = (n + batch_size - 1) // batch_size

    hooks = []
    captured = {}

    def make_hook(layer_idx):
        def _hook(_module, _input, output):
            captured[layer_idx] = output[0].float()

        return _hook

    for layer in layers:
        h = model.model.layers[layer].register_forward_hook(make_hook(layer))
        hooks.append(h)

    try:
        for batch_idx in range(n_batches):
            start = batch_idx * batch_size
            batch_texts = texts[start : start + batch_size]
            bs = len(batch_texts)
            print(f"    batch {batch_idx + 1}/{n_batches} ({bs} samples)", end="\r")

            templated = [
                tok.apply_chat_template(
                    [{"role": "user", "content": t}],
                    tokenize=False,
                    add_generation_prompt=True,
                    enable_thinking=False,
                )
                for t in batch_texts
            ]
            enc = tok(
                templated,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=MAX_INPUT_TOKENS,
            ).to(device)

            # accumulate binary fired flags on GPU, flush to CPU array at end of batch
            batch_fired = {
                layer: torch.zeros(bs, d_sae, dtype=torch.bool, device=device)
                for layer in layers
            }
            gen_tokens = [[] for _ in range(bs)]
            active = torch.ones(bs, dtype=torch.bool, device=device)

            # prefill
            captured.clear()
            out = model(**enc, use_cache=True)
            past_kv = out.past_key_values
            next_tokens = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)

            for _step in range(MAX_GEN_TOKENS):
                active = active & (next_tokens.squeeze(-1) != tok.eos_token_id)
                if not active.any():
                    break

                active_idx = active.nonzero(as_tuple=True)[0]
                for i in active_idx.tolist():
                    gen_tokens[i].append(next_tokens[i, 0].item())

                captured.clear()
                out = model(next_tokens, past_key_values=past_kv, use_cache=True)
                past_kv = out.past_key_values
                next_tokens = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)

                for layer in layers:
                    h = captured[layer]
                    if h.dim() == 3:
                        h = h[:, -1, :]
                    if h.shape[0] == 1 and bs > 1:
                        h = h.expand(bs, -1)
                    h_active = h[active_idx]
                    topk_idx = saes[layer].topk_indices(h_active)  # (n_active, k)
                    # scatter fired features into the per-sample bool mask
                    step_fired = torch.zeros(
                        len(active_idx), d_sae, dtype=torch.bool, device=device
                    )
                    step_fired.scatter_(1, topk_idx, True)
                    batch_fired[layer][active_idx] |= step_fired

            for layer in layers:
                fired_matrices[layer][start : start + bs] = (
                    batch_fired[layer].cpu().numpy().astype(np.float32)
                )

            for i in range(bs):
                n_gen_steps[start + i] = len(gen_tokens[i])
                generations.append(tok.decode(gen_tokens[i], skip_special_tokens=True))

            del past_kv, out, enc, batch_fired, gen_tokens
            if device.startswith("cuda"):
                torch.cuda.empty_cache()

    finally:
        for h in hooks:
            h.remove()

    print(f"    done ({n} samples, {len(layers)} layers)" + " " * 20)
    return fired_matrices, generations, n_gen_steps



# ==========================================================================
# FEATURE DISCOVERY
# ==========================================================================


def discover_features(firing_rates_by_topic, k=N_TOP_FEATURES, layer_override=None):
    """Find the layer + k features with highest cross-topic variance in firing rates."""
    topics = list(firing_rates_by_topic.keys())
    layers = list(firing_rates_by_topic[topics[0]].keys())
    layer_scores = {}

    if layer_override is not None:
        stacked = np.stack([firing_rates_by_topic[t][layer_override] for t in topics])
        var = stacked.var(axis=0)
        return layer_override, np.argsort(var)[::-1][:k], {}

    best_layer, best_score, best_feats = None, -1, None
    for layer in layers:
        stacked = np.stack([firing_rates_by_topic[t][layer] for t in topics])
        var = stacked.var(axis=0)
        top_k_var = np.sort(var)[::-1][:k]
        score = float(top_k_var.mean())
        layer_scores[layer] = score
        if score > best_score:
            best_score = score
            best_layer = layer
            best_feats = np.argsort(var)[::-1][:k]

    return best_layer, best_feats, layer_scores


# ==========================================================================
# MLP TRAINING & EVALUATION
# ==========================================================================


def train_mlp(features, labels, n_classes, device):
    """Train TopicClassifier with AdamW and cross-entropy, printing loss every epoch."""
    n_in = features.shape[1]
    mlp = TopicClassifier(n_in, n_classes, MLP_HIDDEN).to(device)
    # fused=True uses the CUDA fused kernel; requires CUDA and float params
    use_fused = device.startswith("cuda")
    opt = torch.optim.AdamW(mlp.parameters(), lr=MLP_LR, fused=use_fused)
    loss_fn = nn.CrossEntropyLoss()

    X = torch.tensor(features, dtype=torch.float32).to(device)
    y = torch.tensor(labels, dtype=torch.long).to(device)
    n = X.shape[0]

    loss_history = []
    mlp.train()
    for epoch in range(MLP_EPOCHS):
        perm = torch.randperm(n, device=device)
        total_loss, n_batches = 0.0, 0
        for s in range(0, n, MLP_BATCH_SIZE):
            idx = perm[s : s + MLP_BATCH_SIZE]
            loss = loss_fn(mlp(X[idx]), y[idx])
            opt.zero_grad()
            loss.backward()
            opt.step()
            total_loss += loss.item()
            n_batches += 1
        epoch_loss = total_loss / n_batches
        loss_history.append(epoch_loss)
        print(f"    epoch {epoch + 1}/{MLP_EPOCHS}  loss={epoch_loss:.4f}")

    mlp.eval()
    return mlp, loss_history


def compute_per_topic_metrics(preds, labels, topics):
    """Compute precision, recall, F1 per topic using sklearn, return as dict."""
    raw = precision_recall_fscore_support(
        labels, preds, labels=list(range(len(topics))), zero_division="warn"
    )
    # unpack as arrays so Pyright knows they're indexable
    precision, recall, f1, support = (np.asarray(x) for x in raw)
    metrics = {}
    for i, topic in enumerate(topics):
        metrics[topic] = {
            "precision": float(precision[i]),
            "recall": float(recall[i]),
            "f1": float(f1[i]),
            "n_test": int(support[i]),
        }
    return metrics


# ==========================================================================
# MAIN PIPELINE
# ==========================================================================


def run_for_model_size(model_size):
    """Run the full SAE topic-classification pipeline for one model size."""
    cfg = CONFIGS[model_size]
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16 if device.startswith("cuda") else torch.float32
    configure_runtime_speed()

    print("=" * 72)
    print(f"LOAD  Qwen3-{model_size} + Qwen-Scope SAEs on {device}")
    print("=" * 72)
    model, tok = load_model(cfg["model_name"], device, dtype)
    layers = list(
        range(0, model.config.num_hidden_layers, int(cfg.get("layer_stride", 1)))
    )
    saes = load_all_saes(
        cfg["sae_repo"],
        layers,
        model.config.hidden_size,
        cfg["top_k_sae"],
        device,
    )

    # -- Stage 1: load & split data ------------------------------------------
    print("\n" + "=" * 72)
    print("STAGE 1  load topic data")
    print("=" * 72)
    data = load_topic_data(DATASETS, N_POS, SEED, jsonl_path=JSONL_DATA_PATH)
    train_data, test_data = split_data(data, TRAIN_RATIO)
    topics = list(data.keys())
    for topic in topics:
        print(
            f"  {topic}: {len(train_data[topic])} train / {len(test_data[topic])} test"
        )

    # -- Stage 2: generate & fire SAEs inline (all layers) -------------------
    print("\n" + "=" * 72)
    print("STAGE 2  generate & accumulate SAE firing rates (all layers)")
    print("=" * 72)

    train_texts, train_labels = build_labeled_set(train_data, topics)
    test_texts, test_labels = build_labeled_set(test_data, topics)

    print(f"  generating from {len(train_texts)} train prompts ...")
    train_fired, train_gens, train_steps = generate_and_fire(
        train_texts, model, tok, saes, layers, device
    )

    print(f"  generating from {len(test_texts)} test prompts ...")
    test_fired, test_gens, test_steps = generate_and_fire(
        test_texts, model, tok, saes, layers, device
    )

    train_fired, train_labels, train_texts, train_gens, train_steps, n_dead_train = (
        filter_dead_samples(
            train_fired, train_labels, train_texts, train_gens, train_steps, "train"
        )
    )
    test_fired, test_labels, test_texts, test_gens, test_steps, n_dead_test = (
        filter_dead_samples(
            test_fired, test_labels, test_texts, test_gens, test_steps, "test"
        )
    )

    # -- Stage 3: feature discovery from firing matrices ----------------------
    print("\n" + "=" * 72)
    print("STAGE 3  feature discovery")
    print("=" * 72)

    firing_rates_by_topic = {t: {} for t in topics}
    for layer in layers:
        for i, topic in enumerate(topics):
            mask = train_labels == i
            firing_rates_by_topic[topic][layer] = train_fired[layer][mask].mean(axis=0)

    best_layer, best_feats, layer_scores = discover_features(firing_rates_by_topic)
    assert best_feats is not None, (
        "discover_features returned no features (empty layers?)"
    )
    print(f"  best layer: L{best_layer}")
    print(f"  selected {len(best_feats)} features: {best_feats.tolist()}")

    # -- Stage 4: build feature vectors & train MLP ---------------------------
    print("\n" + "=" * 72)
    print("STAGE 4  train & evaluate MLP")
    print("=" * 72)

    train_feat_vectors = train_fired[best_layer][:, best_feats]
    test_feat_vectors = test_fired[best_layer][:, best_feats]

    print(
        f"  feature matrix: train {train_feat_vectors.shape}, test {test_feat_vectors.shape}"
    )

    print("  training MLP ...")
    mlp, loss_history = train_mlp(train_feat_vectors, train_labels, len(topics), device)

    with torch.no_grad():
        X_test = torch.tensor(test_feat_vectors, dtype=torch.float32).to(device)
        preds = mlp(X_test).argmax(dim=-1).cpu().numpy()

    per_topic_metrics = compute_per_topic_metrics(preds, test_labels, topics)
    print_results(preds, test_labels, per_topic_metrics)

    # -- Save checkpoint -------------------------------------------------------
    acc = float((preds == test_labels).mean())
    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_prefix = f"{CHECKPOINT_DIR}/{model_size}_{ts}"

    train_df = build_dataset_df(
        train_texts, train_gens, train_labels, topics, train_steps
    )
    test_df = build_dataset_df(
        test_texts, test_gens, test_labels, topics, test_steps
    )
    train_df.to_parquet(f"{run_prefix}_train.parquet", index=False)
    test_df.to_parquet(f"{run_prefix}_test.parquet", index=False)

    # feature vectors saved as parquet for easy inspection
    feat_col_names = pd.Index([f"feat_{int(fid)}" for fid in best_feats])
    feat_df = pd.DataFrame(train_feat_vectors).set_axis(feat_col_names, axis=1)
    feat_df["label"] = train_labels
    feat_df.to_parquet(f"{run_prefix}_train_features.parquet", index=False)

    feat_test_df = pd.DataFrame(test_feat_vectors).set_axis(feat_col_names, axis=1)
    feat_test_df["label"] = test_labels
    feat_test_df.to_parquet(f"{run_prefix}_test_features.parquet", index=False)

    # move to CPU before saving so the checkpoint loads without CUDA
    torch.save(mlp.cpu().state_dict(), f"{run_prefix}_mlp.pt")

    # all config + scores + metadata in a single JSON
    meta = {
        "model_size": model_size,
        "model_name": cfg["model_name"],
        "sae_repo": cfg["sae_repo"],
        "top_k_sae": cfg["top_k_sae"],
        "layer_stride": cfg.get("layer_stride", 1),
        "topics": topics,
        "n_pos": N_POS,
        "seed": SEED,
        "train_ratio": TRAIN_RATIO,
        "batch_size": BATCH_SIZE,
        "max_gen_tokens": MAX_GEN_TOKENS,
        "max_input_tokens": MAX_INPUT_TOKENS,
        "mlp_hidden": MLP_HIDDEN,
        "mlp_epochs": MLP_EPOCHS,
        "mlp_lr": MLP_LR,
        "best_layer": best_layer,
        "feature_ids": best_feats.tolist(),
        "n_features": len(best_feats),
        "layer_scores": layer_scores,
        "accuracy": acc,
        "per_topic_metrics": per_topic_metrics,
        "loss_history": loss_history,
        "dead_samples_train": n_dead_train,
        "dead_samples_test": n_dead_test,
        "train_samples": len(train_labels),
        "test_samples": len(test_labels),
        "timestamp": ts,
    }
    with open(f"{run_prefix}_meta.json", "w") as f:
        json.dump(meta, f, indent=2)

    print(f"\n  saved: {run_prefix}_mlp.pt")
    print(f"         {run_prefix}_meta.json")
    print(f"         {run_prefix}_train.parquet / _test.parquet")
    print(f"         {run_prefix}_train_features.parquet / _test_features.parquet")

    print("\n" + "=" * 72)
    print(f"Done: Qwen3-{model_size}.")
    print("=" * 72)


def parse_model_sizes(raw):
    """Parse and validate a comma-separated list of model size strings."""
    model_sizes = [item.strip() for item in raw.split(",") if item.strip()]
    unknown = [item for item in model_sizes if item not in CONFIGS]
    if unknown:
        raise ValueError(f"unknown model size(s): {unknown}; valid: {sorted(CONFIGS)}")
    return model_sizes


def main():
    """Entry point: parse args and run pipeline for each requested model size."""
    parser = argparse.ArgumentParser(description="Qwen-Scope SAE topic classification.")
    parser.add_argument(
        "--model-sizes",
        default=",".join(DEFAULT_MODEL_SIZES),
        help=f"Comma-separated model sizes. Valid: {', '.join(CONFIGS)}",
    )
    args = parser.parse_args()

    for model_size in parse_model_sizes(args.model_sizes):
        run_for_model_size(model_size)


if __name__ == "__main__":
    main()