o1 impl

app.py

@@ -45,20 +45,20 @@ class Steer(pv.SourcelessIntervention):
 
     def forward(self, base, source=None, subspaces=None):
         # subspaces is a list of dicts: 
-        #   each has {"idx": int, "internal_mag": float, ...}
+        #   each has {"idx": int, "internal_mag": float, "text": str, ...}
         steer_vec = base
         if subspaces is not None:
             for sp in subspaces:
                 idx = sp["idx"]
-                # Use the internal magnitude for actual steering
-                mag = sp["internal_mag"] 
+                mag = sp["internal_mag"]  # the true scaling factor
                 steering_vec = mag * self.proj.weight[idx].unsqueeze(dim=0)
                 steer_vec = steer_vec + steering_vec
         return steer_vec
 
-# ---------------------------------------------------
-# Load Model & Dictionary if GPU is available
-# ---------------------------------------------------
+
+# ------------------------------------------
+# Load the Model & Dictionary if GPU exists
+# ------------------------------------------
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo won't perform well on CPU.</p>"
 
@@ -91,10 +91,9 @@ if torch.cuda.is_available():
 terminators = [tokenizer.eos_token_id]
 
 
-# ---------------------------------------------------------------------
-# The main generation function, limiting to last 3 conversation turns
-# and then using apply_chat_template
-# ---------------------------------------------------------------------
+# --------------------------------------------
+# Main generation function: keep last 3 turns
+# --------------------------------------------
 @spaces.GPU
 def generate(
     message: str,
@@ -107,31 +106,36 @@ def generate(
     start_idx = max(0, len(chat_history) - 3)
     recent_history = chat_history[start_idx:]
 
-    # Build a list of messages
-    # each tuple is (user_message, model_message)
+    # Convert (user_msg, model_msg) => 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})
 
-    # Now append the new user message
+    # Add the new user message
     messages.append({"role": "user", "content": message})
 
-    input_ids = torch.tensor([tokenizer.apply_chat_template(
-        messages, tokenize=True, add_generation_prompt=True)]).cuda()
+    # Apply the chat template (some HF models expect "assistant" instead of "model")
+    # but let's keep "model" to match your code, if that is required.
+    prompt_dict = tokenizer.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True
+    )
+    input_ids = torch.tensor([prompt_dict["input_ids"]]).cuda()
+    attention_mask = torch.tensor([prompt_dict["attention_mask"]]).cuda()
 
-    # Possibly trim if over max length
+    # Possibly trim if too long
     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},
+        "base": {"input_ids": input_ids, "attention_mask": attention_mask},
         "unit_locations": None,
         "max_new_tokens": max_new_tokens,
         "intervene_on_prompt": True,
-        "subspaces": subspaces_list,  # pass entire structure, using "internal_mag"
+        "subspaces": subspaces_list,
         "streamer": streamer,
         "eos_token_id": terminators,
         "early_stopping": True,
