Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import requests
+from concurrent.futures import ThreadPoolExecutor
 from typing import Dict, List, Tuple
 
 BRAND_EXAMPLES = [
@@ -7,52 +8,74 @@ BRAND_EXAMPLES = [
     "Apple - Think Different. Innovation through unique perspectives.",
     "McDonald's - I'm Lovin' It. Creating moments of joy and satisfaction.",
     "L'Oréal - Because You're Worth It. Embracing self-worth and beauty.",
-    "BMW - The Ultimate Driving Machine. Engineering excellence in motion."
+    "BMW - The Ultimate Driving Machine. Engineering excellence in motion.",
+    "Mastercard - There are some things money can't buy. For everything else, there's MasterCard.",
+    "Google - Don't be evil. Building technology for a better world.",
+    "Amazon - Work Hard. Have Fun. Make History. Revolutionizing how the world shops."
 ]
 
-def search_features(text: str) -> List[Dict]:
+def get_feature_activation(text: str, feature_id: int) -> Dict:
     response = requests.post(
-        "https://www.neuronpedia.org/api/search-all",
+        "https://www.neuronpedia.org/api/activation/new",
         headers={"Content-Type": "application/json"},
         json={
-            "modelId": "gemma-2-2b",
-            "sourceSet": "mlp-16k",
-            "text": text,
-            "selectedLayers": ["0-gemmascope-mlp-16k"],
-            "ignoreBos": True,
-            "densityThreshold": -1,
-            "numResults": 50
+            "feature": {
+                "modelId": "gemma-2-2b",
+                "layer": "0-gemmascope-mlp-16k",
+                "index": str(feature_id)
+            },
+            "customText": text
         }
     )
-    return response.json()['result'] if response.status_code == 200 else []
+    result = response.json()
+    values = result.get('values', [0])
+    return {
+        'id': feature_id,
+        'activation': max(values),
+        'token': result['tokens'][result['maxValueTokenIndex']] if values else None
+    }
+
+def analyze_text(text: str, batch_size: int = 50) -> List[Dict]:
+    features = []
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        futures = []
+        for i in range(0, 16384, batch_size):
+            futures.extend([
+                executor.submit(get_feature_activation, text, idx)
+                for idx in range(i, min(i + batch_size, 16384))
+            ])
+        
+        for future in executor.map(lambda f: f.result(), futures):
+            if future['activation'] > 1.0:
+                features.append(future)
+    
+    return sorted(features, key=lambda x: x['activation'], reverse=True)[:10]
 
 def format_output(features: List[Dict]) -> Tuple[str, str, str]:
     if not features:
         return "No significant neural activations detected", "", ""
     
     output = "# Neural Analysis Results\n\n"
-    sorted_features = sorted(features, key=lambda x: x['maxValue'], reverse=True)[:10]
-    
-    for feature in sorted_features:
-        output += f"### Feature {feature['index']}\n"
-        output += f"- Max Activation: {feature['maxValue']:.2f}\n"
-        output += f"- Layer: {feature['layer']}\n\n"
+    for f in features:
+        output += f"### Feature {f['id']}\n"
+        output += f"- Activation: {f['activation']:.2f}\n"
+        output += f"- Peak Token: {f['token']}\n\n"
 
-    top_feature = sorted_features[0]
-    dashboard_url = f"https://www.neuronpedia.org/gemma-2-2b/0-gemmascope-mlp-16k/{top_feature['index']}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
+    dashboard_url = f"https://www.neuronpedia.org/gemma-2-2b/0-gemmascope-mlp-16k/{features[0]['id']}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
     iframe = f'<iframe src="{dashboard_url}" width="100%" height="600px" frameborder="0"></iframe>'
     
-    return output, iframe, f"Feature {top_feature['index']} Dashboard"
+    return output, iframe, f"Feature {features[0]['id']} Dashboard"
 
 def create_interface():
     with gr.Blocks() as interface:
         gr.Markdown("# Brand Message Neural Analyzer")
+        gr.Markdown("Analyze brand messages and taglines using Gemma's neural features")
         
         with gr.Row():
             with gr.Column():
                 input_text = gr.Textbox(
                     lines=5,
-                    placeholder="Enter brand message to analyze...",
+                    placeholder="Enter brand message or tagline to analyze...",
                     label="Brand Message"
                 )
                 analyze_btn = gr.Button("Analyze Neural Patterns")
@@ -64,7 +87,7 @@ def create_interface():
                 dashboard = gr.HTML()
         
         analyze_btn.click(
-            fn=lambda text: format_output(search_features(text)),
+            fn=lambda text: format_output(analyze_text(text)),
             inputs=input_text,
             outputs=[output_text, dashboard, feature_label]
         )