fix: Merge pull request #33 from jbloomAus/fix/topk-selection-purview

Fix/topk selection purview

sae_dashboard/feature_data_generator.py

@@ -7,6 +7,7 @@
 from jaxtyping import Float, Int
 from sae_lens import SAE
 from sae_lens.config import DTYPE_MAP as DTYPES
+from sae_lens.sae import TopK
 from torch import Tensor, nn
 from tqdm.auto import tqdm
 
@@ -87,11 +88,20 @@ def get_feature_data(  # type: ignore
                 self.encoder.device
             )  # make sure acts are on the correct device
 
-            # Compute feature activations from this
-            with FeatureMaskingContext(self.encoder, feature_indices):
-                feature_acts = self.encoder.encode(primary_acts).to(
+            # For TopK, compute all activations first, then select features
+            if isinstance(self.encoder.activation_fn, TopK):
+                # Get all features' activations
+                all_features_acts = self.encoder.encode(primary_acts)
+                # Then select only the features we're interested in
+                feature_acts = all_features_acts[:, :, feature_indices].to(
                     DTYPES[self.cfg.dtype]
                 )
+            else:
+                # For other activation functions, use the masking context
+                with FeatureMaskingContext(self.encoder, feature_indices):
+                    feature_acts = self.encoder.encode(primary_acts).to(
+                        DTYPES[self.cfg.dtype]
+                    )
 
             self.update_rolling_coefficients(
                 model_acts=primary_acts,

sae_dashboard/neuronpedia/neuronpedia_runner.py

@@ -8,15 +8,15 @@
 
 import numpy as np
 import torch
+import wandb
+import wandb.sdk
 from matplotlib import colors
 from sae_lens.sae import SAE
 from sae_lens.toolkit.pretrained_saes import load_sparsity
 from sae_lens.training.activations_store import ActivationsStore
 from tqdm import tqdm
 from transformer_lens import HookedTransformer
 
-import wandb
-import wandb.sdk
 from sae_dashboard.components_config import (
     ActsHistogramConfig,
     Column,

scripts/run_dashboards_runpod.py

@@ -61,7 +61,7 @@ def copy_output_folder(source: str, destination: str):
 
     def find_saelens_release_from_neuronpedia_id(neuronpedia_id: str) -> str:
         for release, item in directory.items():
-            for _, np_id in item.neuronpedia_id.items():
+            for _, np_id in item.neuronpedia_id.items():  # type: ignore
                 if np_id == neuronpedia_id:
                     return release
         raise ValueError(f"Neuronpedia ID {neuronpedia_id} not found")