@@ -147,11 +151,11 @@ def generate(
         yield "".join(partial_text)
 
 
-# --------------
+# ----------------
 # UI Callbacks
-# --------------
+# ----------------
 def filter_concepts(search_text: str):
-    """Return the first ~500 concepts that match (case-insensitive)."""
+    """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()]
@@ -159,87 +163,75 @@ def filter_concepts(search_text: str):
 
 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.
+    user_slider_val is from [-5..5]. We multiply by 50 internally to get the real magnitude.
     """
     if not selected_concept:
-        return current_list, _build_table_data(current_list)
+        return current_list, _build_table_data(current_list), gr.update(choices=_build_remove_choices(current_list))
 
     concept_idx = concept_id_map[selected_concept]
-
-    # 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.
+    internal_mag = user_slider_val * 50  # scale by 50
     new_entry = {
         "text": selected_concept,
         "idx": concept_idx,
         "display_mag": user_slider_val,
         "internal_mag": internal_mag,
     }
-
-    # 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)
+    return (
+        updated_list,
+        _build_table_data(updated_list),
+        gr.update(choices=_build_remove_choices(updated_list))
+    )
 
-def remove_selected_row(selected_rows, current_list):
+def remove_concept_from_list(concept_to_remove, 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.
+    Remove the chosen concept name from the list.
     """
-    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)
+    if not concept_to_remove:
+        return current_list, _build_table_data(current_list), gr.update(choices=_build_remove_choices(current_list))
+
+    updated_list = [x for x in current_list if x["text"] != concept_to_remove]
+    return (
+        updated_list,
+        _build_table_data(updated_list),
+        gr.update(choices=_build_remove_choices(updated_list))
+    )
 
 def _build_table_data(subspaces):
-    """
-    Build a list of [concept_text, display_mag] to show in the table.
-    """
+    """Return [[concept_name, scaled_mag], ...] for display."""
     return [[x["text"], x["display_mag"]] for x in subspaces]
 
+def _build_remove_choices(subspaces):
+    """Return concept names for the remove dropdown."""
+    return [x["text"] for x in subspaces]
+
 def update_dropdown_choices(search_text):
     filtered = filter_concepts(search_text)
     return gr.update(choices=filtered)
 
-
-# -------------------------
-# Build the Gradio Blocks
-# -------------------------
+# --------------------------------------------------------------------
+# Build the Interface
+# --------------------------------------------------------------------
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
     gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
 
-    # If GPU is available, define a random default concept:
+    # If GPU is available, pick a random concept as default
     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_concept = random.choice(concept_list)
         default_subspaces = [{
             "text": default_concept,
-            "idx": default_concept_idx,
-            "display_mag": 3,       # what user sees
-            "internal_mag": 150.0,  # actual scaling
+            "idx": concept_id_map[default_concept],
+            "display_mag": 3,       # user sees 3
+            "internal_mag": 150.0,  # actual factor
         }]
 
-    # Keep state of subspaces
     selected_subspaces = gr.State(default_subspaces)
 
     with gr.Row():
-        # Left column: Chat
         with gr.Column(scale=5):
+            # Use type="messages" to avoid tuple-format deprecation warnings
             chat_interface = gr.ChatInterface(
                 fn=generate,
                 additional_inputs=[
@@ -250,14 +242,14 @@ with gr.Blocks(css="style.css") as demo:
                         step=1,
                         value=DEFAULT_MAX_NEW_TOKENS,
                     ),
-                    selected_subspaces  # pass the entire subspaces list
+                    selected_subspaces
                 ],
                 title="Model Steering with ReFT-r1 (16K concepts)",
+                type="messages",  # <--- uses openai-style 'role' and 'content'
             )
-
-        # 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')"
@@ -268,30 +260,35 @@ with gr.Blocks(css="style.css") as demo:
                 multiselect=False
             )
             concept_magnitude = gr.Slider(
-                label="Scaled Magnitude (multiplies by 50 internally)",
+                label="Scaled Magnitude (×50 internally)",
                 minimum=-5,
                 maximum=5,
-                step=1.0,
+                step=1,
                 value=3
             )
             add_button = gr.Button("Add Concept")
 
-            # Current subspaces table
+            # Show the table of active subspaces
             active_subspaces_table = gr.Dataframe(
                 headers=["Concept", "Magnitude (scaled)"],
                 datatype=["str", "number"],
+                value=_build_table_data(default_subspaces),
                 interactive=False,
-                row_selectable="single",
-                label="Active Concept Subspaces",
-                value=_build_table_data(default_subspaces)
+                label="Active Concept Subspaces"
             )
 
-            remove_button = gr.Button("Remove Selected Row")
+            # Remove concept by name
+            remove_dropdown = gr.Dropdown(
+                label="Remove a concept",
+                choices=_build_remove_choices(default_subspaces),
+                multiselect=False
+            )
+            remove_button = gr.Button("Remove Selected Concept")
 
     gr.Markdown(LICENSE)
 
     # Wire up events
-    # Whenever user types in search_box, update concept_dropdown
+    # Update concept dropdown when user types in search
     search_box.change(
         fn=update_dropdown_choices,
         inputs=[search_box],
@@ -302,14 +299,14 @@ with gr.Blocks(css="style.css") as demo:
     add_button.click(
         fn=add_concept_to_list,
         inputs=[concept_dropdown, concept_magnitude, selected_subspaces],
-        outputs=[selected_subspaces, active_subspaces_table],
+        outputs=[selected_subspaces, active_subspaces_table, remove_dropdown],
     )
 
-    # Remove selected row from table
+    # Remove a concept
     remove_button.click(
-        fn=remove_selected_row,
-        inputs=[active_subspaces_table, selected_subspaces],
-        outputs=[selected_subspaces, active_subspaces_table],
+        fn=remove_concept_from_list,
+        inputs=[remove_dropdown, selected_subspaces],
+        outputs=[selected_subspaces, active_subspaces_table, remove_dropdown],
     )
 
     demo.queue(max_size=20).launch()