Update app.py

app.py

@@ -1,43 +1,21 @@
 import requests
 import gradio as gr
 
-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, ""
-
+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"
+    }
     try:
-        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 = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
         response.raise_for_status()
-        features = parse_api_response(response.json())
-        first_feature = next(iter(next(iter(features.values()))))
-        dashboard = create_dashboard(first_feature['id'])
-        return features, dashboard
+        return response.json()
     except Exception as e:
-        return None, str(e)
+        return None
 
-def create_dashboard(feature_id: int):
+def create_dashboard(feature_id: int) -> str:
     return f"""
         <div class="dashboard-container p-4">
             <h3 class="text-lg font-semibold mb-4">Feature {feature_id} Dashboard</h3>
@@ -51,77 +29,119 @@ def create_dashboard(feature_id: int):
         </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 = """
-.token-section { margin-bottom: 1.5rem; }
-.token-header { 
+@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 {
     font-size: 1.25rem;
     font-weight: 600;
-    margin-bottom: 0.75rem;
+    margin-top: 1rem;
+    margin-bottom: 0.5rem;
 }
+
 .feature-button {
-    background: #f8fafc;
-    border: 1px solid #e2e8f0;
-    border-radius: 0.5rem;
-    padding: 0.5rem 1rem;
+    display: inline-block;
     margin: 0.25rem;
+    padding: 0.5rem 1rem;
+    background-color: #f3f4f6;
+    border: 1px solid #e5e7eb;
+    border-radius: 0.375rem;
     font-size: 0.875rem;
-    cursor: pointer;
 }
-.feature-button:hover { background: #f1f5f9; }
-.feature-button.selected {
-    background: #e0e7ff;
-    border-color: #6366f1;
+
+.feature-button:hover {
+    background-color: #e5e7eb;
 }
 """
 
-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*")
+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()
-    selected_feature = gr.State()
+    features_state = gr.State([])
     
     with gr.Row():
         with gr.Column(scale=1):
-            input_text = gr.Textbox(lines=5, label="Input Text")
+            input_text = gr.Textbox(
+                lines=5,
+                placeholder="Enter text to analyze...",
+                label="Input Text"
+            )
             analyze_btn = gr.Button("Analyze Features", variant="primary")
-            gr.Examples(["WordLift", "Think Different", "Just Do It"], inputs=input_text)
+            gr.Examples(
+                examples=["WordLift", "Think Different", "Just Do It"],
+                inputs=input_text
+            )
         
         with gr.Column(scale=2):
-            features_container = gr.HTML()
+            @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
+                            )
+            
             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(
-        analyze_text,
+        fn=analyze_text,
         inputs=[input_text],
         outputs=[features_state, dashboard]
-    ).then(
-        render_features,
-        inputs=[features_state, selected_feature],
-        outputs=[features_container]
     )
 
 if __name__ == "__main__":