app-old.py

import gradio as gr
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import hf_hub_download
import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Optional
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class MarketingFeature:
    """Structure to hold marketing-relevant feature information"""

    feature_id: int
    name: str
    category: str
    description: str
    interpretation_guide: str
    layer: int
    threshold: float = 0.1


# Define relevant features
MARKETING_FEATURES = [
    MarketingFeature(
        feature_id=35,
        name="Technical Term Detector",
        category="technical",
        description="Detects technical and specialized terminology",
        interpretation_guide="High activation indicates strong technical focus",
        layer=20,
    ),
    MarketingFeature(
        feature_id=6680,
        name="Compound Technical Terms",
        category="technical",
        description="Identifies complex technical concepts",
        interpretation_guide="Consider simplifying language if activation is too high",
        layer=20,
    ),
    MarketingFeature(
        feature_id=2,
        name="SEO Keyword Detector",
        category="seo",
        description="Identifies potential SEO keywords",
        interpretation_guide="High activation suggests strong SEO potential",
        layer=20,
    ),
]


class JumpReLUSAE(nn.Module):
    def __init__(self, d_model, d_sae):
        super().__init__()
        self.W_enc = nn.Parameter(torch.zeros(d_model, d_sae))
        self.W_dec = nn.Parameter(torch.zeros(d_sae, d_model))
        self.threshold = nn.Parameter(torch.zeros(d_sae))
        self.b_enc = nn.Parameter(torch.zeros(d_sae))
        self.b_dec = nn.Parameter(torch.zeros(d_model))

    def encode(self, input_acts):
        pre_acts = input_acts @ self.W_enc + self.b_enc
        mask = pre_acts > self.threshold
        acts = mask * torch.nn.functional.relu(pre_acts)
        return acts

    def decode(self, acts):
        return acts @ self.W_dec + self.b_dec

    def forward(self, acts):
        acts = self.encode(acts)
        recon = self.decode(acts)
        return recon


class MarketingAnalyzer:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        torch.set_grad_enabled(False)  # Avoid memory issues
        self._initialize_model()

    def _initialize_model(self):
        try:
            self.model = AutoModelForCausalLM.from_pretrained(
                "google/gemma-2-2b", device_map="auto"
            )
            self.tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b")
            self.model.eval()
            logger.info("Model initialized successfully")
        except Exception as e:
            logger.error(f"Error initializing model: {str(e)}")
            raise

    def _load_sae(self, feature_id: int, layer: int = 20):
        """Dynamically load a single SAE"""
        try:
            path = hf_hub_download(
                repo_id="google/gemma-scope-2b-pt-res",
                filename=f"layer_{layer}/width_16k/average_l0_71/params.npz",
                force_download=False,
            )
            params = np.load(path)

            # Create SAE
            d_model = params["W_enc"].shape[0]
            d_sae = params["W_enc"].shape[1]
            sae = JumpReLUSAE(d_model, d_sae).to(self.device)

            # Load parameters
            sae_params = {
                k: torch.from_numpy(v).to(self.device) for k, v in params.items()
            }
            sae.load_state_dict(sae_params)

            return sae
        except Exception as e:
            logger.error(f"Error loading SAE for feature {feature_id}: {str(e)}")
            return None

    def _gather_activations(self, text: str, layer: int):
        inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
        target_act = None

        def hook(mod, inputs, outputs):
            nonlocal target_act
            target_act = outputs[0]
            return outputs

        handle = self.model.model.layers[layer].register_forward_hook(hook)
        with torch.no_grad():
            _ = self.model(**inputs)
        handle.remove()

        return target_act, inputs

    def _get_feature_activations(self, text: str, sae, layer: int = 20):
        """Get activations for a single feature"""
        activations, _ = self._gather_activations(text, layer)
        sae_acts = sae.encode(activations.to(torch.float32))
        sae_acts = sae_acts[:, 1:]  # Skip BOS token

        if sae_acts.numel() > 0:
            mean_activation = float(sae_acts.mean())
            max_activation = float(sae_acts.max())
        else:
            mean_activation = 0.0
            max_activation = 0.0

        return mean_activation, max_activation

    def analyze_content(self, text: str) -> Dict:
        """Analyze content and find most relevant features"""
        results = {
            "text": text,
            "features": {},
            "categories": {},
            "recommendations": [],
        }

        try:
            # Start with a set of potential features to explore
            feature_pool = list(range(1, 16385))  # Full range of features
            sample_size = 50  # Number of features to sample
            sampled_features = np.random.choice(
                feature_pool, sample_size, replace=False
            )

            # Test each feature
            feature_activations = []
            for feature_id in sampled_features:
                sae = self._load_sae(feature_id)
                if sae is None:
                    continue

                mean_activation, max_activation = self._get_feature_activations(
                    text, sae
                )
                feature_activations.append(
                    {
                        "feature_id": feature_id,
                        "mean_activation": mean_activation,
                        "max_activation": max_activation,
                    }
                )

            # Sort by activation and take top features
            top_features = sorted(
                feature_activations, key=lambda x: x["max_activation"], reverse=True
            )[
                :3
            ]  # Keep top 3 features

            # Analyze top features in detail
            for feature_data in top_features:
                feature_id = feature_data["feature_id"]

                # Get neuronpedia data if available (this would be a placeholder)
                feature_name = f"Feature {feature_id}"
                feature_category = "neural"  # Default category

                feature_result = {
                    "name": feature_name,
                    "category": feature_category,
                    "activation_score": feature_data["mean_activation"],
                    "max_activation": feature_data["max_activation"],
                    "interpretation": self._interpret_activation(
                        feature_data["mean_activation"], feature_id
                    ),
                }

                results["features"][feature_id] = feature_result

                if feature_category not in results["categories"]:
                    results["categories"][feature_category] = []
                results["categories"][feature_category].append(feature_result)

            # Generate recommendations based on activations
            if top_features:
                max_activation = max(f["max_activation"] for f in top_features)
                if max_activation > 0.8:
                    results["recommendations"].append(
                        f"Strong activation detected in feature {top_features[0]['feature_id']}. "
                        "Consider exploring this aspect further."
                    )
                elif max_activation < 0.3:
                    results["recommendations"].append(
                        "Low feature activations overall. Content might benefit from more distinctive elements."
                    )

        except Exception as e:
            logger.error(f"Error analyzing content: {str(e)}")
            raise

        return results

    def _interpret_activation(self, activation: float, feature_id: int) -> str:
        """Interpret activation levels for a feature"""
        if activation > 0.8:
            return f"Very strong activation of feature {feature_id}"
        elif activation > 0.5:
            return f"Moderate activation of feature {feature_id}"
        else:
            return f"Limited activation of feature {feature_id}"


