Update app.py

app.py

@@ -9,61 +9,19 @@ def get_features(text: str) -> Dict:
         "text": text,
         "layer": "20-gemmascope-res-16k"
     }
-    
     try:
-        response = requests.post(
-            url, 
-            headers={"Content-Type": "application/json"}, 
-            json=payload
-        )
+        response = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
         response.raise_for_status()
         return response.json()
     except Exception as e:
         return None
 
-def process_features(text: str) -> Tuple[Dict[str, list], str]:
-    """Returns a dictionary of token: [features] and dashboard HTML"""
-    if not text:
-        return {}, ""
-        
-    features_data = get_features(text)
-    if not features_data:
-        return {}, ""
-    
-    token_features = {}
-    for result in features_data['results']:
-        if result['token'] == '<bos>':
-            continue
-            
-        token = result['token']
-        features = []
-        for feature in result['top_features'][:3]:
-            features.append(
-                f"Feature {feature['feature_index']} (Activation: {feature['activation_value']:.2f})"
-            )
-        token_features[token] = features
-
-    # Create initial dashboard
-    if features_data['results']:
-        first_feature = features_data['results'][0]['top_features'][0]
-        dashboard_html = create_dashboard({
-            'feature_id': first_feature['feature_index'],
-            'activation': first_feature['activation_value']
-        })
-    else:
-        dashboard_html = ""
-    
-    return token_features, dashboard_html
-
-def create_dashboard(feature: Dict) -> str:
-    if not feature:
-        return ""
-    
+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['feature_id']} Dashboard</h3>
+            <h3 class="text-lg font-semibold mb-4">Feature {feature_id} Dashboard</h3>
             <iframe
-                src="https://www.neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/{feature['feature_id']}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
+                src="https://www.neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/{feature_id}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
                 width="100%"
                 height="600"
                 frameborder="0"
@@ -72,13 +30,38 @@ def create_dashboard(feature: Dict) -> str:
         </div>
     """
 
+def process_features(text: str) -> Tuple[gr.Column, str]:
+    if not text:
+        return gr.Column(), ""
+        
+    features_data = get_features(text)
+    if not features_data:
+        return gr.Column(), ""
+    
+    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 ""
+
 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;
-}
-
+body { font-family: 'Open Sans', sans-serif !important; }
 .dashboard-container {
     border: 1px solid #e0e5ff;
     border-radius: 8px;
@@ -96,35 +79,32 @@ theme = gr.themes.Soft(
     )
 )
 
-def create_interface():
-    with gr.Blocks(theme=theme, css=css) as interface:
-        gr.Markdown("# Neural Feature 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")
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                input_text = gr.Textbox(
-                    lines=5,
-                    placeholder="Enter text to analyze...",
-                    label="Input Text"
-                )
-                analyze_btn = gr.Button("Analyze Features", variant="primary")
-                gr.Examples(
-                    examples=["WordLift", "Think Different", "Just Do It"],
-                    inputs=input_text
-                )
-            
-            with gr.Column(scale=2):
-                features_text = gr.JSON()
-                dashboard = gr.HTML()
+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")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_text = gr.Textbox(
+                lines=5,
+                placeholder="Enter text to analyze...",
+                label="Input Text"
+            )
+            analyze_btn = gr.Button("Analyze Features", variant="primary")
+            gr.Examples(
+                examples=["WordLift", "Think Different", "Just Do It"],
+                inputs=input_text
+            )
         
-        analyze_btn.click(
-            fn=process_features,
-            inputs=[input_text],
-            outputs=[features_text, dashboard]
-        )
+        with gr.Column(scale=2):
+            features_col = gr.Column()
+            dashboard = gr.HTML()
     
-    return interface
+    analyze_btn.click(
+        fn=process_features,
+        inputs=[input_text],
+        outputs=[features_col, dashboard]
+    )
 
 if __name__ == "__main__":
-    create_interface().launch(share=True)
+    demo.launch(share=True)