o1 impl

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from huggingface_hub import login, hf_hub_download
+import pyreft
 import pyvene as pv
 from threading import Thread
 from typing import Iterator
@@ -66,7 +67,7 @@ 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(
+    pv_model = pyreft.IntervenableModel(
         {
             "component": f"model.layers[20].output",
             "intervention": steer,
@@ -80,8 +81,8 @@ terminators = [tokenizer.eos_token_id] if tokenizer else []
 def generate(
     message: str,
     chat_history: list[tuple[str, str]],
-    max_new_tokens: int,
     subspaces_list: list[dict],
+    max_new_tokens: int=DEFAULT_MAX_NEW_TOKENS,
 ) -> Iterator[str]:
     # limit to last 3 turns
     start_idx = max(0, len(chat_history) - 3)
@@ -211,13 +212,7 @@ with gr.Blocks(css="style.css") as demo:
                 value=3
             )
             add_button = gr.Button("Add Concept")
-            active_subspaces_table = gr.Dataframe(
-                headers=["Concept", "Mag (scaled)"],
-                datatype=["str", "number"],
-                value=_build_table_data(default_subspaces),
-                interactive=False,
-                label="Active Concept Subspaces",
-            )
+
             # Row with the remove dropdown + button
             with gr.Row():
                 remove_dropdown = gr.Dropdown(
@@ -227,27 +222,17 @@ with gr.Blocks(css="style.css") as demo:
                 )
                 remove_button = gr.Button("Remove", variant="secondary")
 
-    # Place the max tokens slider at bottom, smaller
-    with gr.Row():
-        gr.Markdown("**Max New Tokens**", elem_classes=["small-label"])
-        max_token_slider = gr.Slider(
-            minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
-            label="",  # hide the big label
-            container=False,
-        )
-
     # Wire up events
     search_box.change(update_dropdown_choices, [search_box], [concept_dropdown])
     add_button.click(
         add_concept_to_list,
         [concept_dropdown, concept_magnitude, selected_subspaces],
-        [selected_subspaces, active_subspaces_table, remove_dropdown]
+        [selected_subspaces, remove_dropdown]
     )
     remove_button.click(
         remove_concept_from_list,
         [remove_dropdown, selected_subspaces],
-        [selected_subspaces, active_subspaces_table, remove_dropdown]
+        [selected_subspaces, remove_dropdown]
     )
 
     demo.launch()

requirements.txt

@@ -3,4 +3,4 @@ gradio==4.26.0
 spaces==0.26.0
 # TODO: pyreft main
 huggingface-hub>=0.24.0
-pyvene @ git+https://github.com/stanfordnlp/pyvene.git
+pyreft @ git+https://github.com/stanfordnlp/pyreft.git