update

app.py

@@ -1,5 +1,6 @@
 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
@@ -7,7 +8,6 @@ from dataclasses import dataclass
 from typing import List, Dict, Optional
 import logging
 
-# Initialize logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -25,7 +25,7 @@ class MarketingFeature:
     threshold: float = 0.1
 
 
-# Define marketing-relevant features from Gemma Scope
+# Define relevant features
 MARKETING_FEATURES = [
     MarketingFeature(
         feature_id=35,
@@ -33,7 +33,7 @@ MARKETING_FEATURES = [
         category="technical",
         description="Detects technical and specialized terminology",
         interpretation_guide="High activation indicates strong technical focus",
-        layer=6,  # Adjusted for Gemma-2B structure
+        layer=20,
     ),
     MarketingFeature(
         feature_id=6680,
@@ -41,7 +41,7 @@ MARKETING_FEATURES = [
         category="technical",
         description="Identifies complex technical concepts",
         interpretation_guide="Consider simplifying language if activation is too high",
-        layer=6,  # Adjusted for Gemma-2B structure
+        layer=20,
     ),
     MarketingFeature(
         feature_id=2,
@@ -49,57 +49,77 @@ MARKETING_FEATURES = [
         category="seo",
         description="Identifies potential SEO keywords",
         interpretation_guide="High activation suggests strong SEO potential",
-        layer=6,  # Adjusted for Gemma-2B structure
+        layer=20,
     ),
 ]
 
 
-class MarketingAnalyzer:
-    """Main class for analyzing marketing content using Gemma Scope"""
+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")
-        # Store model size as instance variable
-        self.model_size = "2b"
+        torch.set_grad_enabled(False)  # Avoid memory issues
         self._initialize_model()
         self._load_saes()
 
     def _initialize_model(self):
-        """Initialize Gemma model and tokenizer"""
         try:
-            model_name = f"google/gemma-{self.model_size}"
-
-            # Initialize model and tokenizer with token from environment
             self.model = AutoModelForCausalLM.from_pretrained(
-                model_name, device_map="auto"
+                "google/gemma-2-2b", device_map="auto"
             )
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-
+            self.tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b")
             self.model.eval()
-            logger.info(f"Initialized model: {model_name}")
-
+            logger.info("Model initialized successfully")
         except Exception as e:
             logger.error(f"Error initializing model: {str(e)}")
             raise
 
     def _load_saes(self):
-        """Load relevant SAEs from Gemma Scope"""
         self.saes = {}
         for feature in MARKETING_FEATURES:
             try:
-                # Load SAE parameters for each feature
                 path = hf_hub_download(
-                    repo_id=f"google/gemma-scope-{self.model_size}-pt-res",
+                    repo_id="google/gemma-scope-2b-pt-res",
                     filename=f"layer_{feature.layer}/width_16k/average_l0_71/params.npz",
+                    force_download=False,
                 )
                 params = np.load(path)
-                self.saes[feature.feature_id] = {
-                    "params": {
-                        k: torch.from_numpy(v).to(self.device)
-                        for k, v in params.items()
-                    },
-                    "feature": feature,
+
+                # 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)
+
+                self.saes[feature.feature_id] = {"sae": sae, "feature": feature}
                 logger.info(f"Loaded SAE for feature {feature.feature_id}")
             except Exception as e:
                 logger.error(
@@ -107,8 +127,23 @@ class MarketingAnalyzer:
                 )
                 continue
 
+    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 analyze_content(self, text: str) -> Dict:
-        """Analyze marketing content using loaded SAEs"""
         results = {
             "text": text,
             "features": {},
@@ -117,26 +152,22 @@ class MarketingAnalyzer:
         }
 
         try:
-            # Get model activations
-            inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
-            with torch.no_grad():
-                outputs = self.model(**inputs, output_hidden_states=True)
-
-            # Analyze each feature
+            # Get activations for each feature
             for feature_id, sae_data in self.saes.items():
                 feature = sae_data["feature"]
-                layer_output = outputs.hidden_states[feature.layer]
+                sae = sae_data["sae"]
 
-                # Apply SAE
-                activations = self._apply_sae(
-                    layer_output, sae_data["params"], feature.threshold
-                )
+                # Get layer activations
+                activations, inputs = self._gather_activations(text, feature.layer)
+
+                # Skip BOS token and get activations
+                sae_acts = sae.encode(activations.to(torch.float32))
+                sae_acts = sae_acts[:, 1:]  # Skip BOS token
 
-                # Skip BOS token and handle empty activations
-                activations = activations[:, 1:]  # Skip BOS token
-                if activations.numel() > 0:
-                    mean_activation = float(activations.mean())
-                    max_activation = float(activations.max())
+                # Calculate metrics
+                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
@@ -154,12 +185,10 @@ class MarketingAnalyzer:
 
                 results["features"][feature_id] = feature_result
 
-                # Aggregate by category
                 if feature.category not in results["categories"]:
                     results["categories"][feature.category] = []
                 results["categories"][feature.category].append(feature_result)
 
-            # Generate recommendations
             results["recommendations"] = self._generate_recommendations(results)
 
         except Exception as e:
@@ -168,22 +197,9 @@ class MarketingAnalyzer:
 
         return results
 
-    def _apply_sae(
-        self,
-        activations: torch.Tensor,
-        sae_params: Dict[str, torch.Tensor],
-        threshold: float,
-    ) -> torch.Tensor:
-        """Apply SAE to get feature activations"""
-        pre_acts = activations @ sae_params["W_enc"] + sae_params["b_enc"]
-        mask = pre_acts > sae_params["threshold"]
-        acts = mask * torch.nn.functional.relu(pre_acts)
-        return acts
-
     def _interpret_activation(
         self, activation: float, feature: MarketingFeature
     ) -> str:
-        """Interpret activation patterns for a feature"""
         if activation > 0.8:
             return f"Very strong presence of {feature.name.lower()}"
         elif activation > 0.5:
@@ -192,19 +208,15 @@ class MarketingAnalyzer:
             return f"Limited presence of {feature.name.lower()}"
 
     def _generate_recommendations(self, results: Dict) -> List[str]:
-        """Generate content recommendations based on analysis"""
         recommendations = []
 
         try:
-            # Get technical features
             tech_features = [
                 f for f in results["features"].values() if f["category"] == "technical"
             ]
 
-            # Calculate average technical score if we have features
             if tech_features:
                 tech_score = np.mean([f["activation_score"] for f in tech_features])
-
                 if tech_score > 0.8:
                     recommendations.append(
                         "Consider simplifying technical language for broader audience"
@@ -220,7 +232,6 @@ class MarketingAnalyzer:
 
 
 def create_gradio_interface():
-    """Create Gradio interface for marketing analysis"""
     try:
         analyzer = MarketingAnalyzer()
     except Exception as e:
@@ -230,30 +241,26 @@ def create_gradio_interface():
             inputs=gr.Textbox(),
             outputs=gr.Textbox(),
             title="Marketing Content Analyzer (Error)",
-            description="Failed to initialize. Please check if HF_TOKEN is properly set.",
+            description="Failed to initialize.",
         )
 
     def analyze(text):
         results = analyzer.analyze_content(text)
 
-        # Format results for display
         output = "Content Analysis Results\n\n"
 
-        # Overall category scores
         output += "Category Scores:\n"
         for category, features in results["categories"].items():
-            if features:  # Check if we have features for this category
+            if features:
                 avg_score = np.mean([f["activation_score"] for f in features])
                 output += f"{category.title()}: {avg_score:.2f}\n"
 
-        # Feature details
         output += "\nFeature Details:\n"
         for feature_id, feature in results["features"].items():
             output += f"\n{feature['name']}:\n"
             output += f"Score: {feature['activation_score']:.2f}\n"
             output += f"Interpretation: {feature['interpretation']}\n"
 
-        # Recommendations
         if results["recommendations"]:
             output += "\nRecommendations:\n"
             for rec in results["recommendations"]:
@@ -261,28 +268,38 @@ def create_gradio_interface():
 
         return output
 
-    # Create interface with custom theming
-    custom_theme = gr.themes.Soft(
-        primary_hue="indigo", secondary_hue="blue", neutral_hue="gray"
-    )
-
-    interface = gr.Interface(
-        fn=analyze,
-        inputs=gr.Textbox(
-            lines=5,
-            placeholder="Enter your marketing content here...",
-            label="Marketing Content",
-        ),
-        outputs=gr.Textbox(label="Analysis Results"),
-        title="Marketing Content Analyzer",
-        description="Analyze your marketing content using Gemma Scope's neural features",
-        examples=[
-            ["WordLift is an AI-powered SEO tool"],
-            ["Our advanced machine learning algorithms optimize your content"],
-            ["Simple and effective website optimization"],
-        ],
-        theme=custom_theme,
-    )
+    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():
+            input_text = gr.Textbox(
+                lines=5,
+                placeholder="Enter your marketing content here...",
+                label="Marketing Content",
+            )
+            output_text = gr.Textbox(label="Analysis Results")
+
+        analyze_btn = gr.Button("Analyze", variant="primary")
+        analyze_btn.click(fn=analyze, inputs=input_text, outputs=output_text)
+
+        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,
+        )
 
     return interface