app.py

import sys
import subprocess
import pkg_resources

def install(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", package])

required_packages = {
    'gradio': '3.50.2',
    'torch': '2.0.0',
    'transformers': '4.35.0',
    'einops': '0.6.1',
    'tqdm': '4.66.1',
    'bitsandbytes': '0.41.1'
}

for package, version in required_packages.items():
    try:
        pkg_resources.require(f"{package}>={version}")
    except pkg_resources.VersionConflict:
        print(f"{package} version {version} or higher is required. Upgrading...")
        install(f"{package}>={version}")
    except pkg_resources.DistributionNotFound:
        print(f"{package} is not installed. Installing...")
        install(f"{package}>={version}")

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from einops import einsum
from tqdm import tqdm

device = "cuda" if torch.cuda.is_available() else "cpu"
model_name = 'microsoft/Phi-3-mini-4k-instruct'

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    device_map=device, 
    torch_dtype="auto", 
    trust_remote_code=True, 
    quantization_config=quantization_config,
)

tokenizer = AutoTokenizer.from_pretrained(model_name)

def tokenize_instructions(tokenizer, instructions):
    return tokenizer.apply_chat_template(
        instructions,
        padding=True,
        truncation=False,
        return_tensors="pt",
        return_dict=True,
        add_generation_prompt=True,
    ).input_ids

def find_steering_vecs(model, base_toks, target_toks, batch_size=16):
    device = model.device
    num_its = len(range(0, base_toks.shape[0], batch_size))
    steering_vecs = {}
    for i in tqdm(range(0, base_toks.shape[0], batch_size)):
        base_out = model(base_toks[i:i+batch_size].to(device), output_hidden_states=True).hidden_states
        target_out = model(target_toks[i:i+batch_size].to(device), output_hidden_states=True).hidden_states
        for layer in range(len(base_out)):
            if i == 0:
                steering_vecs[layer] = torch.mean(target_out[layer][:,-1,:].detach().cpu() - base_out[layer][:,-1,:].detach().cpu(), dim=0)/num_its
            else:
                steering_vecs[layer] += torch.mean(target_out[layer][:,-1,:].detach().cpu() - base_out[layer][:,-1,:].detach().cpu(), dim=0)/num_its
    return steering_vecs

def do_steering(model, test_toks, steering_vec, scale=1, normalise=True, layer=None, proj=True, batch_size=16):
    def modify_activation():
        def hook(model, input):
            if normalise:
                sv = steering_vec / steering_vec.norm()
            else:
                sv = steering_vec
            sv = torch.clamp(sv, min=-1e3, max=1e3)
            if proj:
                sv = einsum(input[0], sv.view(-1,1), 'b l h, h s -> b l s') * sv
            input[0][:,:,:] = input[0][:,:,:] - scale * sv
        return hook
    
    handles = []
    if steering_vec is not None:
        for i in range(len(model.model.layers)):
            if layer is None or i == layer:
                handles.append(model.model.layers[i].register_forward_pre_hook(modify_activation()))
    
    outs_all = []
    for i in tqdm(range(0, test_toks.shape[0], batch_size)):
        outs = model.generate(test_toks[i:i+batch_size], num_beams=4, do_sample=True, max_new_tokens=60)
        outs_all.append(outs)
    outs_all = torch.cat(outs_all, dim=0)
    
    for handle in handles:
        handle.remove()
    
    return outs_all

def create_steering_vector(towards, away):
    towards_data = [[{"role": "user", "content": text.strip()}] for text in towards.split(',')]
    away_data = [[{"role": "user", "content": text.strip()}] for text in away.split(',')]
    
    towards_toks = tokenize_instructions(tokenizer, towards_data)
    away_toks = tokenize_instructions(tokenizer, away_data)
    
    steering_vecs = find_steering_vecs(model, away_toks, towards_toks)
    return steering_vecs

