Ahhhhh lol

app.py

@@ -2,7 +2,6 @@
 import gradio as gr
 from concept_steerer import ConceptSteerer
 
-# Global model instance (loaded once)
 steerer = None
 
 def get_steerer():
@@ -13,191 +12,159 @@ def get_steerer():
 
 def create_concept(name, pos_examples, neg_examples):
     if not name.strip():
-        return "❌ Concept name cannot be empty."
+        return "❌ Concept name is required."
     pos_list = [p.strip() for p in pos_examples.strip().split('\n') if p.strip()]
     neg_list = [n.strip() for n in neg_examples.strip().split('\n') if n.strip()]
     if not pos_list or not neg_list:
-        return "❌ Please provide at least one example for both positive and negative prompts."
-    
+        return "❌ Provide at least one positive and one negative example."
     try:
-        steerer = get_steerer()
-        steerer.register_concept(name.strip(), pos_list, neg_list, layer=-2)
-        return f"✅ Concept '{name.strip()}' registered successfully!"
+        s = get_steerer()
+        s.register_concept(name.strip(), pos_list, neg_list, layer=-2)
+        return f"✅ Concept '{name}' registered!"
     except Exception as e:
-        return f"❌ Error: {str(e)}"
+        return f"❌ Error: {e}"
 
-def update_ui():
-    """Returns updated UI components based on current concepts."""
-    steerer = get_steerer()
-    concepts = steerer.get_concept_names()
+def build_generation_ui():
+    """Returns prompt input, list of sliders, and output box."""
+    s = get_steerer()
+    concepts = s.get_concept_names()
+    
+    prompt = gr.Textbox(label="Prompt", lines=2, placeholder="Enter your prompt...")
+    output = gr.Textbox(label="Output", lines=8, interactive=False)
     
     if not concepts:
-        return (
-            gr.Markdown("💡 No concepts yet. Create one in the 'Create Concepts' tab!"),
-            gr.Row(visible=False),
-            gr.Button(visible=True)
-        )
+        sliders = [gr.Markdown("ℹ️ No concepts registered. Go to 'Create Concepts' tab.")]
+        return prompt, sliders, output
     
-    # Build sliders dynamically
     sliders = []
-    with gr.Row():
-        for concept in concepts:
-            slider = gr.Slider(-5.0, 5.0, value=0.0, label=concept, step=0.1)
-            sliders.append(slider)
-    
-    return (
-        gr.Markdown(visible=False),  # hide placeholder
-        gr.Row(*sliders, visible=True),
-        gr.Button(visible=True)
-    )
-
-def generate_with_steering(prompt, *slider_values):
-    steerer = get_steerer()
-    concepts = steerer.get_concept_names()
-    if not concepts:
-        # No steering: just generate normally
-        return steerer.generate(prompt, max_new_tokens=150)
+    for concept in concepts:
+        slider = gr.Slider(-5.0, 5.0, value=0.0, step=0.1, label=concept)
+        sliders.append(slider)
     
-    steering_config = {name: float(val) for name, val in zip(concepts, slider_values)}
+    return prompt, sliders, output
+
+def generate_text(prompt, *slider_values):
+    s = get_steerer()
+    concepts = s.get_concept_names()
+    steering = {}
+    for name, val in zip(concepts, slider_values):
+        if abs(val) > 1e-6:  # skip zero
+            steering[name] = float(val)
     try:
-        output = steerer.generate(prompt, steering_config=steering_config, max_new_tokens=150)
-        return output
+        return s.generate(prompt, steering_config=steering, max_new_tokens=150)
     except Exception as e:
-        return f"❌ Generation error: {str(e)}"
-
-# Build the interface
-with gr.Blocks(title="🧠 LLM Concept Steering") as demo:
-    gr.Markdown("# 🧠 LLM Concept Steering Demo")
-    gr.Markdown("Steer Llama 3.2 1B's behavior using activation vectors — no training needed!")
+        return f"❌ Generation failed: {e}"
 
+# Main interface
+with gr.Blocks() as demo:
+    gr.Markdown("# 🧠 LLM Concept Steering — Working Edition")
+    
     with gr.Tab("Create Concepts"):
         with gr.Row():
             with gr.Column():
