[WIP] Upgrade to the latest version of sae_vis

.gitignore

@@ -4,6 +4,7 @@
 wandb/
 .data/
 .checkpoints/
+tests/saes_for_tests/
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

pyproject.toml

@@ -20,7 +20,7 @@ dependencies = [
     "zstandard~=0.22.0",
     "matplotlib>=3.5.3",
     "eindex-callum@git+https://github.com/callummcdougall/eindex",
-    "sae_vis@git+https://github.com/callummcdougall/sae_vis.git@b28a0f7c7e936f4bea05528d952dfcd438533cce"
+    "sae_vis@git+https://github.com/callummcdougall/sae_vis"
 ]
 
 [project.urls]
@@ -98,4 +98,4 @@ reportPrivateImportUsage = false
 filterwarnings = [
     # https://github.com/google/python-fire/pull/447
     "ignore::DeprecationWarning:fire:59",
-]
+]

sparsify/scripts/dashboards.yaml

@@ -1,22 +1,28 @@
 pretrained_sae_paths: null # Paths of the pretrained SAEs to load. Should be a path to a .pt file, or a list of them. Can also be provided as a second argument in the command line.
 sae_config_path: null # Path to the config file used to train the SAEs (if null, we'll assume it's at pretrained_sae_paths[0].parent / "config.yaml")
-n_samples: 10_000
-batch_size: 64
-minibatch_size_features: 128 # Num features in each batch of calculations. Lower to avoid OOM errors
+n_samples: 3000
+batch_size: 16
+minibatch_size_features: 100 # Num features in each batch of calculations. Lower to avoid OOM errors
 data: # DatasetConfig for the data which will be used to generate the dashboards
-  dataset_name: 'apollo-research/roneneldan-TinyStories-tokenizer-gpt2'
+  dataset_name: 'apollo-research/Skylion007-openwebtext-tokenizer-gpt2'
   is_tokenized: True
   tokenizer_name: 'gpt2'
-  split: "validation"
-  n_ctx: 512
+  split: "train"
+  n_ctx: 1024
 save_dir: null # The directory for saving the HTML feature dashboard files
+save_json_data: false
 sae_positions: null # The names of the SAE positions to generate dashboards for. e.g.'blocks.2.hook_resid_post'. If None, then all positions will be generated
 feature_indices: null # The features for which to generate dashboards on each SAE. If none, then we'll generate dashbaords for every feature.
 prompt_centric: # Used to generate prompt-centric (rather than feature-centric) dashboards. Feature-centric dashboards will also be generated for every feature appaearing in these
-  n_random_prompt_dashboards: 50 # The number of random prompts to generate prompt-centric dashboards for.
+  n_random_prompt_dashboards: 10 # The number of random prompts to generate prompt-centric dashboards for.
   data: null # "DatasetConfig for getting random prompts. If None, then non-prompt-centric data will be used
   prompts: # Specific prompts on which to generate prompt-centric feature dashboards. A feature-centric dashboard will be generated for every token position in each prompt.
-    - "Sally met Mike at the show. She brought popcorn for him."
-  str_score: "loss_effect" # The ordering metric for which features are most important in prompt-centric dashboards. Can be one of 'act_size', 'act_quantile', or 'loss_effect'
+    - "Sally met Mike at the show. She brought popcorn for him. They ate it together"
+    - 'Lily asked, "Mommy, can I go on the slide?"'
+    - "It was time for the lecture to begin."
+    - "A man was taken to hospital after the crash"
+    - "CAMPAIGN The campaign will focus on three core goals:"
+    - "new_list = [n**2 for n in numbers if n%2==0]"
+  str_score: "act_quantile" # The ordering metric for which features are most important in prompt-centric dashboards. Can be one of 'act_size', 'act_quantile', or 'loss_effect'
   num_top_features: 10 # How many of the most relevant features to show for each prompt in the prompt-centric dashboards
 seed: 0