Upload 5 files

app.py

@@ -0,0 +1,282 @@
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoConfig
+from pathlib import Path
+import spaces
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+import json
+from model import SAE, SteerableOlmo2ForCausalLM
+
+# Initialize model and tokenizer
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_name = "allenai/OLMo-2-1124-7B-Instruct"
+
+print("Loading model and tokenizer...")
+model = SteerableOlmo2ForCausalLM.from_pretrained(
+    model_name, 
+    torch_dtype=torch.bfloat16
+).to(device)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model_config = AutoConfig.from_pretrained(model_name)
+
+# Load SAE from Hugging Face Hub
+print("Loading SAE from Hugging Face Hub...")
+
+# Download SAE files from your model repository
+sae_weights_path = hf_hub_download(
+    repo_id="open-concept-steering/olmo2-7b-sae-65k-v1",
+    filename="sae_weights.safetensors"
+)
+sae_config_path = hf_hub_download(
+    repo_id="open-concept-steering/olmo2-7b-sae-65k-v1", 
+    filename="sae_config.json"
+)
+
+# Load SAE
+sae_weights = load_file(sae_weights_path, device=device)
+with open(sae_config_path, "r") as f:
+    sae_config = json.load(f)
+
+sae = SAE(sae_config['input_size'], sae_config['hidden_size']).to(device).to(torch.bfloat16)
+sae.load_state_dict(sae_weights)
+
+# Set up steering
+steering_layer = model_config.num_hidden_layers // 2 - 1
+model.set_sae_and_layer(sae, steering_layer)
+
+# Steering features configuration
+STEERING_FEATURES = {
+    "None": {"feature": None, "default": 0, "name": "No Steering"},
+    "batman/bruce wayne": {"feature": 758, "default": 11, "name": "🦸 Superhero/Batman"},
+    "japan": {"feature": 29940, "default": 13, "name": "🗾 Japan"},
+    "baseball": {"feature": 65023, "default": 6, "name": "⚾ Baseball"}
+}
+
+default_system_prompt = "You are OLMo 2, a helpful and harmless AI Assistant built by the Allen Institute for AI."
+
+@spaces.GPU
+def generate_responses(message, history_unsteered, history_steered, steering_type, steering_strength, system_prompt):
+    """Generate both unsteered and steered responses with conversation history"""
+    
+    if not message:
+        return history_unsteered, history_steered, ""
+    
+    # Build messages for unsteered conversation
+    messages_unsteered = []
+    if system_prompt:
+        messages_unsteered.append({"role": "system", "content": system_prompt})
+    
+    # Add conversation history
+    for msg in history_unsteered:
+        messages_unsteered.append({"role": msg["role"], "content": msg["content"]})
+    
+    # Add current message
+    messages_unsteered.append({"role": "user", "content": message})
+    
+    # Format prompt for unsteered
+    formatted_prompt_unsteered = tokenizer.apply_chat_template(
+        messages_unsteered,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    
+    inputs_unsteered = tokenizer(
+        formatted_prompt_unsteered,
+        return_tensors="pt",
+        padding=True,
+        return_attention_mask=True
+    ).to(device)
+    
+    # Generate unsteered response
+    model.clear_steering()
+    with torch.inference_mode():
+        outputs_unsteered = model.generate(
+            input_ids=inputs_unsteered.input_ids,
+            attention_mask=inputs_unsteered.attention_mask,
+            max_new_tokens=256,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    full_response_unsteered = tokenizer.decode(outputs_unsteered[0], skip_special_tokens=False)
+    unsteered_response = full_response_unsteered.split("<|assistant|>")[-1].split("<|endoftext|>")[0].strip()
+    
+    # Update unsteered history
+    history_unsteered.append({"role": "user", "content": message})
+    history_unsteered.append({"role": "assistant", "content": unsteered_response})
+    
+    # Generate steered response
+    if steering_type != "None":
+        # Build messages for steered conversation  
+        messages_steered = []
+        if system_prompt:
+            messages_steered.append({"role": "system", "content": system_prompt})
+        
+        # Add conversation history
+        for msg in history_steered:
+            messages_steered.append({"role": msg["role"], "content": msg["content"]})
+        
+        # Add current message
+        messages_steered.append({"role": "user", "content": message})
+        
+        # Format prompt for steered
+        formatted_prompt_steered = tokenizer.apply_chat_template(
+            messages_steered,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        
+        inputs_steered = tokenizer(
+            formatted_prompt_steered,
+            return_tensors="pt",
+            padding=True,
+            return_attention_mask=True
+        ).to(device)
+        
+        # Apply steering
+        feature_config = STEERING_FEATURES[steering_type]
+        steering_value = feature_config["default"] * steering_strength
+        model.set_steering(feature_config["feature"], steering_value)
+        
+        with torch.inference_mode():
+            outputs_steered = model.generate(
+                input_ids=inputs_steered.input_ids,
+                attention_mask=inputs_steered.attention_mask,
+                max_new_tokens=256,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        full_response_steered = tokenizer.decode(outputs_steered[0], skip_special_tokens=False)
+        steered_response = full_response_steered.split("<|assistant|>")[-1].split("<|endoftext|>")[0].strip()
+        model.clear_steering()
+    else:
+        steered_response = unsteered_response
+    
+    # Update steered history
+    history_steered.append({"role": "user", "content": message})
+    history_steered.append({"role": "assistant", "content": steered_response})
+    
+    return history_unsteered, history_steered, ""
+
+def clear_chats():
+    """Clear both chat histories"""
+    return [], []
+
+# Create Gradio interface
+with gr.Blocks(title="OLMo-2 Feature Steering Demo", theme=gr.themes.Default()) as demo:
+    gr.Markdown("""
+    # 🎛️ OLMo-2 Feature Steering Demo
+    
+    This demo showcases how sparse autoencoders (SAEs) can steer OLMo-2's responses by manipulating specific features.
+    Have a conversation and see how steering changes the model's behavior across multiple turns!
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            steering_type = gr.Dropdown(
+                choices=list(STEERING_FEATURES.keys()),
+                value="None",
+                label="Steering Type",
+                info="Choose a feature to steer the model's response"
+            )
+            
+            steering_strength = gr.Slider(
+                minimum=0.5,
+                maximum=2.0,
+                value=1.0,
+                step=0.1,
+                label="Steering Strength",
+                info="Adjust the intensity of the steering effect (higher = more steering, very high values may cause gobbledygook)"
+            )
+            
+            system_prompt = gr.Textbox(
+                label="System Prompt",
+                value=default_system_prompt,
+                lines=3
+            )
+            
+            clear_btn = gr.Button("🗑️ Clear Chats", variant="secondary")
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### 🤖 Original OLMo")
+            chatbot_unsteered = gr.Chatbot(
+                label="Unsteered",
+                height=500,
+                show_copy_button=True,
+                type="messages"
+            )
+        
+        with gr.Column():
+            gr.Markdown("### 🎯 Steered OLMo")
+            chatbot_steered = gr.Chatbot(
+                label="Steered",
+                height=500,
+                show_copy_button=True,
+                type="messages"
+            )
+    
+    with gr.Row():
+        user_input = gr.Textbox(
+            label="Your Message",
+            placeholder="Type your message here... (Enter to send, Shift+Enter for new line)",
+            lines=2,
+            scale=4
+        )
+        submit_btn = gr.Button("Send", variant="primary", scale=1)
+    
+    # Example questions
+    gr.Examples(
+        examples=[
+            "What's an interesting way to spend a weekend?",
+            "Tell me about your favorite subject.",
+            "What should I do with $5?",
+            "How do you approach solving difficult problems?",
+            "What's something that makes you excited?",
+            "Tell me a story about adventure.",
+            "What advice would you give to someone feeling stuck?"
+        ],
+        inputs=user_input,
+        label="Example Questions"
+    )
+    
+    # Handle submission
+    def submit_message(message, history_unsteered, history_steered, steering_type, steering_strength, system_prompt):
+        return generate_responses(message, history_unsteered, history_steered, steering_type, steering_strength, system_prompt)
+    
+    # Wire up the interface
+    user_input.submit(
+        fn=submit_message,
+        inputs=[user_input, chatbot_unsteered, chatbot_steered, steering_type, steering_strength, system_prompt],
+        outputs=[chatbot_unsteered, chatbot_steered, user_input]
+    )
+    
+    submit_btn.click(
+        fn=submit_message,
+        inputs=[user_input, chatbot_unsteered, chatbot_steered, steering_type, steering_strength, system_prompt],
+        outputs=[chatbot_unsteered, chatbot_steered, user_input]
+    )
+    
+    clear_btn.click(
+        fn=clear_chats,
+        outputs=[chatbot_unsteered, chatbot_steered]
+    )
+    
+    # Update slider visibility based on steering selection
+    def update_slider_visibility(steering_type):
+        return gr.update(visible=(steering_type != "None"))
+    
+    steering_type.change(
+        fn=update_slider_visibility,
+        inputs=steering_type,
+        outputs=steering_strength
+    )
+
+if __name__ == "__main__":
+    demo.launch()

gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

model.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import Olmo2ForCausalLM
+class SAE(nn.Module):
+    def __init__(self, input_size, hidden_size, init_scale=0.1):
+        super().__init__()
+        
+        # Store dimensions
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        
+        # Initialize as before 
+        self.encode = nn.Linear(input_size, hidden_size, bias=True)
+        self.decode = nn.Linear(hidden_size, input_size, bias=True)
+                
+        with torch.no_grad():
+            # Random directions
+            decoder_weights = torch.randn(input_size, hidden_size) 
+            # Normalize columns
+            decoder_weights = decoder_weights / torch.linalg.vector_norm(decoder_weights, dim=0, keepdim=True)
+            # Scale by random values between 0.05 and 1.0
+            scales = torch.rand(hidden_size) * 0.95 + 0.05
+            decoder_weights = decoder_weights * scales
+            
+            self.decode.weight.data = decoder_weights
+            self.encode.weight.data = decoder_weights.T.contiguous()
+            self.encode.bias.data.zero_() #zero in place
+            self.decode.bias.data.zero_()
+
+        self.constrain_weights()
+
+    @property
+    def device(self):
+        """Return the device the model parameters are on"""
+        return next(self.parameters()).device
+
+    def constrain_weights(self):
+        """Constrain the decoder weights to have unit norm."""
+        with torch.no_grad():
+            decoder_norm = torch.linalg.vector_norm(self.decode.weight, dim=0, keepdim=True)
+            self.decode.weight.data = self.decode.weight.data / decoder_norm
+
+    def forward(self, x):
+        features = F.relu(self.encode(x))
+        reconstruction = self.decode(features)
+        return reconstruction, features
+
+    def get_decoder_norms(self):
+        # returns a 1-D tensor (hidden_size,) on the right device/dtype
+        return torch.linalg.vector_norm(self.decode.weight, dim=0)
+
+        
+    @property
+    def W_dec(self):
+        """Return decoder weights for easier access during analysis"""
+        return self.decode.weight
+        
+    def compute_loss(self, x, recon, feats, lambda_):
+        # reconstruction term — sum over feature-dim, mean over batch
+        recon_mse = (recon - x).pow(2).sum(-1).mean()
+
+        # sparsity term — L1 on feature activations * current decoder-column norms
+        sparsity = (feats.abs() * self.get_decoder_norms()).sum(1).mean()
+
+        return recon_mse + lambda_ * sparsity
+
+class SteerableOlmo2ForCausalLM(Olmo2ForCausalLM):
+    def __init__(self, config):
+        super().__init__(config)
+        self.steering_layer = None
+        self.sae = None
+        self.steering_features = {}
+        self.steering_hook = None
+        self.sae_max = None
+
+    def set_sae_and_layer(self, sae, layer):
+        self.sae = sae
+        self.steering_layer = layer
+        self._register_steering_hook()
+
+    def set_sae_max(self, sae_max):
+        self.sae_max = sae_max
+
+    def set_steering(self, feature_idx, value, *, as_multiple_of_max=False):
+        if as_multiple_of_max and self.sae_max is not None:
+            value = float(value) * float(self.sae_max[feature_idx])
+        self.steering_features[feature_idx] = value
+
+    def clear_steering(self):
+        self.steering_features = {}
+
+    @torch.no_grad()
+    def _steering_hook_fn(self, module, input, output):
+        if not self.steering_features or self.sae is None:
+            return output
+
+        hidden_states = output[0]
+        feats = self.sae.encode(hidden_states)
+        recon = self.sae.decode(feats)
+        error = hidden_states - recon
+
+        feats_steered = feats.clone()
+        for idx, clamp_value in self.steering_features.items():
+            feats_steered[..., idx] = clamp_value
+
+        recon_steered = self.sae.decode(feats_steered)
+        hidden_steered = recon_steered + error
+
+        return (hidden_steered,) + output[1:]
+
+    def _register_steering_hook(self):
+        if self.steering_hook is not None:
+            self.steering_hook.remove()
+            self.steering_hook = None
+
+        if self.steering_layer is not None:
+            target_layer = self.model.layers[self.steering_layer]
+            self.steering_hook = target_layer.register_forward_hook(self._steering_hook_fn)
+
+    def remove_steering_hook(self):
+        if self.steering_hook is not None:
+            self.steering_hook.remove()
+            self.steering_hook = None

requirements.txt

@@ -0,0 +1,6 @@
+torch
+transformers
+gradio
+spaces
+safetensors
+huggingface_hub