Update app.py

app.py

@@ -1,10 +1,7 @@
 import gradio as gr
 import requests
-from typing import Dict, Tuple
-import logging
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+from typing import Dict, Tuple, List
+import json
 
 def get_features(text: str) -> Dict:
     """Get neural features from the API using the exact website parameters."""
@@ -24,62 +21,154 @@ def get_features(text: str) -> Dict:
         response.raise_for_status()
         return response.json()
     except Exception as e:
-        logger.error(f"Error in API call: {str(e)}")
         return None
 
-def format_output(text: str) -> Tuple[str, str, str]:
-    data = get_features(text)
+def create_feature_html(feature_id: int, activation: float, selected: bool = False) -> str:
+    """Create HTML for an individual feature card."""
+    border_class = "border-blue-500 border-2" if selected else "border border-gray-200"
+    return f"""
+        <div class="feature-card mb-4 {border_class} rounded-lg shadow hover:shadow-md transition-all cursor-pointer p-4"
+             data-feature-id="{feature_id}" onclick="selectFeature(this, {feature_id})">
+            <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>
+    """
+
+def create_token_section(token: str, features: List[Dict], initial_count: int = 3) -> str:
+    """Create HTML for a token section with its features."""
+    features_html = "".join([
+        create_feature_html(f['feature_index'], f['activation_value'])
+        for f in features[:initial_count]
+    ])
     
-    if not data or 'results' not in data:
-        return "Error analyzing text", "", ""
-        
-    output = "# Neural Feature Analysis\n\n"
+    show_more = ""
+    if len(features) > initial_count:
+        remaining = len(features) - initial_count
+        hidden_features = "".join([
+            create_feature_html(f['feature_index'], f['activation_value'])
+            for f in features[initial_count:]
+        ])
+        show_more = f"""
+            <div class="hidden" id="more-features-{token}">{hidden_features}</div>
+            <button class="text-blue-500 hover:text-blue-700 text-sm mt-2"
+                    onclick="toggleFeatures('{token}', this)">
+                Show {remaining} More Features
+            </button>
+        """
     
-    # Process each token's features
-    for result in data['results']:
-        token = result['token']
-        if token == '<bos>':
-            continue
+    return f"""
+        <div class="mb-6">
+            <h2 class="text-xl font-bold mb-4">Token: {token}</h2>
+            <div id="features-{token}">
+                {features_html}
+            </div>
+            {show_more}
+        </div>
+    """
+
+def create_dashboard_html(feature_id: int, activation: float) -> str:
+    """Create HTML for the feature dashboard."""
+    return f"""
+        <div class="border border-gray-200 rounded-lg 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 create_interface_html(data: Dict) -> str:
+    """Create the complete interface HTML."""
+    js_code = """
+        <script>
+        function selectFeature(element, featureId) {
+            // Remove selection from all features
+            document.querySelectorAll('.feature-card').forEach(card => {
+                card.classList.remove('border-blue-500', 'border-2');
+                card.classList.add('border', 'border-gray-200');
+            });
+            
+            // Add selection to clicked feature
+            element.classList.remove('border', 'border-gray-200');
+            element.classList.add('border-blue-500', 'border-2');
             
-        output += f"## Token: {token}\n\n"
+            // Update dashboard
+            document.getElementById('dashboard-container').innerHTML = 
+                `<iframe src="https://www.neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/${featureId}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
+                         width="100%" height="600" frameborder="0" class="rounded-lg"></iframe>`;
+        }
         
-        for feature in result['top_features']:
-            feature_id = feature['feature_index']
-            activation = feature['activation_value']
-            feature_url = f"https://www.neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/{feature_id}"
+        function toggleFeatures(token, button) {
+            const moreFeatures = document.getElementById(`more-features-${token}`);
+            const featuresContainer = document.getElementById(`features-${token}`);
             
-            output += f"- Feature {feature_id} (Activation: {activation:.2f})\n"
-            output += f"  [View on Neuronpedia]({feature_url})\n\n"
+            if (moreFeatures.classList.contains('hidden')) {
+                moreFeatures.classList.remove('hidden');
+                featuresContainer.innerHTML += moreFeatures.innerHTML;
+                button.textContent = 'Show Less';
+            } else {
+                const allFeatures = featuresContainer.querySelectorAll('.feature-card');
+                for (let i = 3; i < allFeatures.length; i++) {
+                    allFeatures[i].remove();
+                }
+                moreFeatures.classList.add('hidden');
+                button.textContent = `Show ${moreFeatures.querySelectorAll('.feature-card').length} More Features`;
+            }
+        }
+        </script>
+    """
     
-    # Use first non-BOS token's top feature for the dashboard
+    tokens_html = ""
     dashboard_html = ""
-    feature_label = "No features found"
+    first_feature = None
     
     for result in data['results']:
-        if result['token'] != '<bos>' and result['top_features']:
-            top_feature = result['top_features'][0]
-            feature_id = top_feature['feature_index']
-            activation = top_feature['activation_value']
+        if result['token'] == '<bos>':
+            continue
             
-            dashboard_url = f"https://www.neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/{feature_id}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
-            dashboard_html = f'''
-                <div style="border:1px solid #eee;border-radius:8px;padding:1px;background:#fff;">
-                    <iframe 
-                        src="{dashboard_url}" 
-                        width="100%" 
-                        height="600px" 
-                        frameborder="0" 
-                        style="border-radius:8px;"
-                    ></iframe>
+        tokens_html += create_token_section(result['token'], result['top_features'])
+        
+        if not first_feature and result['top_features']:
+            first_feature = result['top_features'][0]
+            dashboard_html = create_dashboard_html(
+                first_feature['feature_index'],
+                first_feature['activation_value']
+            )
+    
+    return f"""
+        <div class="p-6">
+            {js_code}
+            <div class="grid grid-cols-1 lg:grid-cols-2 gap-8">
+                <div class="space-y-6">
+                    {tokens_html}
                 </div>
-            '''
-            feature_label = f"Feature {feature_id} Dashboard (Activation: {activation:.2f})"
-            break
+                <div class="lg:sticky lg:top-6">
+                    <div id="dashboard-container">
+                        {dashboard_html}
+                    </div>
+                </div>
+            </div>
+        </div>
+    """
+
+def analyze_features(text: str) -> Tuple[str, str, str]:
+    data = get_features(text)
+    if not data:
+        return "Error analyzing text", "", ""
     
-    return output, dashboard_html, feature_label
+    interface_html = create_interface_html(data)
+    return interface_html, "", ""
 
 def create_interface():
-    with gr.Blocks() as interface:
+    with gr.Blocks(css="") as interface:
         gr.Markdown("# Neural Feature Analyzer")
         gr.Markdown("*Analyze text using Gemma's interpretable neural features*")
         
@@ -98,14 +187,14 @@ def create_interface():
                 ], inputs=input_text)
             
             with gr.Column():
-                output_text = gr.Markdown()
-                feature_label = gr.Text(show_label=False)
-                dashboard = gr.HTML()
+                output_html = gr.HTML()
+                feature_label = gr.Text(show_label=False, visible=False)
+                dashboard = gr.HTML(visible=False)
         
         analyze_btn.click(
-            fn=lambda text: format_output(text),
+            fn=analyze_features,
             inputs=input_text,
-            outputs=[output_text, dashboard, feature_label]
+            outputs=[output_html, feature_label, dashboard]
         )
     
     return interface