Update app.py

app.py

@@ -1,21 +1,43 @@
 import requests
 import gradio as gr
 
-def get_features(text: str):
-    url = "https://www.neuronpedia.org/api/search-with-topk"
-    payload = {
-        "modelId": "gemma-2-2b", 
-        "text": text,
-        "layer": "20-gemmascope-res-16k"
-    }
+def parse_api_response(data):
+    features_by_token = {}
+    for result in data['results']:
+        token = result['token']
+        if token == '<bos>':
+            continue
+        features_by_token[token] = []
+        for feature in result['top_features']:
+            features_by_token[token].append({
+                'id': feature['feature_index'],
+                'activation': feature['activation_value']
+            })
+    return features_by_token
+
+def analyze_text(text: str):
+    if not text:
+        return None, ""
+
     try:
-        response = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
+        response = requests.post(
+            "https://www.neuronpedia.org/api/search-with-topk",
+            headers={"Content-Type": "application/json"},
+            json={
+                "modelId": "gemma-2-2b",
+                "text": text,
+                "layer": "20-gemmascope-res-16k"
+            }
+        )
         response.raise_for_status()
-        return response.json()
+        features = parse_api_response(response.json())
+        first_feature = next(iter(next(iter(features.values()))))
+        dashboard = create_dashboard(first_feature['id'])
+        return features, dashboard
     except Exception as e:
-        return None
+        return None, str(e)
 
-def create_dashboard(feature_id: int) -> str:
+def create_dashboard(feature_id: int):
     return f"""
         <div class="dashboard-container p-4">
             <h3 class="text-lg font-semibold mb-4">Feature {feature_id} Dashboard</h3>
@@ -29,119 +51,77 @@ 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 [], ""
-        
-    features_data = get_features(text)
-    if not features_data:
-        return [], ""
-
-    features = []
-    first_feature_id = None
-    
-    for result in features_data['results']:
-        if result['token'] == '<bos>':
-            continue
-            
-        token = result['token']
-        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
-                
-            token_features.append({
-                "token": token,
-                "id": feature_id,
-                "activation": feature['activation_value']
-            })
-        
-        features.append({"token": token, "features": token_features})
-    
-    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');
-
-body { font-family: 'Open Sans', sans-serif !important; }
-
-.dashboard-container {
-    border: 1px solid #e0e5ff;
-    border-radius: 8px;
-    background-color: #ffffff;
-}
-
-.token-header {
+.token-section { margin-bottom: 1.5rem; }
+.token-header { 
     font-size: 1.25rem;
     font-weight: 600;
-    margin-top: 1rem;
-    margin-bottom: 0.5rem;
+    margin-bottom: 0.75rem;
 }
-
 .feature-button {
-    display: inline-block;
-    margin: 0.25rem;
+    background: #f8fafc;
+    border: 1px solid #e2e8f0;
+    border-radius: 0.5rem;
     padding: 0.5rem 1rem;
-    background-color: #f3f4f6;
-    border: 1px solid #e5e7eb;
-    border-radius: 0.375rem;
+    margin: 0.25rem;
     font-size: 0.875rem;
+    cursor: pointer;
 }
-
-.feature-button:hover {
-    background-color: #e5e7eb;
+.feature-button:hover { background: #f1f5f9; }
+.feature-button.selected {
+    background: #e0e7ff;
+    border-color: #6366f1;
 }
 """
 
-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")
+def update_dashboard(feature_id, features_state):
+    return create_dashboard(feature_id)
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("# Brand Analyzer")
+    gr.Markdown("*Analyze text using Gemma's interpretable neural features*")
     
-    features_state = gr.State([])
+    features_state = gr.State()
+    selected_feature = gr.State()
     
     with gr.Row():
         with gr.Column(scale=1):
-            input_text = gr.Textbox(
-                lines=5,
-                placeholder="Enter text to analyze...",
-                label="Input Text"
-            )
+            input_text = gr.Textbox(lines=5, label="Input Text")
             analyze_btn = gr.Button("Analyze Features", variant="primary")
-            gr.Examples(
-                examples=["WordLift", "Think Different", "Just Do It"],
-                inputs=input_text
-            )
+            gr.Examples(["WordLift", "Think Different", "Just Do It"], inputs=input_text)
         
         with gr.Column(scale=2):
-            @gr.render(inputs=features_state)
-            def render_features(features):
-                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
-                            )
-            
+            features_container = gr.HTML()
             dashboard = gr.HTML()
-    
+
+    def render_features(features_dict, selected=None):
+        if not features_dict:
+            return ""
+        
+        html = ""
+        for token, features in features_dict.items():
+            html += f'<div class="token-section">'
+            html += f'<div class="token-header">{token}</div>'
+            for feature in features[:3]:
+                selected_class = "selected" if selected == feature['id'] else ""
+                html += f"""
+                <button 
+                    onclick='updateDashboard({feature["id"]})'
+                    class='feature-button {selected_class}'>
+                    Feature {feature["id"]} (Activation: {feature["activation"]:.2f})
+                </button>
+                """
+            html += '</div>'
+        return html
+
     analyze_btn.click(
-        fn=analyze_text,
+        analyze_text,
         inputs=[input_text],
         outputs=[features_state, dashboard]
+    ).then(
+        render_features,
+        inputs=[features_state, selected_feature],
+        outputs=[features_container]
     )
 
 if __name__ == "__main__":