Update app.py

app.py

@@ -1,22 +1,6 @@
 import gradio as gr
 import requests
 from typing import Dict, Tuple, List
-import json
-from dataclasses import dataclass
-from typing import Optional
-
-@dataclass
-class Feature:
-    feature_id: int
-    activation: float
-    token: str
-    position: int
-
-class FeatureState:
-    def __init__(self):
-        self.features_by_token = {}
-        self.expanded_tokens = set()
-        self.selected_feature = None
 
 def get_features(text: str) -> Dict:
     """Get neural features from the API using the exact website parameters."""
@@ -38,72 +22,88 @@ def get_features(text: str) -> Dict:
     except Exception as e:
         return None
 
-def format_feature_list(features: List[Feature], token: str, expanded: bool = False) -> str:
+def format_feature_list(token: str, features: List[Dict], show_all: bool = False) -> str:
     """Format features as HTML list."""
-    display_features = features if expanded else features[:3]
+    feature_count = len(features) if show_all else min(3, len(features))
     features_html = ""
     
-    for feature in display_features:
+    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 cursor-pointer hover:border-blue-500" 
-                 data-feature-id="{feature.feature_id}">
+            <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.feature_id}</span>
-                        <span class="ml-2 text-gray-600">(Activation: {feature.activation:.2f})</span>
+                        <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>
+            """
     
-    if not expanded and len(features) > 3:
-        remaining = len(features) - 3
+    remaining = len(features) - 3
+    if not show_all and remaining > 0:
         features_html += f"""
-            <div class="text-center">
-                <span class="text-blue-500 text-sm">{remaining} more features available</span>
+            <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 format_dashboard(feature: Feature) -> str:
-    """Format the dashboard HTML for a selected feature."""
-    if not feature:
+def analyze_features(text: str) -> str:
+    """Main analysis function that processes text and returns formatted output."""
+    if not text:
         return ""
         
-    return f"""
-        <div class="dashboard-container p-4">
-            <h3 class="text-lg font-semibold mb-4 text-gray-900">
-                Feature {feature.feature_id} Dashboard (Activation: {feature.activation:.2f})
-            </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"
-                width="100%"
-                height="600"
-                frameborder="0"
-                class="rounded-lg"
-            ></iframe>
-        </div>
-    """
-
-def process_features(data: Dict) -> Dict[str, List[Feature]]:
-    """Process API response into features grouped by token."""
-    features_by_token = {}
-    for result in data.get('results', []):
+    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']:
         if result['token'] == '<bos>':
             continue
             
         token = result['token']
-        features = []
-        for idx, feature in enumerate(result.get('top_features', [])):
-            features.append(Feature(
-                feature_id=feature['feature_index'],
-                activation=feature['activation_value'],
-                token=token,
-                position=idx
-            ))
-        features_by_token[token] = features
-    return features_by_token
+        features = result['top_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('</div>')
+    return "\n".join(output)
 
 css = """
 @import url('https://fonts.googleapis.com/css2?family=Open+Sans:wght@300;400;600;700&display=swap');
@@ -119,7 +119,6 @@ body {
 }
 
 .feature-card:hover {
-    border-color: #3452db;
     box-shadow: 0 2px 4px rgba(52, 82, 219, 0.1);
 }
 
@@ -128,6 +127,10 @@ body {
     border-radius: 8px;
     background-color: #ffffff;
 }
+
+.hljs {
+    background: #f5f7ff !important;
+}
 """
 
 theme = gr.themes.Soft(
@@ -147,69 +150,7 @@ theme = gr.themes.Soft(
     )
 )
 
-def analyze_features(text: str, state: Optional[Dict] = None) -> Tuple[str, Dict]:
-    """Main analysis function that processes text and returns formatted output."""
-    if not text:
-        return "", None
-        
-    data = get_features(text)
-    if not data:
-        return "Error analyzing text", None
-        
-    # Process features and build state
-    features_by_token = process_features(data)
-    
-    # Initialize state if needed
-    if not state:
-        state = {
-            'features_by_token': features_by_token,
-            'expanded_tokens': set(),
-            'selected_feature': None
-        }
-        # Select first feature as default
-        first_token = next(iter(features_by_token))
-        if features_by_token[first_token]:
-            state['selected_feature'] = features_by_token[first_token][0]
-    
-    # Build output HTML
-    output = []
-    for token, features in features_by_token.items():
-        expanded = token in state['expanded_tokens']
-        token_html = f"<h2 class='text-xl font-bold mb-4'>Token: {token}</h2>"
-        features_html = format_feature_list(features, token, expanded)
-        
-        output.append(f"<div class='mb-6'>{token_html}{features_html}</div>")
-    
-    # Add dashboard if a feature is selected
-    if state['selected_feature']:
-        output.append(format_dashboard(state['selected_feature']))
-        
-    return "\n".join(output), state
-
-def toggle_expansion(token: str, state: Dict) -> Tuple[str, Dict]:
-    """Toggle expansion state for a token's features."""
-    if token in state['expanded_tokens']:
-        state['expanded_tokens'].remove(token)
-    else:
-        state['expanded_tokens'].add(token)
-    
-    output_html, state = analyze_features(None, state)
-    return output_html, state
-
-def select_feature(feature_id: int, state: Dict) -> Tuple[str, Dict]:
-    """Select a feature and update the dashboard."""
-    for features in state['features_by_token'].values():
-        for feature in features:
-            if feature.feature_id == feature_id:
-                state['selected_feature'] = feature
-                break
-    
-    output_html, state = analyze_features(None, state)
-    return output_html, state
-
 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")
@@ -230,14 +171,13 @@ def create_interface():
             with gr.Column(scale=2):
                 output = gr.HTML()
         
-        # Event handlers
         analyze_btn.click(
             fn=analyze_features,
-            inputs=[input_text, state],
-            outputs=[output, state]
+            inputs=input_text,
+            outputs=output
         )
-        
+    
     return interface
 
 if __name__ == "__main__":
-    create_interface().launch()
+    create_interface().launch(share=True)