fix: update mdl sae for new SAELens api

pyproject.toml

@@ -14,7 +14,7 @@ classifiers = ["Topic :: Scientific/Engineering :: Artificial Intelligence"]
 
 [tool.poetry.dependencies]
 python = "^3.10"
-sae_lens = "^6.22.2"
+sae_lens = "^6.28.0"
 transformer-lens = ">=2.0.0"
 torch = ">=2.1.0"
 einops = ">=0.8.0"

sae_bench/evals/mdl/main.py

@@ -39,6 +39,15 @@ class Decodable(Protocol):
     def decode(self, x: torch.Tensor) -> torch.Tensor: ...
 
 
+def _get_filtered_buffer(
+    activations_store: ActivationsStore, n_batches: int
+) -> torch.Tensor:
+    return torch.cat(
+        [activations_store.get_filtered_llm_batch() for _ in range(n_batches)],
+        dim=0,
+    )
+
+
 def build_bins(
     min_pos_activations_F: torch.Tensor,
     max_activations_F: torch.Tensor,
@@ -104,7 +113,7 @@ def calculate_dl(
     float_entropy_F = torch.zeros(num_features, device=device, dtype=torch.float32)
     bool_entropy_F = torch.zeros(num_features, device=device, dtype=torch.float32)
 
-    x_BSN = activations_store.get_filtered_buffer(config.sae_batch_size)
+    x_BSN = _get_filtered_buffer(activations_store, config.sae_batch_size)
     # previous SAELens version had an extra dim in the middle for layer
     x_BSN = x_BSN.unsqueeze(1)
     feature_activations_BsF = sae.encode(x_BSN).squeeze()
@@ -237,7 +246,7 @@ def check_quantised_features_reach_mse_threshold(
     mse_losses: list[torch.Tensor] = []
 
     for i in range(1):
-        x_BSN = activations_store.get_filtered_buffer(config.sae_batch_size)
+        x_BSN = _get_filtered_buffer(activations_store, config.sae_batch_size)
         # previous SAELens version had an extra dim in the middle for layer
         x_BSN = x_BSN.unsqueeze(1)
         feature_activations_BSF = sae.encode(x_BSN).squeeze()
@@ -347,8 +356,8 @@ def get_min_max_activations() -> tuple[torch.Tensor, torch.Tensor]:
         max_activations_1F = torch.zeros(1, num_features, device=device) + 100
 
         for _ in range(10):
-            neuron_activations_BSN = activations_store.get_filtered_buffer(
-                config.sae_batch_size
+            neuron_activations_BSN = _get_filtered_buffer(
+                activations_store, config.sae_batch_size
             ).unsqueeze(1)
 
             feature_activations_BsF = sae.encode(neuron_activations_BSN).squeeze()

tests/unit/evals/mdl/test_main.py

@@ -0,0 +1,26 @@
+from sae_lens import SAE, ActivationsStore
+from transformer_lens import HookedTransformer
+
+from sae_bench.evals.mdl.main import _get_filtered_buffer
+
+
+def test_get_filtered_buffer_concatenates_n_llm_batches(
+    gpt2_model: HookedTransformer, gpt2_l4_sae: SAE
+):
+    store_batch_size_prompts = 2
+    n_batches = 3
+    context_size = 128
+
+    activations_store = ActivationsStore.from_sae(
+        gpt2_model,
+        gpt2_l4_sae,
+        context_size=context_size,
+        store_batch_size_prompts=store_batch_size_prompts,
+        dataset="roneneldan/TinyStories",
+        device="cpu",
+    )
+
+    buffer = _get_filtered_buffer(activations_store, n_batches=n_batches)
+
+    expected_rows = n_batches * store_batch_size_prompts * context_size
+    assert buffer.shape == (expected_rows, gpt2_l4_sae.cfg.d_in)