app.py

import os
import random
import torch
from tqdm import tqdm
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

# --- Helper functions ---

def load_instructions(dataset_id, column, n_instructions):
    dataset = load_dataset(dataset_id, split="train")
    indices = random.sample(range(len(dataset)), n_instructions * 2)
    return [dataset[i][column] for i in indices[:n_instructions]], [
        dataset[i][column] for i in indices[n_instructions:]
    ]

def generate_response(model, tokenizer, prompt, max_new_tokens=128):
    if hasattr(tokenizer, "apply_chat_template"):
        inputs = tokenizer.apply_chat_template(
            conversation=[{"role": "user", "content": prompt}],
            add_generation_prompt=True,
            return_tensors="pt",
        ).to(model.device)
    else:
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    output_ids = model.generate(
        inputs,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=0.5,
        min_p=0.1,
        repetition_penalty=1.05,
    )
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

def generate_outputs(model, tokenizer, instructions, system_prompt):
    outputs = []
    for instruction in tqdm(instructions, desc="Generating outputs", leave=False):
        if hasattr(tokenizer, "apply_chat_template"):
            inputs = tokenizer.apply_chat_template(
                conversation=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": instruction},
                ],
                add_generation_prompt=True,
                return_tensors="pt",
            ).to(model.device)
        else:
            prompt = system_prompt + "\n" + instruction
            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        out = model.generate(
            inputs,
            use_cache=False,
            max_new_tokens=1,
            return_dict_in_generate=True,
            output_hidden_states=True,
        )
        outputs.append(out["hidden_states"][0])
    return outputs

def orthogonalize_matrix(matrix, vec, weight):
    vec = vec.view(-1).to(matrix.device)
    if matrix.shape[-1] == vec.shape[0]:
        proj = torch.einsum("...d,d->...", matrix, vec).unsqueeze(-1) * vec.unsqueeze(0)
        return matrix - weight * proj
    elif matrix.shape[0] == vec.shape[0]:
        proj = torch.einsum("d...,d->...", matrix, vec).unsqueeze(0) * vec.unsqueeze(-1)
        return matrix - weight * proj
    else:
        raise ValueError(
            f"Matrix shape {matrix.shape} incompatible with vector shape {vec.shape}"
        )

# --- Streamlit UI ---

st.title("LLM Abliteration with Qwen")
st.markdown("Credits: Thanks to **Maxime Labonne**")
st.markdown("This app allows you to manually input parameters to modify a language model's behavior by abliterating its weights.")

st.sidebar.header("Abliteration Parameters")
MODEL_ID = st.sidebar.text_input("Model ID", "Qwen/Qwen2.5-3B-Instruct")
N_INSTRUCTIONS = st.sidebar.number_input("Number of Instructions", min_value=1, value=128, step=1)
TARGET_LAYER = st.sidebar.slider("Target Layer (relative ratio)", 0.0, 1.0, 0.65, step=0.05)
REFUSAL_WEIGHT = st.sidebar.slider("Refusal Weight", 0.0, 2.0, 1.0, step=0.05)
PRIVATE_UPLOAD = st.sidebar.checkbox("Push Model to Hub Privately", value=True)

st.sidebar.header("HF Token")
hf_token = st.sidebar.text_input("Hugging Face Token", type="password")
if hf_token:
    os.environ["HF_TOKEN"] = hf_token

st.sidebar.header("Target Dataset")
target_prompt = st.sidebar.text_area("Target Prompt", "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.")
target_dataset = st.sidebar.text_input("Target Dataset ID", "mlabonne/harmful_behaviors")
target_column = st.sidebar.text_input("Target Column Name", "text")

st.sidebar.header("Baseline Dataset")
baseline_prompt = st.sidebar.text_area("Baseline Prompt", "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.")
baseline_dataset = st.sidebar.text_input("Baseline Dataset ID", "mlabonne/harmless_alpaca")
baseline_column = st.sidebar.text_input("Baseline Column Name", "text")

