o1 impl

app.py

@@ -32,6 +32,8 @@ class Steer(pv.SourcelessIntervention):
             self.embed_dim, kwargs["latent_dim"], bias=False
         )
     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)
         return base + steering_vec
@@ -97,6 +99,7 @@ def generate(
         yield "[Truncated prior text]\n"
 
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+    print(subspaces_list)
     generate_kwargs = {
         "base": {"input_ids": input_ids},
         "unit_locations": None,
@@ -107,7 +110,7 @@ def generate(
                 "idx": int(subspaces_list[0]["idx"]),
                 "mag": int(subspaces_list[0]["internal_mag"])
             }
-        ] if subspaces_list else [],
+        ] if subspaces_list else None,
         "streamer": streamer,
         "do_sample": True
     }