o1 impl

app.py

@@ -25,19 +25,15 @@ def load_jsonl(jsonl_path):
     return jsonl_data
 
 class Steer(pv.SourcelessIntervention):
+    """Steer model via activation addition"""
     def __init__(self, **kwargs):
         super().__init__(**kwargs, keep_last_dim=True)
-        self.proj = torch.nn.Linear(self.embed_dim, kwargs["latent_dim"], bias=False)
-
+        self.proj = torch.nn.Linear(
+                self.embed_dim, kwargs["latent_dim"], bias=False)
     def forward(self, base, source=None, subspaces=None):
-        steer_vec = base
-        if subspaces is not None:
-            for sp in subspaces:
-                idx = sp["idx"]
-                mag = sp["internal_mag"]  # scaled by 50
-                steering_vec = mag * self.proj.weight[idx].unsqueeze(dim=0)
-                steer_vec = steer_vec + steering_vec
-        return steer_vec
+        steering_vec = torch.tensor(subspaces["mag"]) * \
+            self.proj.weight[subspaces["idx"]].unsqueeze(dim=0)
+        return base + steering_vec
 
 # Check GPU
 if not torch.cuda.is_available():
@@ -67,13 +63,9 @@ if torch.cuda.is_available():
     steer = Steer(embed_dim=params.shape[0], latent_dim=params.shape[1])
     steer.proj.weight.data = params.float()
 
-    pv_model = pv.IntervenableModel(
-        {
-            "component": f"model.layers[20].output",
-            "intervention": steer,
-        },
-        model=model,
-    )
+    pv_model = pv.IntervenableModel({
+        "component": f"model.layers[20].output",
+        "intervention": steer}, model=model)
 
 terminators = [tokenizer.eos_token_id] if tokenizer else []
 
@@ -171,7 +163,7 @@ with gr.Blocks(css="style.css") as demo:
     # Pre-populate with a random concept if available
     default_subspaces = []
     if pv_model and concept_list:
-        default_concept = random.choice(concept_list)
+        default_concept = "words related to time travel and its consequences"
         default_subspaces = [{
             "text": default_concept,
             "idx": concept_id_map[default_concept],
@@ -191,7 +183,7 @@ with gr.Blocks(css="style.css") as demo:
             )
         # Right side: concept management
         with gr.Column(scale=3):
-            gr.Markdown("## Steering Concepts")
+            gr.Markdown("# Steering Concepts")
             search_box = gr.Textbox(
                 label="Search concepts",
                 placeholder="e.g. 'time travel'"