def chat(message, history, steering_vec, layer):
    history_formatted = [{"role": "user", "content": message}]

    input_ids = tokenize_instructions(tokenizer, [history_formatted])
    
    generations_baseline = do_steering(model, input_ids.to(device), None)
    for j in range(generations_baseline.shape[0]):
        response_baseline = f"BASELINE: {tokenizer.decode(generations_baseline[j], skip_special_tokens=True, layer=layer)}"

    if steering_vec is not None:
        generation_intervene = do_steering(model, input_ids.to(device), steering_vec[layer].to(device), scale=0.5)
        for j in range(generation_intervene.shape[0]):
            response_intervention = f"INTERVENTION: {tokenizer.decode(generation_intervene[j], skip_special_tokens=True)}"

    response = response_baseline 

    if steering_vec is not None:
        response += "\n\n" + response_intervention
    
    return [(message, response)]

def launch_app():
    with gr.Blocks() as demo:
        steering_vec = gr.State(None)
        layer = gr.State(None)

        away_default = ['hate','i hate this', 'hating the', 'hater', 'hating', 'hated in']

        towards_default = ['love','i love this', 'loving the', 'lover', 'loving', 'loved in']

        instructions = """
        ### Instructions for Using the Steering Chatbot
        
        Welcome to the Steering Chatbot! This app allows you to explore how language models can be guided (or "steered") 
        to generate different types of responses. You will be able to create **steering vectors** that influence the chatbot to either generate responses 
        that favor one set of ideas (like "love") or avoid another set (like "hate").

        #### How to Use the App:

        1. **Define Your "Towards" and "Away" Directions:**
           - In the **"Towards"** text box, enter a list of concepts, words, or phrases (comma-separated) that you want the model to generate responses toward. 
             For example, you might use: `love, happiness, kindness`.
           - In the **"Away"** text box, enter a list of concepts, words, or phrases that you want the model to steer away from. 
             For example: `hate, anger, sadness`.

        2. **Create a Steering Vector:**
           - Click the **"Create Steering Vector"** button to generate a vector that will nudge the model’s responses. 
             This vector will attempt to shift the model’s behavior towards the concepts in the "Towards" box and away from the concepts in the "Away" box.
           - You can also adjust the **layer slider** to choose which layer of the model the steering vector will affect.

        3. **Chat with the Model:**
           - Type a message in the chatbox and press Enter. The model will generate two responses:
             - **Baseline Response:** This is the model’s response without any steering vector applied.
             - **Intervention Response:** This is the response after applying the steering vector.

        4. **Compare Results:**
           - The chatbot will show both the baseline (non-steered) and the intervention (steered) responses. 
             You can compare them to see how much influence the steering vector had on the generated text.

        **Tips:**
        - Try experimenting with different word sets for "Towards" and "Away" to see how it affects the chatbot's behavior.
        - Adjusting the **layer slider** allows you to control at which stage of the model's processing the steering vector is applied, 
          which can lead to different types of modifications in the output.

        Happy chatting!
        """

        instruction_dropdown = gr.Markdown(instructions)
        
        with gr.Row():
            towards = gr.Textbox(label="Towards (comma-separated)", value= ", ".join(sentence.replace(",", "") for sentence in towards_default))
            away = gr.Textbox(label="Away from (comma-separated)", value= ", ".join(sentence.replace(",", "") for sentence in away_default))
        
        with gr.Row():
            create_vector = gr.Button("Create Steering Vector")
            layer_slider = gr.Slider(minimum=0, maximum=len(model.model.layers)-1, step=1, label="Layer", value=0)
        
        def create_vector_and_set_layer(towards, away, layer_value):
            vectors = create_steering_vector(towards, away)
            layer.value = int(layer_value)
            steering_vec.value = vectors
            return f"Steering vector created for layer {layer_value}"
        create_vector.click(create_vector_and_set_layer, [towards, away, layer_slider], gr.Textbox())
        
        chatbot = gr.Chatbot()
        msg = gr.Textbox()

        msg.submit(chat, [msg, chatbot, steering_vec, layer], chatbot)

    demo.launch()

if __name__ == "__main__":
    launch_app()