o1 impl

app.py

@@ -1,4 +1,4 @@
-import os, json
+import os, json, random
 import torch
 import gradio as gr
 import spaces
@@ -44,13 +44,14 @@ class Steer(pv.SourcelessIntervention):
         self.proj = torch.nn.Linear(self.embed_dim, kwargs["latent_dim"], bias=False)
 
     def forward(self, base, source=None, subspaces=None):
-        # subspaces is a list of dicts: each has {"idx": int, "mag": float}
+        # subspaces is a list of dicts: 
+        #   each has {"idx": int, "internal_mag": float, ...}
         steer_vec = base
         if subspaces is not None:
             for sp in subspaces:
                 idx = sp["idx"]
-                mag = sp["mag"]
-                # each idx is a row in self.proj.weight
+                # Use the internal magnitude for actual steering
+                mag = sp["internal_mag"] 
                 steering_vec = mag * self.proj.weight[idx].unsqueeze(dim=0)
                 steer_vec = steer_vec + steering_vec
         return steer_vec
@@ -107,39 +108,30 @@ def generate(
     recent_history = chat_history[start_idx:]
 
     # Build a list of messages
-    # each tuple is (user_message, assistant_message)
+    # each tuple is (user_message, model_message)
     messages = []
-    for user_msg, assistant_msg in recent_history:
+    for user_msg, model_msg in recent_history:
         messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": assistant_msg})
+        messages.append({"role": "model", "content": model_msg})
 
     # Now append the new user message
     messages.append({"role": "user", "content": message})
 
-    # Convert messages into model input tokens with a generation prompt
-    prompt = tokenizer.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True  # appends a final "Assistant:" for the model to continue
-    )
-
-    # Retrieve input_ids and mask
-    input_ids = torch.tensor([prompt["input_ids"]]).cuda()
-    attention_mask = torch.tensor([prompt["attention_mask"]]).cuda()
+    input_ids = torch.tensor([tokenizer.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True)]).cuda()
 
     # Possibly trim if over max length
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        attention_mask = attention_mask[:, -MAX_INPUT_TOKEN_LENGTH:]
         yield "\n[Warning: Truncated conversation exceeds max allowed input tokens]\n"
 
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = {
-        "base": {"input_ids": input_ids, "attention_mask": attention_mask},
+        "base": {"input_ids": input_ids},
         "unit_locations": None,
         "max_new_tokens": max_new_tokens,
         "intervene_on_prompt": True,
-        "subspaces": subspaces_list,
+        "subspaces": subspaces_list,  # pass entire structure, using "internal_mag"
         "streamer": streamer,
         "eos_token_id": terminators,
         "early_stopping": True,
@@ -159,32 +151,62 @@ def generate(
 # UI Callbacks
 # --------------
 def filter_concepts(search_text: str):
+    """Return the first ~500 concepts that match (case-insensitive)."""
     if not search_text.strip():
         return concept_list[:500]
     filtered = [c for c in concept_list if search_text.lower() in c.lower()]
     return filtered[:500]
 
-def add_concept_to_list(selected_concept, magnitude, current_list):
-    """When 'Add Concept' is clicked, add the chosen concept and magnitude to subspaces."""
+def add_concept_to_list(selected_concept, user_slider_val, current_list):
+    """
+    When 'Add Concept' is clicked, add the chosen concept with the 
+    scaled magnitude to the subspaces list.
+
+    user_slider_val is from [-5..5], we multiply by 50 internally.
+    """
     if not selected_concept:
-        return current_list, current_list, gr.update(choices=[str(x["idx"]) for x in current_list])
+        return current_list, _build_table_data(current_list)
+
     concept_idx = concept_id_map[selected_concept]
-    new_entry = {"idx": concept_idx, "mag": magnitude}
-    updated_list = current_list + [new_entry]
 
-    remove_choices = [str(x["idx"]) for x in updated_list]
-    table_data = [[x['idx'], x['mag']] for x in updated_list]
-    return updated_list, table_data, gr.update(choices=remove_choices)
+    # Multiply slider by 50 internally
+    internal_mag = user_slider_val * 50
+
+    # We'll store both displayed magnitude (for the table) and the internal
+    # magnitude for the model. Also store 'text' for easy display.
+    new_entry = {
+        "text": selected_concept,
+        "idx": concept_idx,
+        "display_mag": user_slider_val,
+        "internal_mag": internal_mag,
+    }
 
-def remove_concept_from_list(rem_concept_idx_str, current_list):
-    """Remove the chosen concept from the list. Index is a string from remove_dropdown."""
-    if not rem_concept_idx_str:
-        return current_list, current_list, gr.update()
-    rem_idx = int(rem_concept_idx_str)
-    updated_list = [x for x in current_list if x["idx"] != rem_idx]
-    remove_choices = [str(x["idx"]) for x in updated_list]
-    table_data = [[x['idx'], x['mag']] for x in updated_list]
-    return updated_list, table_data, gr.update(choices=remove_choices)
+    # Avoid duplicates if you prefer:
+    # e.g. check if concept_idx already in current_list. We'll skip that for now.
+    updated_list = current_list + [new_entry]
+    return updated_list, _build_table_data(updated_list)
+
+def remove_selected_row(selected_rows, current_list):
+    """
+    Removes the row selected from the table.
+    selected_rows is a list of selected row indices, 
+    e.g. [1] meaning row with index 1 is selected.
+    """
+    if not selected_rows:
+        return current_list, _build_table_data(current_list)
+    row_idx = selected_rows[0]  # single selection
+    # Safely remove if in range
+    if 0 <= row_idx < len(current_list):
+        updated_list = current_list[:row_idx] + current_list[row_idx+1:]
+        return updated_list, _build_table_data(updated_list)
+    else:
+        return current_list, _build_table_data(current_list)
+
+def _build_table_data(subspaces):
+    """
+    Build a list of [concept_text, display_mag] to show in the table.
+    """
+    return [[x["text"], x["display_mag"]] for x in subspaces]
 
 def update_dropdown_choices(search_text):
     filtered = filter_concepts(search_text)
@@ -198,81 +220,96 @@ with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
     gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
 
-    selected_subspaces = gr.State([])
+    # If GPU is available, define a random default concept:
+    default_subspaces = []
+    if torch.cuda.is_available() and len(concept_list) > 0:
+        default_index = random.randint(0, len(concept_list) - 1)
+        default_concept = concept_list[default_index]
+        default_concept_idx = concept_id_map[default_concept]
+        # default slider is 3 => 3*50=150 internally
+        default_subspaces = [{
+            "text": default_concept,
+            "idx": default_concept_idx,
+            "display_mag": 3,       # what user sees
+            "internal_mag": 150.0,  # actual scaling
+        }]
+
+    # Keep state of subspaces
+    selected_subspaces = gr.State(default_subspaces)
 
     with gr.Row():
-        with gr.Column():
-            # Searching / selecting a concept
+        # Left column: Chat
+        with gr.Column(scale=5):
+            chat_interface = gr.ChatInterface(
+                fn=generate,
+                additional_inputs=[
+                    gr.Slider(
+                        label="Max new tokens",
+                        minimum=1,
+                        maximum=MAX_MAX_NEW_TOKENS,
+                        step=1,
+                        value=DEFAULT_MAX_NEW_TOKENS,
+                    ),
+                    selected_subspaces  # pass the entire subspaces list
+                ],
+                title="Model Steering with ReFT-r1 (16K concepts)",
+            )
+
+        # Right column: concept searching, adding, table display, removal
+        with gr.Column(scale=4):
+            gr.Markdown("## Steering Concepts")
             search_box = gr.Textbox(
                 label="Search concepts",
                 placeholder="Type text to filter concepts (e.g. 'sports')"
             )
             concept_dropdown = gr.Dropdown(
                 label="Filtered Concepts",
-                choices=[],
+                choices=[],  # dynamically populated
                 multiselect=False
             )
             concept_magnitude = gr.Slider(
-                label="Magnitude",
-                minimum=-300.0,
-                maximum=300.0,
+                label="Scaled Magnitude (multiplies by 50 internally)",
+                minimum=-5,
+                maximum=5,
                 step=1.0,
-                value=150.0
+                value=3
             )
             add_button = gr.Button("Add Concept")
 
-            # Removal
-            remove_dropdown = gr.Dropdown(
-                label="Remove from active list",
-                choices=[],
-                multiselect=False
-            )
-            remove_button = gr.Button("Remove Selected")
-
-        with gr.Column():
-            # Display currently active subspaces
+            # Current subspaces table
             active_subspaces_table = gr.Dataframe(
-                headers=["idx", "magnitude"],
-                datatype=["number", "number"],
+                headers=["Concept", "Magnitude (scaled)"],
+                datatype=["str", "number"],
                 interactive=False,
-                label="Active Concept Subspaces"
+                row_selectable="single",
+                label="Active Concept Subspaces",
+                value=_build_table_data(default_subspaces)
             )
 
-    # The Chat Interface
-    chat_interface = gr.ChatInterface(
-        fn=generate,
-        additional_inputs=[
-            gr.Slider(
-                label="Max new tokens",
-                minimum=1,
-                maximum=MAX_MAX_NEW_TOKENS,
-                step=1,
-                value=DEFAULT_MAX_NEW_TOKENS,
-            ),
-            selected_subspaces
-        ],
-        title="Model Steering with ReFT-r1 (16K concepts)",
-    )
+            remove_button = gr.Button("Remove Selected Row")
 
     gr.Markdown(LICENSE)
 
     # Wire up events
+    # Whenever user types in search_box, update concept_dropdown
     search_box.change(
         fn=update_dropdown_choices,
         inputs=[search_box],
         outputs=[concept_dropdown]
     )
 
+    # Add concept
     add_button.click(
         fn=add_concept_to_list,
         inputs=[concept_dropdown, concept_magnitude, selected_subspaces],
-        outputs=[selected_subspaces, active_subspaces_table, remove_dropdown],
+        outputs=[selected_subspaces, active_subspaces_table],
     )
 
+    # Remove selected row from table
     remove_button.click(
-        fn=remove_concept_from_list,
-        inputs=[remove_dropdown, selected_subspaces],
-        outputs=[selected_subspaces, active_subspaces_table, remove_dropdown],
+        fn=remove_selected_row,
+        inputs=[active_subspaces_table, selected_subspaces],
+        outputs=[selected_subspaces, active_subspaces_table],
     )
 
     demo.queue(max_size=20).launch()