if st.button("Run Abliteration"):
    st.write("### Loading Model and Tokenizer")
    if torch.cuda.is_available():
        if torch.cuda.get_device_capability()[0] >= 8:
            torch_dtype = torch.bfloat16
            attn_implementation = "flash_attention_2"
        else:
            torch_dtype = torch.float16
            attn_implementation = "eager"
    else:
        torch_dtype = torch.float32
        attn_implementation = "eager"

    with st.spinner("Loading model..."):
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            device_map="auto",
            torch_dtype=torch_dtype,
            attn_implementation=attn_implementation,
            trust_remote_code=True
        ).eval()
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    
    try:
        total_layers = len(model.model.layers)
    except AttributeError:
        st.error("The loaded model does not have a `model.layers` attribute.")
        st.stop()
    LAYER_IDX = int(TARGET_LAYER * total_layers)
    
    st.write("### Before Abliteration Response")
    prompt_example = "How to write a computer virus?"
    response_before = generate_response(model, tokenizer, prompt_example)
    st.text_area("Response before abliteration:", response_before, height=150)

    st.write("### Loading Instructions")
    with st.spinner("Loading target instructions..."):
        target_instructions, _ = load_instructions(target_dataset, target_column, N_INSTRUCTIONS)
    with st.spinner("Loading baseline instructions..."):
        baseline_instructions, _ = load_instructions(baseline_dataset, baseline_column, N_INSTRUCTIONS)

    st.write("### Generating Hidden States")
    with st.spinner("Generating baseline hidden states..."):
        baseline_outputs = generate_outputs(model, tokenizer, baseline_instructions, system_prompt=baseline_prompt)
    with st.spinner("Generating target hidden states..."):
        target_outputs = generate_outputs(model, tokenizer, target_instructions, system_prompt=target_prompt)

    target_hidden = [output[LAYER_IDX][:, -1, :] for output in target_outputs]
    baseline_hidden = [output[LAYER_IDX][:, -1, :] for output in baseline_outputs]

    st.write("### Calculating Refusal Direction")
    target_mean = torch.stack(target_hidden).mean(dim=0)
    baseline_mean = torch.stack(baseline_hidden).mean(dim=0)
    refusal_dir = target_mean - baseline_mean
    refusal_dir = refusal_dir / refusal_dir.norm()

    del target_outputs, baseline_outputs, target_hidden, baseline_hidden

    st.write("### Orthogonalizing Model Weights")
    refusal_dir = refusal_dir.view(-1).to(model.device)
    stats = {"embed_tokens": False, "attention_o_proj": 0, "mlp_proj": 0}

    if hasattr(model.model, "embed_tokens"):
        model.model.embed_tokens.weight.data = orthogonalize_matrix(
            model.model.embed_tokens.weight.data, refusal_dir, REFUSAL_WEIGHT
        )
        stats["embed_tokens"] = True

    for layer in tqdm(model.model.layers, desc="Orthogonalizing weights", leave=False):
        if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "o_proj"):
            layer.self_attn.o_proj.weight.data = orthogonalize_matrix(
                layer.self_attn.o_proj.weight.data, refusal_dir, REFUSAL_WEIGHT
            )
            stats["attention_o_proj"] += 1

        if hasattr(layer, "mlp"):
            proj_name = (
                "down_proj"
                if hasattr(layer.mlp, "down_proj")
                else "c_proj"
                if hasattr(layer.mlp, "c_proj")
                else None
            )
            if proj_name:
                getattr(layer.mlp, proj_name).weight.data = orthogonalize_matrix(
                    getattr(layer.mlp, proj_name).weight.data, refusal_dir, REFUSAL_WEIGHT
                )
                stats["mlp_proj"] += 1

    del refusal_dir

    if (
        not stats["embed_tokens"]
        and stats["attention_o_proj"] == 0
        and stats["mlp_proj"] == 0
    ):
        st.error("Failed to orthogonalize any model weights. Model not abliterated.")
        st.stop()

    st.write(f"Orthogonalization stats: {stats}")

    st.write("### After Abliteration Response")
    response_after = generate_response(model, tokenizer, prompt_example)
    st.text_area("Response after abliteration:", response_after, height=150)

    st.write("### (Optional) Pushing Model to Hugging Face Hub")
    if st.checkbox("Push model to HF Hub?"):
        try:
            model_name = MODEL_ID.split("/")[-1] + "-abliterated"
            model.push_to_hub(model_name, private=PRIVATE_UPLOAD)
            tokenizer.push_to_hub(model_name, private=PRIVATE_UPLOAD)
            st.success(f"Model pushed as {model_name}")
        except Exception as e:
            st.error(f"Error while pushing model: {e}")

    st.success("Abliteration process complete!")