-                name_input = gr.Textbox(label="Concept Name", placeholder="e.g., Formality")
-                pos_input = gr.Textbox(label="Positive Prompts (one per line)", lines=4)
-                neg_input = gr.Textbox(label="Negative Prompts (one per line)", lines=4)
+                name_in = gr.Textbox(label="Concept Name")
+                pos_in = gr.Textbox(label="Positive Prompts (one per line)", lines=4)
+                neg_in = gr.Textbox(label="Negative Prompts (one per line)", lines=4)
                 create_btn = gr.Button("Register Concept")
             with gr.Column():
-                status = gr.Textbox(label="Status", interactive=False)
+                status_out = gr.Textbox(label="Status", interactive=False)
         create_btn.click(
             create_concept,
-            inputs=[name_input, pos_input, neg_input],
-            outputs=status
+            inputs=[name_in, pos_in, neg_in],
+            outputs=status_out
         )
-
-    with gr.Tab("Generate with Steering"):
-        prompt_box = gr.Textbox(label="Enter your prompt", lines=3, placeholder="e.g., Write a short story about a robot.")
-        gen_btn = gr.Button("Generate")
+    
+    with gr.Tab("Generate"):
+        # We will dynamically replace this row
+        with gr.Row() as dynamic_row:
+            pass
         
-        # Placeholder for dynamic sliders
-        placeholder_msg = gr.Markdown("Loading concepts...")
-        slider_row = gr.Row(visible=False)
+        gen_btn = gr.Button("Generate")
         output_box = gr.Textbox(label="Output", lines=8, interactive=False)
 
-        # Update sliders when tab is loaded or after concept creation
-        demo.load(update_ui, None, [placeholder_msg, slider_row, gen_btn])
-
-        # Connect generate button — we pass all children of slider_row as inputs
-        # But since we can't reference them before creation, we use a trick:
-        # We'll re-render the generate button's dependency after UI update.
-        # Simpler: just call generate_with_steering with all possible inputs.
-        # Gradio handles missing inputs gracefully if we set default None.
-
-        # Instead, we use a wrapper that fetches current slider values via State (not needed here)
-        # Better: restructure to use a single State for config — but for simplicity, we do this:
-
-        # We'll make the generate button depend on all possible sliders by rebuilding the UI
-        # But an easier fix: use a single "Generate" that reads from a hidden state.
-        # However, the cleanest modern way is to use **gr.on()** or just accept that we need to
-        # capture the slider row's children.
-
-        # WORKING SOLUTION: Use a lambda that captures the current UI
-        def make_generate_fn():
-            # This will be called after sliders exist
-            children = slider_row.children
-            if not children:
-                inputs = [prompt_box]
-            else:
-                inputs = [prompt_box] + children
-            gen_btn.click(
-                generate_with_steering,
-                inputs=inputs,
-                outputs=output_box
-            )
-
-        # Unfortunately, we can't easily re-bind click after load in pure Python.
-        # So here's a robust alternative: use a single button and read concepts at generate time.
-
-        # REVISED PLAN: Don't use dynamic inputs to .click(). Instead, store slider values in State.
-        # But to keep it simple and working NOW, let's use a different approach:
-
-        # ✅ FINAL WORKING FIX: Use a single "steering_dict" stored in a Gradio State.
-        # We'll add invisible number inputs that update a dict via change handlers.
-
-        pass  # We'll refactor below
-
-# ---- REFACTORED WORKING VERSION BELOW ----
-
-steering_state = gr.State({})
-
-def update_steering_state(name, value, state):
-    state = state or {}
-    state[name] = float(value)
-    return state
-
-def generate_from_state(prompt, state):
-    steerer = get_steerer()
-    try:
-        output = steerer.generate(prompt, steering_config=state or {}, max_new_tokens=150)
-        return output
-    except Exception as e:
-        return f"❌ Error: {str(e)}"
+        # Function to refresh the Generate tab UI
+        def refresh_ui():
+            prompt, sliders, _ = build_generation_ui()
+            return [prompt] + sliders
+
+        # Load initial UI
+        demo.load(
+            refresh_ui,
+            inputs=None,
+            outputs=None,
+            # We need to define the exact output components.
+            # So we pre-define a fixed max number of sliders (e.g., 10)
+        )
 
