Update app.py

app.py

@@ -1,8 +1,7 @@
 import gradio as gr
 import requests
-from typing import Dict, List, Tuple
 
-def get_features(text: str) -> Dict:
+def get_features(text: str):
     url = "https://www.neuronpedia.org/api/search-with-topk"
     payload = {
         "modelId": "gemma-2-2b",
@@ -30,43 +29,70 @@ def create_dashboard(feature_id: int) -> str:
         </div>
     """
 
-def process_features(text: str) -> Tuple[gr.Column, str]:
+def analyze_text(text: str):
     if not text:
-        return gr.Column(), ""
-        
+        return gr.update(visible=False), ""
+    
     features_data = get_features(text)
     if not features_data:
-        return gr.Column(), ""
-    
+        return gr.update(visible=False), ""
+
+    html = "<div class='features-list'>"
     first_feature_id = None
-    with gr.Column() as col:
-        for result in features_data['results']:
-            if result['token'] == '<bos>':
-                continue
-                
-            gr.Markdown(f"### {result['token']}")
-            for i, feature in enumerate(result['top_features'][:3]):
-                feature_id = feature['feature_index']
-                if first_feature_id is None:
-                    first_feature_id = feature_id
-                gr.Button(
-                    f"Feature {feature_id} (Activation: {feature['activation_value']:.2f})",
-                    elem_id=str(feature_id)
-                ).click(
-                    fn=lambda fid=feature_id: create_dashboard(fid),
-                    outputs=dashboard
-                )
     
-    return col, create_dashboard(first_feature_id) if first_feature_id else ""
+    for result in features_data['results']:
+        if result['token'] == '<bos>':
+            continue
+            
+        token = result['token']
+        html += f"<h3>{token}</h3>"
+        
+        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 ""
+    
+    return html, initial_dashboard
 
 css = """
 @import url('https://fonts.googleapis.com/css2?family=Open+Sans:wght@300;400;600;700&display=swap');
+
 body { font-family: 'Open Sans', sans-serif !important; }
+
 .dashboard-container {
     border: 1px solid #e0e5ff;
     border-radius: 8px;
     background-color: #ffffff;
 }
+
+.features-list h3 {
+    margin-top: 1rem;
+    font-weight: 600;
+}
+
+.feature-button {
+    display: block;
+    margin: 0.5rem 0;
+    padding: 0.5rem 1rem;
+    background-color: #f3f4f6;
+    border: 1px solid #e5e7eb;
+    border-radius: 0.375rem;
+    cursor: pointer;
+}
+
+.feature-button:hover {
+    background-color: #e5e7eb;
+}
 """
 
 theme = gr.themes.Soft(
@@ -79,6 +105,9 @@ theme = gr.themes.Soft(
     )
 )
 
+def update_dashboard(feature_id: int):
+    return create_dashboard(feature_id)
+
 with gr.Blocks(theme=theme, 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")
@@ -97,13 +126,16 @@ with gr.Blocks(theme=theme, css=css) as demo:
             )
         
         with gr.Column(scale=2):
-            features_col = gr.Column()
+            features_html = gr.HTML()
             dashboard = gr.HTML()
+            
+            # Handle feature selection via JavaScript
+            dashboard.change(fn=update_dashboard, inputs=gr.Textbox(visible=False), outputs=dashboard)
     
     analyze_btn.click(
-        fn=process_features,
+        fn=analyze_text,
         inputs=[input_text],
-        outputs=[features_col, dashboard]
+        outputs=[features_html, dashboard]
     )
 
 if __name__ == "__main__":