Update app.py

app.py

@@ -1,8 +1,110 @@
 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."""
+    url = "https://www.neuronpedia.org/api/search-with-topk"
+    payload = {
+        "modelId": "gemma-2-2b",
+        "text": text,
+        "layer": "20-gemmascope-res-16k"
+    }
+    
+    try:
+        response = requests.post(
+            url, 
+            headers={"Content-Type": "application/json"}, 
+            json=payload
+        )
+        response.raise_for_status()
+        return response.json()
+    except Exception as e:
+        return None
+
+def format_feature_list(features: List[Feature], token: str, expanded: bool = False) -> str:
+    """Format features as HTML list."""
+    display_features = features if expanded else features[:3]
+    features_html = ""
+    
+    for feature in display_features:
+        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="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>
+                    </div>
+                </div>
+            </div>
+        """
+    
+    if not expanded and len(features) > 3:
+        remaining = len(features) - 3
+        features_html += f"""
+            <div class="text-center">
+                <span class="text-blue-500 text-sm">{remaining} more features available</span>
+            </div>
+        """
+    
+    return features_html
+
+def format_dashboard(feature: Feature) -> str:
+    """Format the dashboard HTML for a selected feature."""
+    if not feature:
+        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', []):
+        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
 
-# Define custom CSS with Open Sans font and color theme
 css = """
 @import url('https://fonts.googleapis.com/css2?family=Open+Sans:wght@300;400;600;700&display=swap');
 
@@ -10,14 +112,6 @@ body {
     font-family: 'Open Sans', sans-serif !important;
 }
 
-.primary-btn {
-    background-color: #3452db !important;
-}
-
-.primary-btn:hover {
-    background-color: #2a41af !important;
-}
-
 .feature-card {
     border: 1px solid #e0e5ff;
     background-color: #ffffff;
@@ -29,25 +123,6 @@ body {
     box-shadow: 0 2px 4px rgba(52, 82, 219, 0.1);
 }
 
-.feature-card.selected {
-    border: 2px solid #3452db;
-    background-color: #eef1ff;
-}
-
-.show-more-btn {
-    color: #3452db;
-    font-weight: 600;
-}
-
-.show-more-btn:hover {
-    color: #2a41af;
-}
-
-.token-header {
-    color: #152156;
-    font-weight: 700;
-}
-
 .dashboard-container {
     border: 1px solid #e0e5ff;
     border-radius: 8px;
@@ -55,7 +130,6 @@ body {
 }
 """
 
