Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import requests
 from typing import Dict, List, Tuple
-import json
 
 def get_features(text: str) -> Dict:
     """Get neural features from the API."""
@@ -23,77 +22,77 @@ def get_features(text: str) -> Dict:
     except Exception as e:
         return None
 
-def format_features(features_data: Dict, expanded_tokens: List[str], selected_feature: Dict) -> str:
-    """Format features as HTML with expanded state."""
-    if not features_data or 'results' not in features_data:
-        return ""
+def process_features(text: str, state: Dict) -> Tuple[gr.Tabs, gr.HTML, Dict]:
+    """Process features and return UI components."""
+    if not text:
+        return None, None, state
+        
+    features_data = get_features(text)
+    if not features_data:
+        return None, None, state
     
-    output = ['<div class="p-6">']
+    # Update state with new features data
+    state['features_data'] = features_data
+    if 'expanded_tokens' not in state:
+        state['expanded_tokens'] = []
     
-    # Process each token's features
+    # Create tabs for each token
+    tokens_data = []
     for result in features_data['results']:
         if result['token'] == '<bos>':
             continue
-        
+            
         token = result['token']
         features = result['top_features']
-        is_expanded = token in expanded_tokens
-        feature_count = len(features) if is_expanded else min(3, len(features))
-        
-        output.append(f'<div class="mb-8"><h2 class="text-xl font-bold mb-4">Token: {token}</h2>')
+        is_expanded = token in state.get('expanded_tokens', [])
+        num_features = len(features) if is_expanded else min(3, len(features))
         
-        # Display features
-        for idx in range(feature_count):
-            feature = features[idx]
-            feature_id = feature['feature_index']
-            activation = feature['activation_value']
-            is_selected = selected_feature and selected_feature.get('feature_id') == feature_id
+        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"
+            ))
             
-            selected_class = "border-blue-500 border-2" if is_selected else ""
-            
-            output.append(f"""
-                <div class="feature-card p-4 rounded-lg mb-4 hover:border-blue-500 {selected_class}">
-                    <div class="flex justify-between items-center">
-                        <div>
-                            <span class="font-semibold">Feature {feature_id}</span>
-                            <span class="ml-2 text-gray-600">(Activation: {activation:.2f})</span>
-                        </div>
-                    </div>
-                </div>
-            """)
-        
-        # Show more/less button if needed
         if len(features) > 3:
-            action = "less" if is_expanded else f"{len(features) - 3} more"
-            output.append(f"""
-                <div class="text-center mb-4">
-                    <button class="text-blue-600 hover:text-blue-800 text-sm"
-                            onclick="gradio('toggle_expansion', '{token}')">
-                        Show {action} features
-                    </button>
-                </div>
-            """)
-        
-        output.append('</div>')
+            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]
+        first_feature = features_data['results'][0]['top_features'][0]
+        state['selected_feature'] = {
+            '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]
+    )
     
-    output.append('</div>')
-    return "\n".join(output)
+    return token_tabs, dashboard_html, state
 
-def format_dashboard(feature: Dict) -> str:
-    """Format the feature dashboard."""
+def create_dashboard(feature: Dict) -> str:
+    """Create dashboard HTML for a feature."""
     if not feature:
         return ""
     
-    feature_id = feature['feature_id']
-    activation = feature['activation']
-    
     return f"""
         <div class="dashboard-container p-4">
-            <h3 class="text-lg font-semibold mb-4">
-                Feature {feature_id} Dashboard (Activation: {activation:.2f})
-            </h3>
+            <h3 class="text-lg font-semibold mb-4">Feature {feature['feature_id']} Dashboard</h3>
             <iframe
-                src="https://www.neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/{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['feature_id']}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
                 width="100%"
                 height="600"
                 frameborder="0"
@@ -102,47 +101,31 @@ def format_dashboard(feature: Dict) -> str:
         </div>
     """
 
-def analyze_features(text: str, state: Dict) -> Tuple[str, str, Dict]:
-    """Process text and update state."""
-    if not text:
-        return "", "", state
-    
-    features_data = get_features(text)
-    if not features_data:
-        return "Error analyzing text", "", state
-    
-    # Update state
-    state['features_data'] = features_data
-    if not state.get('expanded_tokens'):
-        state['expanded_tokens'] = []
-    
-    # Select first feature by default if none selected
-    if not state.get('selected_feature'):
-        for result in features_data['results']:
-            if result['token'] != '<bos>' and result['top_features']:
-                first_feature = result['top_features'][0]
+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': first_feature['feature_index'],
-                    'activation': first_feature['activation_value']
+                    'feature_id': feature_id,
+                    'activation': feature['activation_value']
                 }
-                break
-    
-    features_html = format_features(features_data, state['expanded_tokens'], state['selected_feature'])
-    dashboard_html = format_dashboard(state['selected_feature'])
-    
-    return features_html, dashboard_html, state
+                return create_dashboard(state['selected_feature']), state
+                
+    return None, state
 
-def toggle_expansion(token: str, state: Dict) -> Tuple[str, str, Dict]:
-    """Toggle expansion state for a token."""
-    if token in state['expanded_tokens']:
+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)
-    
-    features_html = format_features(state['features_data'], state['expanded_tokens'], state['selected_feature'])
-    dashboard_html = format_dashboard(state['selected_feature'])
-    
-    return features_html, dashboard_html, state
+        
+    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');
@@ -151,16 +134,6 @@ body {
     font-family: 'Open Sans', sans-serif !important;
 }
 
-.feature-card {
-    border: 1px solid #e0e5ff;
-    background-color: #ffffff;
-    transition: all 0.2s ease;
-}
-
-.feature-card:hover {
-    box-shadow: 0 2px 4px rgba(52, 82, 219, 0.1);
-}
-
 .dashboard-container {
     border: 1px solid #e0e5ff;
     border-radius: 8px;
@@ -179,12 +152,7 @@ theme = gr.themes.Soft(
 )
 
 def create_interface():
-    # Initialize state
-    state = gr.State({
-        'features_data': None,
-        'expanded_tokens': [],
-        'selected_feature': None
-    })
+    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")
@@ -204,22 +172,14 @@ def create_interface():
                 )
             
             with gr.Column(scale=2):
-                features_html = gr.HTML()
-                dashboard_html = gr.HTML()
+                token_tabs = gr.Tabs()
+                dashboard = gr.HTML()
         
-        # Event handlers
         analyze_btn.click(
-            fn=analyze_features,
+            fn=process_features,
             inputs=[input_text, state],
-            outputs=[features_html, dashboard_html, state]
+            outputs=[token_tabs, dashboard, state]
         )
-        
-        # Custom JavaScript function for token expansion
-        interface.load(None, None, None, _js="""
-            function toggle_expansion(token) {
-                // Function will be called from HTML onclick
-            }
-        """)
     
     return interface