Update app.py

app.py

@@ -4,7 +4,7 @@ import requests
 def get_features(text: str):
     url = "https://www.neuronpedia.org/api/search-with-topk"
     payload = {
-        "modelId": "gemma-2-2b",
+        "modelId": "gemma-2-2b", 
         "text": text,
         "layer": "20-gemmascope-res-16k"
     }
@@ -29,15 +29,18 @@ def create_dashboard(feature_id: int) -> str:
         </div>
     """
 
+def handle_feature_click(feature_id):
+    return create_dashboard(feature_id)
+
 def analyze_text(text: str):
     if not text:
-        return gr.update(visible=False), ""
-    
+        return [], ""
+        
     features_data = get_features(text)
     if not features_data:
-        return gr.update(visible=False), ""
+        return [], ""
 
-    html = "<div class='features-list'>"
+    features = []
     first_feature_id = None
     
     for result in features_data['results']:
@@ -45,24 +48,22 @@ def analyze_text(text: str):
             continue
             
         token = result['token']
-        html += f"<h3>{token}</h3>"
+        token_features = []
         
         for feature in result['top_features'][:3]:
             feature_id = feature['feature_index']
             if first_feature_id is None:
                 first_feature_id = feature_id
-            
-            html += f"""
-            <button onclick='document.dispatchEvent(new CustomEvent("select_feature", 
-                {{detail: {{feature_id: {feature_id}}}}}))' class='feature-button'>
-                Feature {feature_id} (Activation: {feature['activation_value']:.2f})
-            </button>
-            """
-    
-    html += "</div>"
-    initial_dashboard = create_dashboard(first_feature_id) if first_feature_id else ""
+                
+            token_features.append({
+                "token": token,
+                "id": feature_id,
+                "activation": feature['activation_value']
+            })
+        
+        features.append({"token": token, "features": token_features})
     
-    return html, initial_dashboard
+    return features, create_dashboard(first_feature_id) if first_feature_id else ""
 
 css = """
 @import url('https://fonts.googleapis.com/css2?family=Open+Sans:wght@300;400;600;700&display=swap');
@@ -75,19 +76,21 @@ body { font-family: 'Open Sans', sans-serif !important; }
     background-color: #ffffff;
 }
 
-.features-list h3 {
-    margin-top: 1rem;
+.token-header {
+    font-size: 1.25rem;
     font-weight: 600;
+    margin-top: 1rem;
+    margin-bottom: 0.5rem;
 }
 
 .feature-button {
-    display: block;
-    margin: 0.5rem 0;
+    display: inline-block;
+    margin: 0.25rem;
     padding: 0.5rem 1rem;
     background-color: #f3f4f6;
     border: 1px solid #e5e7eb;
     border-radius: 0.375rem;
-    cursor: pointer;
+    font-size: 0.875rem;
 }
 
 .feature-button:hover {
@@ -95,23 +98,12 @@ body { font-family: 'Open Sans', sans-serif !important; }
 }
 """
 
-theme = gr.themes.Soft(
-    primary_hue=gr.themes.colors.Color(
-        name="blue",
-        c50="#eef1ff", c100="#e0e5ff", c200="#c3cbff",
-        c300="#a5b2ff", c400="#8798ff", c500="#6a7eff",
-        c600="#3452db", c700="#2a41af", c800="#1f3183",
-        c900="#152156", c950="#0a102b",
-    )
-)
-
-def update_dashboard(feature_id: int):
-    return create_dashboard(feature_id)
-
-with gr.Blocks(theme=theme, css=css) as demo:
+with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     gr.Markdown("# Brand Analyzer", elem_classes="text-2xl font-bold mb-2")
     gr.Markdown("*Analyze text using Gemma's interpretable neural features*", elem_classes="text-gray-600 mb-6")
     
+    features_state = gr.State([])
+    
     with gr.Row():
         with gr.Column(scale=1):
             input_text = gr.Textbox(
@@ -126,17 +118,32 @@ with gr.Blocks(theme=theme, css=css) as demo:
             )
         
         with gr.Column(scale=2):
-            features_html = gr.HTML()
-            dashboard = gr.HTML()
+            @gr.render(every=0.5)
+            def render_features():
+                features = features_state.value
+                if not features:
+                    return
+                
+                for token_group in features:
+                    gr.Markdown(f"### {token_group['token']}")
+                    with gr.Row():
+                        for feature in token_group['features']:
+                            btn = gr.Button(
+                                f"Feature {feature['id']} (Activation: {feature['activation']:.2f})",
+                                elem_classes=["feature-button"]
+                            )
+                            btn.click(
+                                fn=lambda fid=feature['id']: handle_feature_click(fid),
+                                outputs=dashboard
+                            )
             
-            # Handle feature selection via JavaScript
-            dashboard.change(fn=update_dashboard, inputs=gr.Textbox(visible=False), outputs=dashboard)
+            dashboard = gr.HTML()
     
     analyze_btn.click(
         fn=analyze_text,
         inputs=[input_text],
-        outputs=[features_html, dashboard]
+        outputs=[features_state, dashboard]
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=False)