-# Create custom theme
 theme = gr.themes.Soft(
     primary_hue=gr.themes.colors.Color(
         name="blue",
@@ -73,229 +147,96 @@ theme = gr.themes.Soft(
     )
 )
 
-def get_features(text: str) -> Dict:
-    """Get neural features from the API using the exact website parameters."""
-    url = "https://www.neuronpedia.org/api/search-with-topk"
-    payload = {
-        "modelId": "gemma-2-2b",
-        "text": text,
-        "layer": "20-gemmascope-res-16k"
-    }
-    
-    try:
-        response = requests.post(
-            url, 
-            headers={"Content-Type": "application/json"}, 
-            json=payload
-        )
-        response.raise_for_status()
-        return response.json()
-    except Exception as e:
-        return None
-
-def create_feature_html(feature_id: int, activation: float, selected: bool = False) -> str:
-    """Create HTML for an individual feature card."""
-    selected_class = "selected" if selected else ""
-    return f"""
-        <div class="feature-card {selected_class} p-4 rounded-lg mb-4" 
-             data-feature-id="{feature_id}" 
-             onclick="selectFeature(this, {feature_id}, {activation})">
-            <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]
-    ])
-    
-    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 id="toggle-btn-{token}" 
-                    class="show-more-btn text-sm mt-2"
-                    onclick="toggleFeatures('{token}')">
-                Show {remaining} More Features
-            </button>
-        """
-    
-    return f"""
-        <div class="mb-6">
-            <h2 class="token-header text-xl 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="dashboard-container p-4">
-            <h3 class="text-lg font-semibold mb-4 text-gray-900">
-                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 with JavaScript functionality."""
-    js_code = """
-        <script>
-        function updateDashboard(featureId, activation) {
-            const dashboardContainer = document.getElementById('dashboard-container');
-            dashboardContainer.innerHTML = `
-                <div class="dashboard-container p-4">
-                    <h3 class="text-lg font-semibold mb-4 text-gray-900">
-                        Feature ${featureId} Dashboard (Activation: ${activation.toFixed(2)})
-                    </h3>
-                    <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>
-                </div>
-            `;
-        }
-
-        function selectFeature(element, featureId, activation) {
-            // Update selected state visually
-            document.querySelectorAll('.feature-card').forEach(card => {
-                card.classList.remove('selected');
-            });
-            element.classList.add('selected');
-            
-            // Update dashboard
-            updateDashboard(featureId, activation);
-        }
+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
         
-        function toggleFeatures(token) {
-            const moreFeatures = document.getElementById(`more-features-${token}`);
-            const featuresContainer = document.getElementById(`features-${token}`);
-            const toggleButton = document.getElementById(`toggle-btn-${token}`);
-            
-            if (moreFeatures.classList.contains('hidden')) {
-                // Show additional features
-                moreFeatures.classList.remove('hidden');
-                const additionalFeatures = moreFeatures.innerHTML;
-                featuresContainer.insertAdjacentHTML('beforeend', additionalFeatures);
-                toggleButton.textContent = 'Show Less';
-            } else {
-                // Hide additional features
-                const allFeatures = featuresContainer.querySelectorAll('.feature-card');
-                Array.from(allFeatures).slice(3).forEach(card => card.remove());
-                moreFeatures.classList.add('hidden');
-                toggleButton.textContent = `Show ${moreFeatures.children.length} More Features`;
-            }
+    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
         }
-        </script>
-    """
+        # 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]
     
-    tokens_html = ""
-    dashboard_html = ""
-    first_feature = None
+    # 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>")
     
-    for result in data['results']:
-        if result['token'] == '<bos>':
-            continue
-            
-        tokens_html += create_token_section(result['token'], result['top_features'])
+    # Add dashboard if a feature is selected
+    if state['selected_feature']:
+        output.append(format_dashboard(state['selected_feature']))
         
-        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 "\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)
     
-    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>
-                <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", "", ""
+    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
     
-    interface_html = create_interface_html(data)
-    return interface_html, "", ""
+    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(
-            "# Brand Feature Analyzer",
-            elem_classes="text-2xl font-bold text-gray-900 mb-2"
-        )
-        gr.Markdown(
-            "*Analyze your brand using Gemma's interpretable neural features*",
-            elem_classes="text-gray-600 mb-6"
-        )
+        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")
         
         with gr.Row():
-            with gr.Column():
+            with gr.Column(scale=1):
                 input_text = gr.Textbox(
                     lines=5,
                     placeholder="Enter text to analyze...",
-                    label="Input Text",
-                    elem_classes="mb-4"
-                )
-                analyze_btn = gr.Button(
-                    "Analyze Features",
-                    variant="primary",
-                    elem_classes="primary-btn"
+                    label="Input Text"
                 )
-                # Examples without elem_classes
+                analyze_btn = gr.Button("Analyze Features", variant="primary")
                 gr.Examples(
                     examples=["WordLift", "Think Different", "Just Do It"],
                     inputs=input_text
                 )
             
-            with gr.Column():
-                output_html = gr.HTML()
-                feature_label = gr.Text(show_label=False, visible=False)
-                dashboard = gr.HTML(visible=False)
+            with gr.Column(scale=2):
+                output = gr.HTML()
         
+        # Event handlers
         analyze_btn.click(
             fn=analyze_features,
-            inputs=input_text,
-            outputs=[output_html, feature_label, dashboard]
+            inputs=[input_text, state],
+            outputs=[output, state]
         )
-    
+        
     return interface
 
 if __name__ == "__main__":