fix: ensure that load_from_pretrained config typing is correct

pyproject.toml

@@ -36,6 +36,7 @@ automated-interpretability = "^0.0.3"
 python-dotenv = "^1.0.1"
 pyyaml = "^6.0.1"
 pytest-profiling = "^1.7.0"
+typeguard = "^2.13.3"
 
 
 [tool.poetry.group.dev.dependencies]

sae_lens/toolkit/pretrained_sae_loaders.py

@@ -1,9 +1,11 @@
 import json
-from typing import Protocol
+from dataclasses import fields
+from typing import Any, Protocol
 
 import torch
 from huggingface_hub import hf_hub_download
 from safetensors import safe_open
+from typeguard import TypeCheckError, check_type
 
 from sae_lens import __version__
 from sae_lens.training.config import LanguageModelSAERunnerConfig
@@ -44,9 +46,21 @@ def load_pretrained_sae_lens_sae_components(
 ) -> tuple[LanguageModelSAERunnerConfig, dict[str, torch.Tensor]]:
     with open(cfg_path, "r") as f:
         config = json.load(f)
-    var_names = LanguageModelSAERunnerConfig.__init__.__code__.co_varnames
-    # filter config for varnames
-    config = {k: v for k, v in config.items() if k in var_names}
+
+    config_fields = fields(LanguageModelSAERunnerConfig)
+    fields_by_name = {f.name: f for f in config_fields}
+
+    def is_valid_config_field(var_name: str, val: Any) -> bool:
+        if var_name not in fields_by_name:
+            return False
+        try:
+            check_type(val, fields_by_name[var_name].type)
+            return True
+        except TypeCheckError:
+            return False
+
+    # filter out any invalid config fields
+    config = {k: v for k, v in config.items() if is_valid_config_field(k, v)}
     config["verbose"] = False
     config["device"] = device
 

sae_lens/training/sparse_autoencoder.py

@@ -6,7 +6,8 @@
 import json
 import os
 import pickle
-from typing import Callable, NamedTuple, Optional
+from dataclasses import fields
+from typing import Any, Callable, NamedTuple, Optional
 
 import einops
 import torch

tests/unit/training/test_session_loader.py

@@ -1,10 +1,12 @@
 import os
 import tempfile
+from dataclasses import fields
 
 import pytest
 import torch
 from huggingface_hub import hf_hub_download
 from transformer_lens import HookedTransformer
+from typeguard import TypeCheckError, check_type
 
 from sae_lens.training.activations_store import ActivationsStore
 from sae_lens.training.config import LanguageModelSAERunnerConfig
@@ -105,3 +107,13 @@ def test_load_pretrained_sae_from_huggingface():
     assert isinstance(activation_store, ActivationsStore)
     assert sae.cfg.hook_point_layer == layer
     assert sae.cfg.model_name == "gpt2-small"
+
+    for field in fields(LanguageModelSAERunnerConfig):
+        val = getattr(sae.cfg, field.name)
+        try:
+            check_type(val, field.type)
+        except TypeCheckError as e:
+            # reraise to get nicer error message so we know what field is messed up
+            raise ValueError(
+                f"Field {field.name} with value {val} is not of type {field.type} in the config.\n{e}"
+            )