Update app.py

app.py

@@ -3,7 +3,6 @@ import requests
 from typing import Dict, List, Tuple
 
 def get_features(text: str) -> Dict:
-    """Get neural features from the API."""
     url = "https://www.neuronpedia.org/api/search-with-topk"
     payload = {
         "modelId": "gemma-2-2b",
@@ -22,69 +21,44 @@ def get_features(text: str) -> Dict:
     except Exception as e:
         return None
 
-def process_features(text: str, state: Dict) -> Tuple[gr.Tabs, gr.HTML, Dict]:
-    """Process features and return UI components."""
+def process_features(text: str) -> Tuple[List[gr.Tab], str]:
     if not text:
-        return None, None, state
+        return [], ""
         
     features_data = get_features(text)
     if not features_data:
-        return None, None, state
+        return [], ""
     
-    # Update state with new features data
-    state['features_data'] = features_data
-    if 'expanded_tokens' not in state:
-        state['expanded_tokens'] = []
-    
-    # Create tabs for each token
-    tokens_data = []
+    tabs = []
     for result in features_data['results']:
         if result['token'] == '<bos>':
             continue
             
         token = result['token']
         features = result['top_features']
-        is_expanded = token in state.get('expanded_tokens', [])
-        num_features = len(features) if is_expanded else min(3, len(features))
         
-        feature_list = []
-        for feature in features[:num_features]:
-            feature_list.append(gr.Button(
-                f"Feature {feature['feature_index']} (Activation: {feature['activation_value']:.2f})",
-                variant="secondary"
-            ))
-            
-        if len(features) > 3:
-            expand_btn = gr.Button(
-                f"Show {'less' if is_expanded else 'more'} features",
-                variant="secondary",
-                size="sm"
-            )
-            feature_list.append(expand_btn)
-            
-        tokens_data.append((token, feature_list))
-    
-    # Select first feature for dashboard
-    if tokens_data and 'selected_feature' not in state:
-        first_token = tokens_data[0][0]
+        with gr.Tab(token):
+            feature_list = []
+            for feature in features[:3]:
+                gr.Button(
+                    f"Feature {feature['feature_index']} (Activation: {feature['activation_value']:.2f})",
+                    variant="secondary"
+                )
+        tabs.append(gr.Tab)
+
+    # Create initial dashboard for first feature
+    if features_data['results']:
         first_feature = features_data['results'][0]['top_features'][0]
-        state['selected_feature'] = {
+        dashboard_html = create_dashboard({
             'feature_id': first_feature['feature_index'],
             'activation': first_feature['activation_value']
-        }
-        dashboard_html = create_dashboard(state['selected_feature'])
+        })
     else:
-        dashboard_html = None
-    
-    # Create tabs component
-    token_tabs = gr.Tabs(
-        [gr.Tab(token, feature_list) for token, feature_list in tokens_data]
-    )
+        dashboard_html = ""
     
-    return token_tabs, dashboard_html, state
+    return tabs, dashboard_html
 
 def create_dashboard(feature: Dict) -> str:
-    """Create dashboard HTML for a feature."""
     if not feature:
         return ""
     
@@ -101,32 +75,6 @@ def create_dashboard(feature: Dict) -> str:
         </div>
     """
 
-def select_feature(feature_id: str, state: Dict) -> Tuple[str, Dict]:
-    """Handle feature selection."""
-    if not state.get('features_data'):
-        return None, state
-        
-    for result in state['features_data']['results']:
-        for feature in result['top_features']:
-            if feature['feature_index'] == feature_id:
-                state['selected_feature'] = {
-                    'feature_id': feature_id,
-                    'activation': feature['activation_value']
-                }
-                return create_dashboard(state['selected_feature']), state
-                
-    return None, state
-
-def toggle_token_expansion(token: str, state: Dict) -> Tuple[gr.Tabs, Dict]:
-    """Toggle token expansion state."""
-    if token in state.get('expanded_tokens', []):
-        state['expanded_tokens'].remove(token)
-    else:
-        state['expanded_tokens'].append(token)
-        
-    tabs, _, state = process_features(state['current_text'], state)
-    return tabs, state
-
 css = """
 @import url('https://fonts.googleapis.com/css2?family=Open+Sans:wght@300;400;600;700&display=swap');
 
@@ -152,8 +100,6 @@ theme = gr.themes.Soft(
 )
 
 def create_interface():
-    state = gr.State({})
-    
     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")
@@ -177,11 +123,11 @@ def create_interface():
         
         analyze_btn.click(
             fn=process_features,
-            inputs=[input_text, state],
-            outputs=[token_tabs, dashboard, state]
+            inputs=[input_text],
+            outputs=[token_tabs, dashboard]
         )
     
     return interface
 
 if __name__ == "__main__":
-    create_interface().launch(share=True)
+    create_interface().launch()