remove sae references

Author: luciaquirke

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,7 +205,7 @@ 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,
 ):
@@ -239,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)
@@ -275,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 (
@@ -286,14 +286,14 @@ async def run(
             run_cfg,
             cache_cfg,
             model,
-            hookpoint_to_sae_encode,
+            hookpoint_to_sparse_encode,
             latents_path,
             tokenizer,
         )
     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/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",
     ")"
    ]