Update app.py

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
 import requests
-from typing import Dict, Tuple, List
+from typing import Dict, List, Tuple
+import json
 
 def get_features(text: str) -> Dict:
-    """Get neural features from the API using the exact website parameters."""
+    """Get neural features from the API."""
     url = "https://www.neuronpedia.org/api/search-with-topk"
     payload = {
         "modelId": "gemma-2-2b",
@@ -22,89 +23,127 @@ def get_features(text: str) -> Dict:
     except Exception as e:
         return None
 
-def format_feature_list(token: str, features: List[Dict], show_all: bool = False) -> str:
-    """Format features as HTML list."""
-    feature_count = len(features) if show_all else min(3, len(features))
-    features_html = ""
-    
-    for idx in range(feature_count):
-        feature = features[idx]
-        feature_id = feature['feature_index']
-        activation = feature['activation_value']
-        
-        features_html += f"""
-            <div class="feature-card p-4 rounded-lg mb-4 hover:border-blue-500">
-                <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>
-                    <a href="https://www.neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/{feature_id}" 
-                       target="_blank" 
-                       class="text-blue-600 hover:text-blue-800">
-                        View on Neuronpedia →
-                    </a>
-                </div>
-            </div>
-        """
-        
-        # Add dashboard for first feature only
-        if idx == 0:
-            features_html += f"""
-                <div class="dashboard-container mb-6 p-4">
-                    <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_id}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
-                        width="100%"
-                        height="600"
-                        frameborder="0"
-                        class="rounded-lg"
-                    ></iframe>
-                </div>
-            """
-    
-    remaining = len(features) - 3
-    if not show_all and remaining > 0:
-        features_html += f"""
-            <div class="text-sm text-gray-600 mb-4">
-                {remaining} more features available.
-                <a href="https://www.neuronpedia.org/gemma-2-2b" target="_blank" class="text-blue-600 hover:text-blue-800">
-                    View all on Neuronpedia →
-                </a>
-            </div>
-        """
-    
-    return features_html
-
-def analyze_features(text: str) -> str:
-    """Main analysis function that processes text and returns formatted output."""
-    if not text:
+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 ""
-        
-    data = get_features(text)
-    if not data:
-        return "Error analyzing text"
     
     output = ['<div class="p-6">']
     
     # Process each token's features
-    for result in data['results']:
+    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>
-                {format_feature_list(token, features)}
-            </div>
-        """)
+        output.append(f'<div class="mb-8"><h2 class="text-xl font-bold mb-4">Token: {token}</h2>')
+        
+        # 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
+            
+            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>')
     
     output.append('</div>')
     return "\n".join(output)
 
+def format_dashboard(feature: Dict) -> str:
+    """Format the feature dashboard."""
+    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>
+            <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"
+                width="100%"
+                height="600"
+                frameborder="0"
+                class="rounded-lg"
+            ></iframe>
+        </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]
+                state['selected_feature'] = {
+                    'feature_id': first_feature['feature_index'],
+                    'activation': first_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
+
+def toggle_expansion(token: str, state: Dict) -> Tuple[str, str, Dict]:
+    """Toggle expansion state for a token."""
+    if token in state['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
+
 css = """
 @import url('https://fonts.googleapis.com/css2?family=Open+Sans:wght@300;400;600;700&display=swap');
 
@@ -127,30 +166,26 @@ body {
     border-radius: 8px;
     background-color: #ffffff;
 }
-
-.hljs {
-    background: #f5f7ff !important;
-}
 """
 
 theme = gr.themes.Soft(
     primary_hue=gr.themes.colors.Color(
         name="blue",
-        c50="#eef1ff",
-        c100="#e0e5ff",
-        c200="#c3cbff",
-        c300="#a5b2ff",
-        c400="#8798ff",
-        c500="#6a7eff",
-        c600="#3452db",
-        c700="#2a41af",
-        c800="#1f3183",
-        c900="#152156",
-        c950="#0a102b",
+        c50="#eef1ff", c100="#e0e5ff", c200="#c3cbff",
+        c300="#a5b2ff", c400="#8798ff", c500="#6a7eff",
+        c600="#3452db", c700="#2a41af", c800="#1f3183",
+        c900="#152156", c950="#0a102b",
     )
 )
 
 def create_interface():
+    # Initialize state
+    state = gr.State({
+        'features_data': None,
+        'expanded_tokens': [],
+        'selected_feature': None
+    })
+    
     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")
@@ -169,13 +204,22 @@ def create_interface():
                 )
             
             with gr.Column(scale=2):
-                output = gr.HTML()
+                features_html = gr.HTML()
+                dashboard_html = gr.HTML()
         
+        # Event handlers
         analyze_btn.click(
             fn=analyze_features,
-            inputs=input_text,
-            outputs=output
+            inputs=[input_text, state],
+            outputs=[features_html, dashboard_html, 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