generating_sae_dashboards.ipynb is not working

[hijohnnylin]

Two issues with generating_sae_dashboards.ipynb causing it to fail:

1. eindex is missing.
   Potential Fix (my current local workaroud): add back eindex
   Not creating a pull request due to possibility that the correct fix is for eindex to be a dependency in sae_vis, not SAELens.

2. Error running "Use runner" cell with no modifications: ModuleNotFoundError: No module named 'sae_training'

Full error:

{
	"name": "ModuleNotFoundError",
	"message": "No module named 'sae_training'",
	"stack": "---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[2], line 6
      3 FILENAME = f\"final_sparse_autoencoder_gpt2-small_blocks.{layer}.hook_resid_pre_24576.pt\"
      4 path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
----> 6 obj = torch.load(path, map_location=device)
      7 state_dict = obj[\"state_dict\"]
      8 assert set(state_dict.keys()) == {\"W_enc\", \"b_enc\", \"W_dec\", \"b_dec\"}

File ~/Documents/Projects/SAELens/.venv/lib/python3.12/site-packages/torch/serialization.py:1026, in load(f, map_location, pickle_module, weights_only, mmap, **pickle_load_args)
   1024             except RuntimeError as e:
   1025                 raise pickle.UnpicklingError(UNSAFE_MESSAGE + str(e)) from None
-> 1026         return _load(opened_zipfile,
   1027                      map_location,
   1028                      pickle_module,
   1029                      overall_storage=overall_storage,
   1030                      **pickle_load_args)
   1031 if mmap:
   1032     raise RuntimeError(\"mmap can only be used with files saved with \"
   1033                        \"`torch.save(_use_new_zipfile_serialization=True), \"
   1034                        \"please torch.save your checkpoint with this option in order to use mmap.\")

File ~/Documents/Projects/SAELens/.venv/lib/python3.12/site-packages/torch/serialization.py:1438, in _load(zip_file, map_location, pickle_module, pickle_file, overall_storage, **pickle_load_args)
   1436 unpickler = UnpicklerWrapper(data_file, **pickle_load_args)
   1437 unpickler.persistent_load = persistent_load
-> 1438 result = unpickler.load()
   1440 torch._utils._validate_loaded_sparse_tensors()
   1441 torch._C._log_api_usage_metadata(
   1442     \"torch.load.metadata\", {\"serialization_id\": zip_file.serialization_id()}
   1443 )

File ~/Documents/Projects/SAELens/.venv/lib/python3.12/site-packages/torch/serialization.py:1431, in _load.<locals>.UnpicklerWrapper.find_class(self, mod_name, name)
   1429         pass
   1430 mod_name = load_module_mapping.get(mod_name, mod_name)
-> 1431 return super().find_class(mod_name, name)

ModuleNotFoundError: No module named 'sae_training'"
}

[jbloomAus]

ahh this is so annoying. It's the old artefacts we need to load which apparently need to be compatible with the old library. I'll attempt a fix today.

[Istlemin]

Similar issues with logits_lens_with_features.ipynb (Error loading the state dictionary from .pt file: No module named 'sae_training') and evaluating_your_sae.ipynb (No module named 'sae_lens.analysis.visualizer')

[jbloomAus]

Apologies. I'm currently working on resaving the old SAE pickle files in a new format that gets us off the backward compatibility issues. For now, installing eindex should get this to work.

[neverix]

Not creating a pull request due to possibility that the correct fix is for eindex to be a dependency in sae_vis, not SAELens.

It is and it seems to have been patched on GitHub but not pushed yet to PyPI