Merge pull request #57 from chanind/fix-slash-in-model-name-autointerp

fix: gracefully handle slashes in model filename for autointerp

Author: adamkarvonen

sae_bench/evals/autointerp/main.py

@@ -55,6 +55,10 @@ def str_bool(b: bool) -> str:
     return "Y" if b else ""
 
 
+def escape_slash(s: str) -> str:
+    return s.replace("/", "_")
+
+
 class Example:
     """
     Data for a single example sequence.
@@ -523,7 +527,7 @@ def run_eval_single_sae(
 
     os.makedirs(artifacts_folder, exist_ok=True)
 
-    tokens_filename = f"{config.model_name}_{config.total_tokens}_tokens_{config.llm_context_size}_ctx.pt"
+    tokens_filename = f"{escape_slash(config.model_name)}_{config.total_tokens}_tokens_{config.llm_context_size}_ctx.pt"
     tokens_path = os.path.join(artifacts_folder, tokens_filename)
 
     if os.path.exists(tokens_path):

tests/conftest.py

@@ -1,4 +1,5 @@
 import pytest
+import torch
 from sae_lens import SAE
 from transformer_lens import HookedTransformer
 from transformers import GPT2LMHeadModel
@@ -21,6 +22,15 @@ def gpt2_l4_sae() -> SAE:
     )[0]
 
 
+@pytest.fixture
+def gpt2_l4_sae_sparsity() -> torch.Tensor:
+    sparsity = SAE.from_pretrained(
+        "gpt2-small-res-jb", "blocks.4.hook_resid_pre", device="cpu"
+    )[2]
+    assert sparsity is not None
+    return sparsity
+
+
 @pytest.fixture
 def gpt2_l5_sae() -> SAE:
     return SAE.from_pretrained(

tests/unit/evals/autointerp/test_main.py

@@ -0,0 +1,69 @@
+from pathlib import Path
+from unittest.mock import patch
+
+import torch
+from sae_lens import SAE
+from transformer_lens import HookedTransformer
+
+from sae_bench.evals.autointerp.eval_config import AutoInterpEvalConfig
+from sae_bench.evals.autointerp.main import run_eval_single_sae
+
+
+def test_run_eval_single_sae_saves_tokens_to_artifacts_folder(
+    gpt2_l4_sae: SAE,
+    gpt2_model: HookedTransformer,
+    gpt2_l4_sae_sparsity: torch.Tensor,
+    tmp_path: Path,
+):
+    artifacts_folder = tmp_path / "artifacts"
+
+    config = AutoInterpEvalConfig(
+        model_name="gpt2",
+        dataset_name="roneneldan/TinyStories",
+        n_latents=2,
+        total_tokens=255,
+        llm_context_size=128,
+    )
+    with patch("sae_bench.evals.autointerp.main.AutoInterp.run", return_value={}):
+        run_eval_single_sae(
+            config=config,
+            sae=gpt2_l4_sae,
+            model=gpt2_model,
+            device="cpu",
+            artifacts_folder=str(artifacts_folder),
+            api_key="fake_api_key",
+            sae_sparsity=gpt2_l4_sae_sparsity,
+        )
+
+    assert (artifacts_folder / "gpt2_255_tokens_128_ctx.pt").exists()
+    tokenized_dataset = torch.load(artifacts_folder / "gpt2_255_tokens_128_ctx.pt")
+    assert tokenized_dataset.shape == (2, 128)
+
+
+def test_run_eval_single_sae_saves_handles_slash_in_model_name(
+    gpt2_l4_sae: SAE,
+    gpt2_model: HookedTransformer,
+    gpt2_l4_sae_sparsity: torch.Tensor,
+    tmp_path: Path,
+):
+    artifacts_folder = tmp_path / "artifacts"
+
+    config = AutoInterpEvalConfig(
+        model_name="openai/gpt2",
+        dataset_name="roneneldan/TinyStories",
+        n_latents=2,
+        total_tokens=255,
+        llm_context_size=128,
+    )
+    with patch("sae_bench.evals.autointerp.main.AutoInterp.run", return_value={}):
+        run_eval_single_sae(
+            config=config,
+            sae=gpt2_l4_sae,
+            model=gpt2_model,
+            device="cpu",
+            artifacts_folder=str(artifacts_folder),
+            api_key="fake_api_key",
+            sae_sparsity=gpt2_l4_sae_sparsity,
+        )
+
+    assert (artifacts_folder / "openai_gpt2_255_tokens_128_ctx.pt").exists()