Initial Gradio App

requirements.txt

@@ -0,0 +1,132 @@
+absl-py==2.1.0
+accelerate==1.0.0
+aiofiles==23.2.1
+aiohappyeyeballs==2.4.3
+aiohttp==3.10.9
+aiosignal==1.3.1
+annotated-types==0.7.0
+anyio==4.6.2.post1
+appnope==0.1.4
+asttokens==2.4.1
+astunparse==1.6.3
+attrs==24.2.0
+bitsandbytes==0.42.0
+certifi==2024.8.30
+charset-normalizer==3.3.2
+click==8.1.7
+comm==0.2.2
+datasets==3.0.1
+debugpy==1.8.6
+decorator==5.1.1
+diffusers==0.30.3
+dill==0.3.8
+einops==0.8.0
+executing==2.1.0
+fastapi==0.115.2
+ffmpy==0.4.0
+filelock==3.16.1
+flatbuffers==24.3.25
+frozenlist==1.4.1
+fsspec==2024.6.1
+gast==0.6.0
+google-pasta==0.2.0
+gradio==5.1.0
+gradio_client==1.4.0
+grpcio==1.67.0
+h11==0.14.0
+h5py==3.12.1
+httpcore==1.0.6
+httpx==0.27.2
+huggingface-hub==0.25.1
+idna==3.10
+importlib_metadata==8.5.0
+ipykernel==6.29.5
+ipython==8.28.0
+jedi==0.19.1
+Jinja2==3.1.4
+jupyter_client==8.6.3
+jupyter_core==5.7.2
+keras==3.6.0
+libclang==18.1.1
+loralib==0.1.2
+Markdown==3.7
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib-inline==0.1.7
+mdurl==0.1.2
+ml-dtypes==0.4.1
+mlx==0.18.1
+mlx-lm==0.19.0
+mpmath==1.3.0
+multidict==6.1.0
+multiprocess==0.70.16
+namex==0.0.8
+nest-asyncio==1.6.0
+networkx==3.3
+numpy==1.26.4
+opt_einsum==3.4.0
+optree==0.13.0
+orjson==3.10.7
+packaging==24.1
+pandas==2.2.3
+parso==0.8.4
+peft @ git+https://github.com/huggingface/peft.git@a724834ac43b9478b066d3ec8b421489151f3815
+pexpect==4.9.0
+pillow==10.4.0
+platformdirs==4.3.6
+prompt_toolkit==3.0.48
+propcache==0.2.0
+protobuf==4.25.5
+psutil==6.0.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pyarrow==17.0.0
+pydantic==2.9.2
+pydantic_core==2.23.4
+pydub==0.25.1
+Pygments==2.18.0
+python-dateutil==2.9.0.post0
+python-multipart==0.0.12
+pytz==2024.2
+PyYAML==6.0.2
+pyzmq==26.2.0
+regex==2024.9.11
+requests==2.32.3
+rich==13.9.2
+ruff==0.6.9
+safetensors==0.4.5
+scipy==1.14.1
+semantic-version==2.10.0
+sentencepiece==0.2.0
+setuptools==75.1.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+stack-data==0.6.3
+starlette==0.40.0
+sympy==1.13.3
+tensorboard==2.17.1
+tensorboard-data-server==0.7.2
+tensorflow==2.17.0
+termcolor==2.5.0
+tokenizers==0.20.0
+tomlkit==0.12.0
+torch==2.4.1
+tornado==6.4.1
+tqdm==4.66.5
+traitlets==5.14.3
+transformers @ git+https://github.com/huggingface/transformers.git@698b36da72ae8377fb08ade92b131069898007c2
+typer==0.12.5
+typing_extensions==4.12.2
+tzdata==2024.2
+urllib3==2.2.3
+uvicorn==0.32.0
+wcwidth==0.2.13
+websockets==12.0
+Werkzeug==3.0.4
+wheel==0.44.0
+wrapt==1.16.0
+xxhash==3.5.0
+yarl==1.14.0
+zipp==3.20.2
+

steering_gradio.py

@@ -0,0 +1,180 @@
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+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'
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    device_map=device, 
+    torch_dtype="auto", 
+    trust_remote_code=True, 
+)
+
+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
+            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" if i % 2 == 0 else "assistant", "content": msg} for i, msg in enumerate(history)]
+    history_formatted.append({"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=1)
+        for j in range(generation_intervene.shape[0]):
+            response_intervention = f"INTERVENTION: {tokenizer.decode(generation_intervene[j], skip_special_tokens=True)}"
+
+    response = response_baseline + "\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 = [
+            "Apples are a popular fruit enjoyed by people around the world.",
+            "The apple tree originated in Central Asia and has been cultivated for thousands of years.",
+            "There are over 7,500 known cultivars of apples.",
+            "Apples are members of the rose family, Rosaceae.",
+            # "The science of apple cultivation is called pomology.",
+            # "Apple trees typically take 4-5 years to produce their first fruit.",
+            # "The phrase 'An apple a day keeps the doctor away' originated in Wales in the 19th century.",
+            # "Apples are rich in antioxidants, flavonoids, and dietary fiber.",
+            # "The most popular apple variety in the United States is the Gala apple.",
+            # "Apple seeds contain a compound called amygdalin, which can release cyanide when digested.",
+            # "The apple is the official state fruit of New York.",
+            # "Apples can be eaten raw, cooked, or pressed for juice.",
+            # "The largest apple ever picked weighed 4 pounds 1 ounce.",
+            # "Apples float in water because 25 percent of their volume is air.",
+            # "The apple blossom is the state flower of Michigan.",
+            # "China is the world's largest producer of apples.",
+            # "The average apple tree can produce up to 840 pounds of apples per year.",
+            # "Apples ripen six to ten times faster at room temperature than if they are refrigerated.",
+            # "The first apple trees in North America were planted by pilgrims in Massachusetts Bay Colony.",
+            # "Apples are harvested by hand in orchards."
+        ]
+
+        towards_default = [
+            "The United States of America is the world's third-largest country by total area.",
+            "America declared its independence from Great Britain on July 4, 1776.",
+            "The U.S. Constitution, written in 1787, is the oldest written national constitution still in use.",
+            "The United States has 50 states and one federal district, Washington D.C.",
+            # "America's national motto is 'In God We Trust,' adopted in 1956.",
+            # "The bald eagle is the national bird and symbol of the United States.",
+            # "The Statue of Liberty, a gift from France, stands in New York Harbor as a symbol of freedom.",
+            # "American culture has had a significant influence on global entertainment and technology.",
+            # "The United States is home to many diverse ecosystems, from deserts to tropical rainforests.",
+            # "America is often referred to as a 'melting pot' due to its diverse immigrant population.",
+            # "The U.S. has the world's largest economy by nominal GDP.",
+            # "American football, derived from rugby, is the most popular sport in the United States.",
+            # "The Grand Canyon, located in Arizona, is one of America's most famous natural landmarks.",
+            # "The U.S. sent the first humans to walk on the moon in 1969.",
+            # "America's system of government is a federal republic with a presidential system.",
+            # "The American flag, known as the Stars and Stripes, has 13 stripes and 50 stars.",
+            # "English is the de facto national language, but the U.S. has no official language at the federal level.",
+            # "The United States is home to many world-renowned universities, including Harvard and MIT.",
+            # "America's national anthem, 'The Star-Spangled Banner,' was written in 1814.",
+            # "The U.S. Interstate Highway System, started in 1956, is one of the largest public works projects in history."
+        ]
+        
+        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()
+
+
+    # clean up
+    # nicer baseline vs intervention
+    # auto clear after messgae