o1 impl

app.py

@@ -13,7 +13,7 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 login(token=HF_TOKEN)
 
 MAX_MAX_NEW_TOKENS = 2048
-DEFAULT_MAX_NEW_TOKENS = 512  # smaller default to save memory
+DEFAULT_MAX_NEW_TOKENS = 128  # smaller default to save memory
 MAX_INPUT_TOKEN_LENGTH = 4096
 
 def load_jsonl(jsonl_path):
@@ -29,7 +29,8 @@ class Steer(pv.SourcelessIntervention):
     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.embed_dim, kwargs["latent_dim"], bias=False
+        )
     def forward(self, base, source=None, subspaces=None):
         steering_vec = torch.tensor(subspaces["mag"]) * \
             self.proj.weight[subspaces["idx"]].unsqueeze(dim=0)
@@ -82,9 +83,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(
@@ -101,7 +102,12 @@ def generate(
         "unit_locations": None,
         "max_new_tokens": max_new_tokens,
         "intervene_on_prompt": True,
-        "subspaces": [{"idx": int(subspaces_list[0]["idx"]), "mag": int(subspaces_list[0]["internal_mag"])}],
+        "subspaces": [
+            {
+                "idx": int(subspaces_list[0]["idx"]),
+                "mag": int(subspaces_list[0]["internal_mag"])
+            }
+        ] if subspaces_list else [],
         "streamer": streamer,
         "do_sample": True
     }
@@ -121,8 +127,14 @@ def filter_concepts(search_text: str):
     return filtered[:500]
 
 def add_concept_to_list(selected_concept, user_slider_val, current_list):
+    """
+    Return exactly 2 values:
+      1) The updated list of concepts (list of dicts).
+      2) A Gradio update for the removal dropdown’s choices.
+    """
     if not selected_concept:
-        return current_list, _build_table_data(current_list), gr.update(choices=_build_remove_choices(current_list))
+        return current_list, gr.update(choices=_build_remove_choices(current_list))
+
     idx = concept_id_map[selected_concept]
     internal_mag = user_slider_val * 50
     new_entry = {
@@ -132,24 +144,18 @@ def add_concept_to_list(selected_concept, user_slider_val, current_list):
         "internal_mag": internal_mag,
     }
     updated_list = current_list + [new_entry]
-    return (
-        updated_list,
-        _build_table_data(updated_list),
-        gr.update(choices=_build_remove_choices(updated_list))
-    )
+    return updated_list, gr.update(choices=_build_remove_choices(updated_list))
 
 def remove_concept_from_list(selected_text, current_list):
+    """
+    Return exactly 2 values:
+      1) The updated list of concepts (list of dicts).
+      2) A Gradio update for the removal dropdown’s choices.
+    """
     if not selected_text:
-        return current_list, _build_table_data(current_list), gr.update(choices=_build_remove_choices(current_list))
+        return current_list, gr.update(choices=_build_remove_choices(current_list))
     updated_list = [x for x in current_list if x["text"] != selected_text]
-    return (
-        updated_list,
-        _build_table_data(updated_list),
-        gr.update(choices=_build_remove_choices(updated_list))
-    )
-
-def _build_table_data(subspaces):
-    return [[x["text"], x["display_mag"]] for x in subspaces]
+    return updated_list, gr.update(choices=_build_remove_choices(updated_list))
 
 def _build_remove_choices(subspaces):
     return [x["text"] for x in subspaces]
@@ -211,12 +217,23 @@ with gr.Blocks(css="style.css") as demo:
                 remove_button = gr.Button("Remove", variant="secondary")
 
     # Wire up events
-    search_box.change(update_dropdown_choices, [search_box], [concept_dropdown])
+    # When the search box changes, update the concept dropdown choices:
+    search_box.change(
+        update_dropdown_choices,
+        [search_box],
+        [concept_dropdown]
+    )
+
+    # When "Add Concept" is clicked, add the concept + magnitude to the list,
+    # and update the "Remove" dropdown choices.
     add_button.click(
         add_concept_to_list,
         [concept_dropdown, concept_magnitude, selected_subspaces],
         [selected_subspaces, remove_dropdown]
     )
+
+    # When "Remove" is clicked, remove the selected concept from the list,
+    # and update the "Remove" dropdown choices.
     remove_button.click(
         remove_concept_from_list,
         [remove_dropdown, selected_subspaces],