Merge pull request #77 from EleutherAI/cleanup

Add CI tests; remove SAE references

Author: luciaquirke

.github/workflows/tests.yml

@@ -0,0 +1,26 @@
+name: Tests
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+
+    - name: Set up Python
+      uses: actions/setup-python@v3
+      with:
+        python-version: '3.10'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install ".[dev]"
+
+    - name: Run tests
+      run: pytest

delphi/__main__.py

@@ -54,9 +54,9 @@ def load_artifacts(run_cfg: RunConfig):
         token=run_cfg.hf_token,
     )
 
-    hookpoint_to_sae_encode = load_sparse_coders(model, run_cfg, compile=True)
+    hookpoint_to_sparse_encode = load_sparse_coders(model, run_cfg, compile=True)
 
-    return run_cfg.hookpoints, hookpoint_to_sae_encode, model
+    return run_cfg.hookpoints, hookpoint_to_sparse_encode, model
 
 
 async def process_cache(
@@ -205,10 +205,9 @@ def populate_cache(
     run_cfg: RunConfig,
     cfg: CacheConfig,
     model: PreTrainedModel,
-    hookpoint_to_sae_encode: dict[str, Callable],
+    hookpoint_to_sparse_encode: dict[str, Callable],
     latents_path: Path,
     tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
-    filter_bos: bool,
 ):
     """
     Populates an on-disk cache in `latents_path` with SAE latent activations.
@@ -224,7 +223,7 @@ def populate_cache(
     )
     tokens = data["input_ids"]
 
-    if filter_bos:
+    if run_cfg.filter_bos:
         if tokenizer.bos_token_id is None:
             print("Tokenizer does not have a BOS token, skipping BOS filtering")
         else:
@@ -240,7 +239,7 @@ def populate_cache(
 
     cache = LatentCache(
         model,
-        hookpoint_to_sae_encode,
+        hookpoint_to_sparse_encode,
         batch_size=cfg.batch_size,
     )
     cache.run(cfg.n_tokens, tokens)
@@ -276,7 +275,7 @@ async def run(
 
     latent_range = torch.arange(run_cfg.max_latents) if run_cfg.max_latents else None
 
-    hookpoints, hookpoint_to_sae_encode, model = load_artifacts(run_cfg)
+    hookpoints, hookpoint_to_sparse_encode, model = load_artifacts(run_cfg)
     tokenizer = AutoTokenizer.from_pretrained(run_cfg.model, token=run_cfg.hf_token)
 
     if (
@@ -287,15 +286,14 @@ async def run(
             run_cfg,
             cache_cfg,
             model,
-            hookpoint_to_sae_encode,
+            hookpoint_to_sparse_encode,
             latents_path,
             tokenizer,
-            filter_bos=run_cfg.filter_bos,
         )
     else:
         print(f"Files found in {latents_path}, skipping cache population...")
 
-    del model, hookpoint_to_sae_encode
+    del model, hookpoint_to_sparse_encode
 
     if (
         not glob(str(scores_path / ".*")) + glob(str(scores_path / "*"))

delphi/sparse_coders/__init__.py

@@ -28,7 +28,7 @@ def load_sparse_coders(
 
     # Add SAE hooks to the model
     if "gemma" not in run_cfg.sparse_model:
-        hookpoint_to_sae_encode = load_sparsify_sparse_coders(
+        hookpoint_to_sparse_encode = load_sparsify_sparse_coders(
             model,
             run_cfg.sparse_model,
             run_cfg.hookpoints,
@@ -54,7 +54,7 @@ def load_sparse_coders(
             sae_sizes.append(sae_size)
             l0s.append(l0)
 
-        hookpoint_to_sae_encode = load_gemma_autoencoders(
+        hookpoint_to_sparse_encode = load_gemma_autoencoders(
             model_path=model_path,
             ae_layers=layers,
             average_l0s=l0s,
@@ -64,4 +64,4 @@ def load_sparse_coders(
             device=model.device,
         )
 
-    return hookpoint_to_sae_encode
+    return hookpoint_to_sparse_encode

delphi/tests/conftest.py

@@ -62,11 +62,13 @@ def cache_setup(
         sparse_model="EleutherAI/sae-pythia-70m-32k",
         hookpoints=["layers.1"],
     )
-    hookpoint_to_sae_encode = load_sparse_coders(model, run_cfg_gemma)
+    hookpoint_to_sparse_encode = load_sparse_coders(model, run_cfg_gemma)
 
     # Define cache config and initialize cache
     cache_cfg = CacheConfig(batch_size=1, ctx_len=16, n_tokens=100)
-    cache = LatentCache(model, hookpoint_to_sae_encode, batch_size=cache_cfg.batch_size)
+    cache = LatentCache(
+        model, hookpoint_to_sparse_encode, batch_size=cache_cfg.batch_size
+    )
 
     # Generate mock tokens and run the cache
     tokens = mock_dataset

examples/caching_activations.ipynb

@@ -87,7 +87,7 @@
     "    hookpoints=[\"layer_10/width_16k/average_l0_39\"],\n",
     ")\n",
     "\n",
-    "hookpoint_to_sae_encode = load_sparse_coders(model, run_cfg)"
+    "hookpoint_to_sparse_encode = load_sparse_coders(model, run_cfg)"
    ]
   },
   {
@@ -134,7 +134,7 @@
     "\n",
     "cache = LatentCache(\n",
     "    model,\n",
-    "    hookpoint_to_sae_encode,\n",
+    "    hookpoint_to_sparse_encode,\n",
     "    batch_size = cfg.batch_size,\n",
     ")"
    ]

pyproject.toml

@@ -25,6 +25,9 @@ dependencies = [
     "sentence_transformers"
 ]
 
+[project.optional-dependencies]
+dev = ["pytest"]
+
 [tool.pyright]
 include = ["delphi*"]
 reportPrivateImportUsage = false