app.py

import requests
import gradio as gr
from enum import Enum

class Model(Enum):
    GEMMA = "gemma-2-2b"
    GPT2 = "gpt2-small"

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": model.value,
        "text": text,
        "layer": MODEL_CONFIGS[model]
    }
    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_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/{model_path}/{layer_name}/{feature_id}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
                width="100%"
                height="600"
                frameborder="0"
                class="rounded-lg"
            ></iframe>
        </div>
    """

def handle_feature_click(feature_id: int, model: Model):
    return create_dashboard(feature_id, model)

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, model)
    if not features_data:
        return [], ""

    features = []
    first_feature_id = None
    
    for result in features_data['results']:
        if result['token'] == '<bos>':
            continue
            
        token = result['token']
        token_features = []
        
        for feature in result['top_features'][:3]:
            feature_id = feature['feature_index']
            if first_feature_id is None:
                first_feature_id = feature_id
                
            token_features.append({
                "token": token,
                "id": feature_id,
                "activation": feature['activation_value']
            })
        
        features.append({"token": token, "features": token_features})
    
    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');
body { font-family: 'Open Sans', sans-serif !important; }
.dashboard-container {
    border: 1px solid #e0e5ff;
    border-radius: 8px;
    background-color: #ffffff;
}
.token-header {
    font-size: 1.25rem;
    font-weight: 600;
    margin-top: 1rem;
    margin-bottom: 0.5rem;
}
.feature-button {
    display: inline-block;
    margin: 0.25rem;
    padding: 0.5rem 1rem;
    background-color: #f3f4f6;
    border: 1px solid #e5e7eb;
    border-radius: 0.375rem;
    font-size: 0.875rem;
}
.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:
    gr.Markdown("# Brand Analyzer", elem_classes="text-2xl font-bold mb-2")
    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):
            input_text = gr.Textbox(
                lines=5,
                placeholder="Enter text to analyze...",
                label="Input Text"
            )
            analyze_btn = gr.Button("Analyze Features", variant="primary")
            gr.Examples(
                examples=["WordLift", "Think Different", "Just Do It"],
                inputs=input_text
            )
        
        with gr.Column(scale=2):
            @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():
                        for feature in token_group['features']:
                            btn = gr.Button(
                                f"Feature {feature['id']} (Activation: {feature['activation']:.2f})",
                                elem_classes=["feature-button"]
                            )
                            btn.click(
                                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, selected_model],
        outputs=[features_state, dashboard]
    )

if __name__ == "__main__":
    demo.launch(share=False)