Update app.py

app.py

@@ -1,25 +1,18 @@
 import gradio as gr
 import requests
-from typing import Dict, Tuple, List
-from operator import itemgetter
-from collections import Counter
+from typing import Dict, Tuple
 import logging
 
-# Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-def get_features(text: str, k: int = 5) -> Dict:
-    """Get neural features from the API with detailed logging."""
+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",
-        "layer": "20-gemmascope-res-16k",  # Updated to match website
-        "sourceSet": "gemma-scope",
         "text": text,
-        "k": k,
-        "maxDensity": 0.01,
-        "ignoreBos": True
+        "layer": "20-gemmascope-res-16k"
     }
     
     try:
@@ -29,106 +22,61 @@ def get_features(text: str, k: int = 5) -> Dict:
             json=payload
         )
         response.raise_for_status()
-        data = response.json()
-        
-        # Log the raw response for analysis
-        logger.info(f"API Response: {data}")
-        
-        # Analyze feature distribution
-        all_features = []
-        feature_counter = Counter()
-        
-        for result in data['results']:
-            token = result['token']
-            logger.info(f"\nToken: {token}")
-            
-            for feature in result['top_features']:
-                feature_id = feature['feature_index']
-                activation = feature['activation_value']
-                logger.info(f"Feature {feature_id}: {activation}")
-                
-                all_features.append({
-                    'token': token,
-                    'feature_id': feature_id,
-                    'activation': activation,
-                    'feature_data': feature.get('feature', {})
-                })
-                feature_counter[feature_id] += 1
-        
-        # Log feature frequency analysis
-        logger.info("\nFeature Frequencies:")
-        for feature_id, count in feature_counter.most_common():
-            logger.info(f"Feature {feature_id}: {count} occurrences")
-        
-        return data, all_features, feature_counter
-        
+        return response.json()
     except Exception as e:
         logger.error(f"Error in API call: {str(e)}")
-        return None, [], Counter()
+        return None
 
 def format_output(text: str) -> Tuple[str, str, str]:
-    data, all_features, feature_counter = get_features(text)
+    data = get_features(text)
     
-    if not data:
+    if not data or 'results' not in data:
         return "Error analyzing text", "", ""
-    
-    # Sort features by frequency first, then by maximum activation within each feature
-    feature_activations = {}
-    for feature in all_features:
-        feature_id = feature['feature_id']
-        activation = feature['activation']
-        if feature_id not in feature_activations or activation > feature_activations[feature_id]['activation']:
-            feature_activations[feature_id] = feature
-    
-    # Get top features by frequency, then sort by activation
-    most_common_features = [
-        feature_activations[feature_id]
-        for feature_id, _ in feature_counter.most_common()
-    ]
-    
-    # Sort by activation within the most common features
-    top_features = sorted(most_common_features, key=itemgetter('activation'), reverse=True)[:5]
-    
-    # Format output
+        
     output = "# Neural Feature Analysis\n\n"
-    output += "## Top 5 Most Active Features\n\n"
     
-    for idx, feat in enumerate(top_features, 1):
-        feature_url = f"https://www.neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/{feat['feature_id']}"
-        
-        feature_info = ""
-        if 'name' in feat['feature_data']:
-            feature_info = f" - {feat['feature_data']['name']}"
-        elif 'description' in feat['feature_data']:
-            feature_info = f" - {feat['feature_data']['description']}"
+    # Process each token's features
+    for result in data['results']:
+        token = result['token']
+        if token == '<bos>':
+            continue
+            
+        output += f"## Token: {token}\n\n"
         
-        output += f"### {idx}. Feature {feat['feature_id']}{feature_info}\n"
-        output += f"- **Token:** '{feat['token']}'\n"
-        output += f"- **Activation:** {feat['activation']:.2f}\n"
-        output += f"- **Frequency:** {feature_counter[feat['feature_id']]} occurrences\n"
-        output += f"- [View on Neuronpedia]({feature_url})\n\n"
+        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}"
+            
+            output += f"- Feature {feature_id} (Activation: {activation:.2f})\n"
+            output += f"  [View on Neuronpedia]({feature_url})\n\n"
     
-    # Use highest activation feature for dashboard
-    if top_features:
-        top_feature = top_features[0]
-        dashboard_url = f"https://www.neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/{top_feature['feature_id']}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
-        iframe = 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>
-            </div>
-        '''
-        feature_label = f"Feature {top_feature['feature_id']} Dashboard (Activation: {top_feature['activation']:.2f})"
-    else:
-        iframe = ""
-        feature_label = "No significant features found"
+    # Use first non-BOS token's top feature for the dashboard
+    dashboard_html = ""
+    feature_label = "No features found"
+    
+    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']
+            
+            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>
+                </div>
+            '''
+            feature_label = f"Feature {feature_id} Dashboard (Activation: {activation:.2f})"
+            break
     
-    return output, iframe, feature_label
+    return output, dashboard_html, feature_label
 
 def create_interface():
     with gr.Blocks() as interface:
@@ -145,8 +93,8 @@ def create_interface():
                 analyze_btn = gr.Button("Analyze Features", variant="primary")
                 gr.Examples([
                     "WordLift",
-                    "Nike - Just Do It. The power of determination.",
-                    "Apple - Think Different. Innovation redefined.",
+                    "Think Different",
+                    "Just Do It"
                 ], inputs=input_text)
             
             with gr.Column():