-with gr.Blocks() as demo2:
-    gr.Markdown("# 🧠 LLM Concept Steering (Fixed!)")
-    gr.Markdown("Now fully compatible with modern Gradio.")
+# 🚨 CRITICAL: Gradio requires fixed output structure.
+# So we pre-define up to 10 sliders (more than enough for a demo).
+MAX_CONCEPTS = 10
 
+with gr.Blocks() as demo_fixed:
+    gr.Markdown("# 🧠 LLM Concept Steering — Working Edition")
+    
     with gr.Tab("Create Concepts"):
         with gr.Row():
             with gr.Column():
                 name_in = gr.Textbox(label="Concept Name")
-                pos_in = gr.Textbox(label="Positive Prompts", lines=3)
-                neg_in = gr.Textbox(label="Negative Prompts", lines=3)
-                create_btn2 = gr.Button("Register")
+                pos_in = gr.Textbox(label="Positive Prompts (one per line)", lines=4)
+                neg_in = gr.Textbox(label="Negative Prompts (one per line)", lines=4)
+                create_btn = gr.Button("Register Concept")
             with gr.Column():
-                status2 = gr.Textbox(label="Result")
-        create_btn2.click(create_concept, [name_in, pos_in, neg_in], status2)
-
+                status_out = gr.Textbox(label="Status", interactive=False)
+        create_btn.click(
+            create_concept,
+            inputs=[name_in, pos_in, neg_in],
+            outputs=status_out
+        ).then(  # After registering, refresh the other tab
+            None, None, None, 
+            _js="() => { document.querySelectorAll('.tabs button')[1].click(); }"  # Switch to Generate tab
+        )
+    
     with gr.Tab("Generate"):
-        prompt_in = gr.Textbox(label="Prompt", lines=2)
-        state_holder = gr.State({})
-        output_out = gr.Textbox(label="Output", lines=8)
-
-        # Container for sliders
-        slider_box = gr.Box()
-
-        def build_sliders():
-            steerer = get_steerer()
-            concepts = steerer.get_concept_names()
-            if not concepts:
-                return [gr.Markdown("No concepts. Create one first!")]
-            components = []
-            for c in concepts:
-                num = gr.Number(value=0.0, label=c, minimum=-5, maximum=5, step=0.1)
-                num.change(
-                    lambda v, name=c, s=state_holder: update_steering_state(name, v, s),
-                    num,
-                    state_holder
-                )
-                components.append(num)
-            return components
-
-        demo2.load(build_sliders, None, slider_box)
+        prompt_input = gr.Textbox(label="Prompt", lines=2, placeholder="Enter your prompt...")
+        
+        # Pre-define 10 optional sliders
+        sliders = []
+        for i in range(MAX_CONCEPTS):
+            slider = gr.Slider(-5.0, 5.0, value=0.0, step=0.1, label=f"Concept {i+1}", visible=False)
+            sliders.append(slider)
+        
+        gen_btn = gr.Button("Generate")
+        output_box = gr.Textbox(label="Output", lines=8, interactive=False)
 
-        gen_btn2 = gr.Button("Generate")
-        gen_btn2.click(generate_from_state, [prompt_in, state_holder], output_out)
+        def update_sliders():
+            s = get_steerer()
+            concepts = s.get_concept_names()
+            updates = []
+            for i in range(MAX_CONCEPTS):
+                if i < len(concepts):
+                    updates.append(gr.Slider(visible=True, label=concepts[i]))
+                else:
+                    updates.append(gr.Slider(visible=False))
+            return updates
+
+        def gen_wrapper(prompt, *slider_vals):
+            s = get_steerer()
+            concepts = s.get_concept_names()
+            steering = {
+                name: float(val) 
+                for name, val in zip(concepts, slider_vals[:len(concepts)])
+                if abs(val) > 1e-6
+            }
+            return s.generate(prompt, steering_config=steering, max_new_tokens=150)
+
+        # Update sliders on tab load
+        demo_fixed.load(update_sliders, None, sliders)
+        
+        # Connect generate
+        gen_btn.click(
+            gen_wrapper,
+            inputs=[prompt_input] + sliders,
+            outputs=output_box
+        )
 
-demo = demo2
+demo = demo_fixed
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", share=False)