update

app.py

@@ -83,7 +83,6 @@ class MarketingAnalyzer:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         torch.set_grad_enabled(False)  # Avoid memory issues
         self._initialize_model()
-        self._load_saes()
 
     def _initialize_model(self):
         try:
@@ -97,35 +96,31 @@ class MarketingAnalyzer:
             logger.error(f"Error initializing model: {str(e)}")
             raise
 
-    def _load_saes(self):
-        self.saes = {}
-        for feature in MARKETING_FEATURES:
-            try:
-                path = hf_hub_download(
-                    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)
+    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)
+            # 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)
+            # 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(
-                    f"Error loading SAE for feature {feature.feature_id}: {str(e)}"
-                )
-                continue
+            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)
@@ -143,7 +138,23 @@ class MarketingAnalyzer:
 
         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": {},
@@ -152,44 +163,74 @@ class MarketingAnalyzer:
         }
 
         try:
-            # Get activations for each feature
-            for feature_id, sae_data in self.saes.items():
-                feature = sae_data["feature"]
-                sae = sae_data["sae"]
-
-                # 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
-
-                # 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
-
-                # Record results
+            # 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": mean_activation,
-                    "max_activation": max_activation,
+                    "name": feature_name,
+                    "category": feature_category,
+                    "activation_score": feature_data["mean_activation"],
+                    "max_activation": feature_data["max_activation"],
                     "interpretation": self._interpret_activation(
-                        mean_activation, feature
+                        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)
-
-            results["recommendations"] = self._generate_recommendations(results)
+                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)}")
@@ -197,38 +238,14 @@ class MarketingAnalyzer:
 
         return results
 
-    def _interpret_activation(
-        self, activation: float, feature: MarketingFeature
-    ) -> str:
+    def _interpret_activation(self, activation: float, feature_id: int) -> str:
+        """Interpret activation levels for a feature"""
         if activation > 0.8:
-            return f"Very strong presence of {feature.name.lower()}"
+            return f"Very strong activation of feature {feature_id}"
         elif activation > 0.5:
-            return f"Moderate presence of {feature.name.lower()}"
+            return f"Moderate activation of feature {feature_id}"
         else:
-            return f"Limited presence of {feature.name.lower()}"
-
-    def _generate_recommendations(self, results: Dict) -> List[str]:
-        recommendations = []
-
-        try:
-            tech_features = [
-                f for f in results["features"].values() if f["category"] == "technical"
-            ]
-
-            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"
-                    )
-                elif tech_score < 0.3:
-                    recommendations.append(
-                        "Could benefit from more specific technical details"
-                    )
-        except Exception as e:
-            logger.error(f"Error generating recommendations: {str(e)}")
-
-        return recommendations
+            return f"Limited activation of feature {feature_id}"
 
 
 def create_gradio_interface():