o1 impl

app.py

@@ -34,8 +34,11 @@ class Steer(pv.SourcelessIntervention):
     def forward(self, base, source=None, subspaces=None):
         if subspaces is None:
             return base
-        steering_vec = torch.tensor(subspaces["mag"]) * \
-            self.proj.weight[subspaces["idx"]].unsqueeze(dim=0)
+        steering_vec = []
+        for idx, mag in zip(subspaces["idx"], subspaces["mag"]):
+            steering_vec.append(self.proj.weight[idx].unsqueeze(dim=0))
+        steering_vec = torch.cat(steering_vec, dim=0).mean(dim=0)
+        steering_vec = mag * steering_vec
         return base + steering_vec
 
 # Check GPU
@@ -91,9 +94,9 @@ def generate(
 
     # build list of messages
     messages = []
-    # for user_msg, model_msg in recent_history:
-    #     messages.append({"role": "user", "content": user_msg})
-    #     messages.append({"role": "model", "content": model_msg})
+    for user_msg, model_msg in recent_history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "model", "content": model_msg})
     messages.append({"role": "user", "content": message})
 
     input_ids = torch.tensor([tokenizer.apply_chat_template(
@@ -113,8 +116,8 @@ def generate(
         "intervene_on_prompt": True,
         "subspaces": [
             {
-                "idx": int(subspaces_list[0]["idx"]),
-                "mag": int(subspaces_list[0]["internal_mag"])
+                "idx": [int(sl["idx"]) for sl in subspaces_list],
+                "mag": [int(sl["internal_mag"]) for sl in subspaces_list]
             }
         ] if subspaces_list else None,
         "streamer": streamer,