common.py

"""Shared helpers for refusal-direction extraction and ablation.

Both compute_refusal_dir.py and inference.py import from here so the model
choice and loading config live in exactly one place.
"""
import os

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# Pick the model once; both scripts use it.
MODEL_ID = "tiiuae/Falcon3-1B-Instruct"
# MODEL_ID = "Qwen/Qwen3-1.7B"
# MODEL_ID = "stabilityai/stablelm-2-zephyr-1_6b"
# MODEL_ID = "Qwen/Qwen1.5-1.8B-Chat"
# MODEL_ID = "Qwen/Qwen-1_8B-chat"
# MODEL_ID = "google/gemma-1.1-2b-it"
# MODEL_ID = "google/gemma-1.1-7b-it"
# MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"

# Let every backend use all CPU cores and the fast matmul kernels.
torch.set_num_threads(os.cpu_count() or 1)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True


def load_model_and_tokenizer(model_id: str = MODEL_ID):
    """Load a 4-bit quantized causal LM and its tokenizer, ready for inference."""
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        trust_remote_code=True,
        dtype=torch.float16,
        device_map="cuda",
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
        ),
    )
    model.eval()

    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    # Left padding lets us read the last-token hidden state at index -1 for a
    # whole batch in one shot.
    tokenizer.padding_side = "left"
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    return model, tokenizer


def refusal_dir_path(model_id: str = MODEL_ID) -> str:
    """Filename used to cache the refusal direction for a given model."""
    return model_id.replace("/", "_") + "_refusal_dir.pt"


def project_out(x: torch.Tensor, direction: torch.Tensor) -> torch.Tensor:
    """Remove the component of `x` that lies along the unit vector `direction`.

    proj = (x . d) d ;  result = x - proj.  Assumes ``direction`` is normalized.
    Equivalent to the original einsum, but a plain matmul does the same job.
    """
    return x - (x @ direction).unsqueeze(-1) * direction