feat: deprecates eval_hook (#511)

Author: anthonyduong9

sae_lens/config.py

@@ -1,5 +1,6 @@
 import json
 import math
+import warnings
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, cast
@@ -125,7 +126,7 @@ class LanguageModelSAERunnerConfig(Generic[T_TRAINING_SAE_CONFIG]):
         model_name (str): The name of the model to use. This should be the name of the model in the Hugging Face model hub.
         model_class_name (str): The name of the class of the model to use. This should be either `HookedTransformer` or `HookedMamba`.
         hook_name (str): The name of the hook to use. This should be a valid TransformerLens hook.
-        hook_eval (str): NOT CURRENTLY IN USE. The name of the hook to use for evaluation.
+        hook_eval (str): DEPRECATED: Will be removed in v7.0.0. NOT CURRENTLY IN USE. The name of the hook to use for evaluation.
         hook_head_index (int, optional): When the hook is for an activation with a head index, we can specify a specific head to use here.
         dataset_path (str): A Hugging Face dataset path.
         dataset_trust_remote_code (bool): Whether to trust remote code when loading datasets from Huggingface.
@@ -264,6 +265,14 @@ class LanguageModelSAERunnerConfig(Generic[T_TRAINING_SAE_CONFIG]):
     exclude_special_tokens: bool | list[int] = False
 
     def __post_init__(self):
+        if self.hook_eval != "NOT_IN_USE":
+            warnings.warn(
+                "The 'hook_eval' field is deprecated and will be removed in v7.0.0. "
+                "It is not currently used and can be safely removed from your config.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         if self.use_cached_activations and self.cached_activations_path is None:
             self.cached_activations_path = _default_cached_activations_path(
                 self.dataset_path,

tests/helpers.py

@@ -27,7 +27,6 @@ class LanguageModelSAERunnerConfigDict(TypedDict, total=False):
     model_name: str
     model_class_name: str
     hook_name: str
-    hook_eval: str
     hook_head_index: int | None
     dataset_path: str
     dataset_trust_remote_code: bool
@@ -118,7 +117,6 @@ def _get_default_runner_config() -> LanguageModelSAERunnerConfigDict:
         "model_name": TINYSTORIES_MODEL,
         "model_class_name": "HookedTransformer",
         "hook_name": "blocks.0.hook_mlp_out",
-        "hook_eval": "NOT_IN_USE",
         "hook_head_index": None,
         "dataset_path": NEEL_NANDA_C4_10K_DATASET,
         "streaming": False,

tests/training/test_config.py

@@ -34,6 +34,17 @@ def test_sae_training_runner_config_seqpos(
         )
 
 
+def test_LanguageModelSAERunnerConfig_hook_eval_deprecated_usage():
+    with pytest.warns(
+        DeprecationWarning,
+        match="The 'hook_eval' field is deprecated and will be removed in v7.0.0. ",
+    ):
+        LanguageModelSAERunnerConfig(
+            sae=StandardTrainingSAEConfig(d_in=10, d_sae=10),
+            hook_eval="blocks.0.hook_output",
+        )
+
+
 @pytest.mark.parametrize("seqpos_slice, expected_error", test_cases_for_seqpos)
 def test_cache_activations_runner_config_seqpos(
     seqpos_slice: tuple[int, int],

tutorials/mamba_train_example.py

@@ -23,7 +23,6 @@
     model_name="state-spaces/mamba-370m",
     model_class_name="HookedMamba",
     hook_name="blocks.39.hook_ssm_input",
-    hook_eval="blocks.39.hook_ssm_output",  # we compare this when replace hook_point activations with autoencode.decode(autoencoder.encode( hook_point activations))
     dataset_path="NeelNanda/openwebtext-tokenized-9b",
     is_dataset_tokenized=True,
     # SAE Parameters