Update app.py

app.py

@@ -1,60 +1,22 @@
-Hugging Face's logo
-Hugging Face
-Search models, datasets, users...
-Models
-Datasets
-Spaces
-Posts
-Docs
-Enterprise
-Pricing
-
-
-
-Spaces:
-
-WordLift
-/
-brand-llms
-
-
-like
-0
-
-Logs
-App
-Files
-Community
-1
-Settings
-brand-llms
-/
-app.py
-
-cyberandy's picture
-cyberandy
-Update app.py
-3bc7e87
-verified
-21 days ago
-raw
-
-Copy download link
-history
-blame
-edit
-delete
-
-4.55 kB
 import requests
 import gradio as gr
+from enum import Enum
+
+class Model(Enum):
+    GEMMA = "gemma-2-2b"
+    GPT2 = "gpt2-small"
 
-def get_features(text: str):
+MODEL_CONFIGS = {
+    Model.GEMMA: "20-gemmascope-res-16k",
+    Model.GPT2: "9-res-jb"
+}
+
+def get_features(text: str, model: Model):
     url = "https://www.neuronpedia.org/api/search-with-topk"
     payload = {
-        "modelId": "gemma-2-2b", 
+        "modelId": model.value,
         "text": text,
-        "layer": "20-gemmascope-res-16k"
+        "layer": MODEL_CONFIGS[model]
     }
     try:
         response = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
@@ -63,12 +25,15 @@ def get_features(text: str):
     except Exception as e:
         return None
 
-def create_dashboard(feature_id: int) -> str:
+def create_dashboard(feature_id: int, model: Model) -> str:
+    model_path = model.value.lower()
+    layer_name = MODEL_CONFIGS[model].lower()
+    
     return f"""
         <div class="dashboard-container 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"
+                src="https://www.neuronpedia.org/{model_path}/{layer_name}/{feature_id}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
                 width="100%"
                 height="600"
                 frameborder="0"
@@ -77,14 +42,16 @@ def create_dashboard(feature_id: int) -> str:
         </div>
     """
 
-def handle_feature_click(feature_id):
-    return create_dashboard(feature_id)
+def handle_feature_click(feature_id: int, model: Model):
+    return create_dashboard(feature_id, model)
 
-def analyze_text(text: str):
+def analyze_text(text: str, selected_model: str):
+    model = Model.GEMMA if selected_model == "Gemini" else Model.GPT2
+    
     if not text:
         return [], ""
         
-    features_data = get_features(text)
+    features_data = get_features(text, model)
     if not features_data:
         return [], ""
 
@@ -111,7 +78,7 @@ def analyze_text(text: str):
         
         features.append({"token": token, "features": token_features})
     
-    return features, create_dashboard(first_feature_id) if first_feature_id else ""
+    return features, create_dashboard(first_feature_id, model) if first_feature_id else ""
 
 css = """
 @import url('https://fonts.googleapis.com/css2?family=Open+Sans:wght@300;400;600;700&display=swap');
@@ -139,6 +106,11 @@ body { font-family: 'Open Sans', sans-serif !important; }
 .feature-button:hover {
     background-color: #e5e7eb;
 }
+.model-selector {
+    display: flex;
+    gap: 1rem;
+    margin-bottom: 1rem;
+}
 """
 
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
@@ -146,6 +118,11 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     gr.Markdown("*Analyze text using interpretable neural features*", elem_classes="text-gray-600 mb-6")
     
     features_state = gr.State([])
+    selected_model = gr.State("Gemini")  # Default to Gemini
+    
+    with gr.Row(elem_classes="model-selector"):
+        gemini_btn = gr.Button("🧬 Gemini", variant="primary" if selected_model.value == "Gemini" else "secondary")
+        openai_btn = gr.Button("🤖 OpenAI", variant="secondary")
     
     with gr.Row():
         with gr.Column(scale=1):
@@ -161,11 +138,13 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             )
         
         with gr.Column(scale=2):
-            @gr.render(inputs=features_state)
-            def render_features(features):
+            @gr.render(inputs=[features_state, selected_model])
+            def render_features(features, current_model):
                 if not features:
                     return
                 
+                model = Model.GEMMA if current_model == "Gemini" else Model.GPT2
+                
                 for token_group in features:
                     gr.Markdown(f"### {token_group['token']}")
                     with gr.Row():
@@ -175,17 +154,32 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                                 elem_classes=["feature-button"]
                             )
                             btn.click(
-                                fn=lambda fid=feature['id']: handle_feature_click(fid),
+                                fn=lambda fid=feature['id']: handle_feature_click(fid, model),
                                 outputs=dashboard
                             )
             
             dashboard = gr.HTML()
     
+    def update_model(new_model):
+        return new_model
+    
+    gemini_btn.click(
+        fn=lambda: update_model("Gemini"),
+        outputs=selected_model,
+        queue=False
+    )
+    
+    openai_btn.click(
+        fn=lambda: update_model("OpenAI"),
+        outputs=selected_model,
+        queue=False
+    )
+    
     analyze_btn.click(
         fn=analyze_text,
-        inputs=[input_text],
+        inputs=[input_text, selected_model],
         outputs=[features_state, dashboard]
     )
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch(share=False)