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()