def create_gradio_interface():
    try:
        analyzer = MarketingAnalyzer()
    except Exception as e:
        logger.error(f"Failed to initialize analyzer: {str(e)}")
        return gr.Interface(
            fn=lambda x: "Error: Failed to initialize model. Please check authentication.",
            inputs=gr.Textbox(),
            outputs=gr.Textbox(),
            title="Marketing Content Analyzer (Error)",
            description="Failed to initialize.",
        )

    def analyze(text):
        results = analyzer.analyze_content(text)

        output = "# Content Analysis Results\n\n"

        output += "## Category Scores\n"
        for category, features in results["categories"].items():
            if features:
                avg_score = np.mean([f["activation_score"] for f in features])
                output += f"**{category.title()}**: {avg_score:.2f}\n"

        output += "\n## Feature Details\n"
        for feature_id, feature in results["features"].items():
            output += f"\n### {feature['name']} (Feature {feature_id})\n"
            output += f"**Score**: {feature['activation_score']:.2f}\n\n"
            output += f"**Interpretation**: {feature['interpretation']}\n\n"
            # Add feature explanation from Neuronpedia reference
            output += f"[View feature details on Neuronpedia](https://neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/{feature_id})\n\n"

        if results["recommendations"]:
            output += "\n## Recommendations\n"
            for rec in results["recommendations"]:
                output += f"- {rec}\n"

        feature_id = max(
            results["features"].items(), key=lambda x: x[1]["activation_score"]
        )[0]

        # Build dashboard URL for the highest activating feature
        dashboard_url = f"https://neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/{feature_id}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"

        return output, dashboard_url, feature_id

    with gr.Blocks(
        theme=gr.themes.Default(
            font=[gr.themes.GoogleFont("Open Sans"), "Arial", "sans-serif"],
            primary_hue="indigo",
            secondary_hue="blue",
            neutral_hue="gray",
        )
    ) as interface:
        gr.Markdown("# Marketing Content Analyzer")
        gr.Markdown(
            "Analyze your marketing content using Gemma Scope's neural features"
        )

        with gr.Row():
            with gr.Column(scale=1):
                input_text = gr.Textbox(
                    lines=5,
                    placeholder="Enter your marketing content here...",
                    label="Marketing Content",
                )
                analyze_btn = gr.Button("Analyze", variant="primary")
                gr.Examples(
                    examples=[
                        "WordLift is an AI-powered SEO tool",
                        "Our advanced machine learning algorithms optimize your content",
                        "Simple and effective website optimization",
                    ],
                    inputs=input_text,
                )

            with gr.Column(scale=2):
                output_text = gr.Markdown(label="Analysis Results")
                with gr.Group():
                    gr.Markdown("## Feature Dashboard")
                    feature_id_text = gr.Text(
                        label="Currently viewing feature", show_label=False
                    )
                    dashboard_frame = gr.HTML(
                        value="Analysis results will appear here",
                        label="Feature Dashboard",
                    )

        def update_dashboard(text):
            output, dashboard_url, feature_id = analyze(text)
            return (
                output,
                f"<iframe src='{dashboard_url}' width='100%' height='600px' frameborder='0' style='border: 1px solid #eee; border-radius: 8px;'></iframe>",
                f"Currently viewing Feature {feature_id} - Most active feature in your content",
            )

        analyze_btn.click(
            fn=update_dashboard,
            inputs=input_text,
            outputs=[output_text, dashboard_frame, feature_id_text],
        )

    return interface


if __name__ == "__main__":
    iface = create_gradio_interface()
    iface.launch()