claude 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 = 128  # smaller default to save memory
+DEFAULT_MAX_NEW_TOKENS = 256  # smaller default to save memory
 MAX_INPUT_TOKEN_LENGTH = 4096
 
 def load_jsonl(jsonl_path):
@@ -133,6 +133,9 @@ def generate(
         partial_text.append(token_str)
         yield "".join(partial_text)
 
+def _build_remove_choices(subspaces):
+    return [f"(+{x['display_mag']:.1f}*) {x['text']}" for x in subspaces]
+
 def filter_concepts(search_text: str):
     if not search_text.strip():
         return concept_list[:500]
@@ -140,11 +143,6 @@ 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, gr.update(choices=_build_remove_choices(current_list))
 
@@ -156,78 +154,69 @@ def add_concept_to_list(selected_concept, user_slider_val, current_list):
         "display_mag": user_slider_val,
         "internal_mag": internal_mag,
     }
-    updated_list = current_list + [new_entry]
+    # Add to the beginning of the list
+    updated_list = [new_entry] + current_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, gr.update(choices=_build_remove_choices(current_list))
-    updated_list = [x for x in current_list if x["text"] != selected_text]
+    
+    # Remove based on the full formatted text
+    updated_list = [x for x in current_list if f"(+{x['display_mag']:.1f}*) {x['text']}" != selected_text]
     return updated_list, gr.update(choices=_build_remove_choices(updated_list))
 
-def _build_remove_choices(subspaces):
-    return [x["text"] for x in subspaces]
-
 def update_dropdown_choices(search_text):
     filtered = filter_concepts(search_text)
     return gr.update(choices=filtered)
 
 with gr.Blocks(css="style.css") as demo:
-
-    # Pre-populate with a random concept if available
-    default_subspaces = []
-    if pv_model and concept_list:
-        default_concept = "words related to time travel and its consequences"
-        default_subspaces = [{
-            "text": default_concept,
-            "idx": concept_id_map[default_concept],
-            "display_mag": 3,
-            "internal_mag": 150.0,
-        }]
-
-    selected_subspaces = gr.State(default_subspaces)
+    # Remove default subspaces
+    selected_subspaces = gr.State([])
+    
     with gr.Row():
         # Left side: bigger chat area
         with gr.Column(scale=7):
             chat_interface = gr.ChatInterface(
                 fn=generate,
-                title="LM Steering with ReFT-r1 (16K concepts)",
+                title="Language Model Concept Steering",
+                description="Steer responses by selecting concepts on the right →",
                 type="messages",
                 additional_inputs=[selected_subspaces],
             )
+        
         # Right side: concept management
         with gr.Column(scale=3):
-            gr.Markdown("# Steering Concepts")
-            search_box = gr.Textbox(
-                label="Search concepts",
-                placeholder="e.g. 'time travel'"
-            )
-            concept_dropdown = gr.Dropdown(
-                label="Filtered Concepts",
-                choices=[]
-            )
-            concept_magnitude = gr.Slider(
-                label="Steering Factor",
-                minimum=-5,
-                maximum=5,
-                step=1,
-                value=3
-            )
-            add_button = gr.Button("Add Concept")
+            gr.Markdown("## Steer Model Responses")
+            
+            # Concept Search and Selection
+            with gr.Group():
+                search_box = gr.Textbox(
+                    label="Search Concepts",
+                    placeholder="Find concepts to steer the model (e.g. 'time travel')",
+                )
+                concept_dropdown = gr.Dropdown(
+                    label="Select a Concept",
+                    interactive=True,
+                )
+                concept_magnitude = gr.Slider(
+                    label="Steering Intensity",
+                    minimum=-5,
+                    maximum=5,
+                    step=0.1,  # Allow 1 decimal point
+                    value=3,
+                )
+                add_button = gr.Button("Add Concept to Steering")
 
-            # Row with the remove dropdown + button
-            with gr.Row():
+            # Current Steering Concepts
+            gr.Markdown("## Current Steering Concepts")
+            with gr.Group():
                 remove_dropdown = gr.Dropdown(
-                    label="Remove concept",
-                    choices=_build_remove_choices(default_subspaces),
-                    multiselect=False
+                    label="Select a Current Steering Concept to Stop",
+                    choices=[],
+                    multiselect=False,
                 )
-                remove_button = gr.Button("Remove", variant="secondary")
+                remove_button = gr.Button("Remove Current Steering Concept", variant="secondary")
 
     # Wire up events
     # When the search box changes, update the concept dropdown choices:
@@ -253,4 +242,5 @@ with gr.Blocks(css="style.css") as demo:
         [selected_subspaces, remove_dropdown]
     )
 
-    demo.launch()
+    demo.launch(share=True)
+