sae-lens 6.2.0__py3-none-any.whl → 6.3.1__py3-none-any.whl

sae_lens/__init__.py

@@ -1,5 +1,5 @@
 # ruff: noqa: E402
-__version__ = "6.2.0"
+__version__ = "6.3.1"
 
 import logging
 

sae_lens/config.py

@@ -46,6 +46,29 @@ def dict_field(default: dict[str, Any] | None, **kwargs: Any) -> Any:  # type: i
     return simple_parsing.helpers.dict_field(default, type=json_dict, **kwargs)
 
 
+def special_token(s: str) -> Any:
+    """Parse special token value from string."""
+    if s.lower() == "none":
+        return None
+    if s in ["bos", "eos", "sep"]:
+        return s
+    try:
+        return int(s)
+    except ValueError:
+        raise ValueError(
+            f"Expected 'bos', 'eos', 'sep', an integer, or 'none', got {s}"
+        )
+
+
+def special_token_field(
+    default: int | Literal["bos", "eos", "sep"] | None, **kwargs: Any
+) -> Any:  # type: ignore
+    """
+    Helper to wrap simple_parsing.helpers.field so we can load special token fields from the command line.
+    """
+    return simple_parsing.helpers.field(default=default, type=special_token, **kwargs)
+
+
 @dataclass
 class LoggingConfig:
     # WANDB
@@ -116,6 +139,8 @@ class LanguageModelSAERunnerConfig(Generic[T_TRAINING_SAE_CONFIG]):
         training_tokens (int): The number of training tokens.
         store_batch_size_prompts (int): The batch size for storing activations. This controls how many prompts are in the batch of the language model when generating activations.
         seqpos_slice (tuple[int | None, ...]): Determines slicing of activations when constructing batches during training. The slice should be (start_pos, end_pos, optional[step_size]), e.g. for Othello we sometimes use (5, -5). Note, step_size > 0.
+        disable_concat_sequences (bool): Whether to disable concatenating sequences and ignore sequences shorter than the context size. If True, disables concatenating and ignores short sequences.
+        sequence_separator_token (int | Literal["bos", "eos", "sep"] | None): If not `None`, this token will be placed between sentences in a batch to act as a separator. By default, this is the `<bos>` token.
         device (str): The device to use. Usually "cuda".
         act_store_device (str): The device to use for the activation store. "cpu" is advised in order to save VRAM. Defaults to "with_model" which uses the same device as the main model.
         seed (int): The seed to use.
@@ -178,6 +203,10 @@ class LanguageModelSAERunnerConfig(Generic[T_TRAINING_SAE_CONFIG]):
     training_tokens: int = 2_000_000
     store_batch_size_prompts: int = 32
     seqpos_slice: tuple[int | None, ...] = (None,)
+    disable_concat_sequences: bool = False
+    sequence_separator_token: int | Literal["bos", "eos", "sep"] | None = (
+        special_token_field(default="bos")
+    )
 
     # Misc
     device: str = "cpu"
@@ -564,6 +593,9 @@ class PretokenizeRunnerConfig:
     begin_sequence_token: int | Literal["bos", "eos", "sep"] | None = None
     sequence_separator_token: int | Literal["bos", "eos", "sep"] | None = "bos"
 
+    # sequence processing
+    disable_concat_sequences: bool = False
+
     # if saving locally, set save_path
     save_path: str | None = None
 

sae_lens/llm_sae_training_runner.py

@@ -17,6 +17,7 @@ from sae_lens.config import HfDataset, LanguageModelSAERunnerConfig
 from sae_lens.constants import ACTIVATIONS_STORE_STATE_FILENAME, RUNNER_CFG_FILENAME
 from sae_lens.evals import EvalConfig, run_evals
 from sae_lens.load_model import load_model
+from sae_lens.saes.batchtopk_sae import BatchTopKTrainingSAEConfig
 from sae_lens.saes.gated_sae import GatedTrainingSAEConfig
 from sae_lens.saes.jumprelu_sae import JumpReLUTrainingSAEConfig
 from sae_lens.saes.sae import (
@@ -202,6 +203,12 @@ class LanguageModelSAETrainingRunner:
         )
         self.sae.cfg.metadata.prepend_bos = self.cfg.prepend_bos
         self.sae.cfg.metadata.exclude_special_tokens = self.cfg.exclude_special_tokens
+        self.sae.cfg.metadata.sequence_separator_token = (
+            self.cfg.sequence_separator_token
+        )
+        self.sae.cfg.metadata.disable_concat_sequences = (
+            self.cfg.disable_concat_sequences
+        )
 
     def _compile_if_needed(self):
         # Compile model and SAE
@@ -285,7 +292,7 @@ def _parse_cfg_args(
     architecture_parser.add_argument(
         "--architecture",
         type=str,
-        choices=["standard", "gated", "jumprelu", "topk"],
+        choices=["standard", "gated", "jumprelu", "topk", "batchtopk"],
         default="standard",
         help="SAE architecture to use",
     )
@@ -346,6 +353,7 @@ def _parse_cfg_args(
         "gated": GatedTrainingSAEConfig,
         "jumprelu": JumpReLUTrainingSAEConfig,
         "topk": TopKTrainingSAEConfig,
+        "batchtopk": BatchTopKTrainingSAEConfig,
     }
 
     sae_config_type = sae_config_map[architecture]

sae_lens/pretokenize_runner.py

@@ -35,6 +35,7 @@ class PretokenizedDatasetMetadata:
     begin_batch_token: int | Literal["bos", "eos", "sep"] | None
     begin_sequence_token: int | Literal["bos", "eos", "sep"] | None
     sequence_separator_token: int | Literal["bos", "eos", "sep"] | None
+    disable_concat_sequences: bool
 
 
 def metadata_from_config(cfg: PretokenizeRunnerConfig) -> PretokenizedDatasetMetadata:
@@ -52,6 +53,7 @@ def metadata_from_config(cfg: PretokenizeRunnerConfig) -> PretokenizedDatasetMet
         begin_batch_token=cfg.begin_batch_token,
         begin_sequence_token=cfg.begin_sequence_token,
         sequence_separator_token=cfg.sequence_separator_token,
+        disable_concat_sequences=cfg.disable_concat_sequences,
     )
 
 
@@ -99,6 +101,7 @@ def pretokenize_dataset(
                     sequence_separator_token_id=get_special_token_from_cfg(
                         cfg.sequence_separator_token, tokenizer
                     ),
+                    disable_concat_sequences=cfg.disable_concat_sequences,
                 )
             )
         }

sae_lens/tokenization_and_batching.py

@@ -64,6 +64,7 @@ def concat_and_batch_sequences(
     begin_batch_token_id: int | None = None,
     begin_sequence_token_id: int | None = None,
     sequence_separator_token_id: int | None = None,
+    disable_concat_sequences: bool = False,
 ) -> Generator[torch.Tensor, None, None]:
     """
     Generator to concat token sequences together from the tokens_interator, yielding
@@ -75,8 +76,15 @@ def concat_and_batch_sequences(
         begin_batch_token_id: If provided, this token will be at position 0 of each batch
         begin_sequence_token_id: If provided, this token will be the first token of each sequence
         sequence_separator_token_id: If provided, this token will be inserted between concatenated sequences
+        disable_concat_sequences: If True, disable concatenating sequences and ignore sequences shorter than context_size
         max_batches: If not provided, the iterator will be run to completion.
     """
+    if disable_concat_sequences:
+        for tokens in tokens_iterator:
+            if len(tokens) >= context_size:
+                yield tokens[:context_size]
+        return
+
     batch: torch.Tensor | None = None
     for tokens in tokens_iterator:
         if len(tokens.shape) != 1:

sae_lens/training/activations_store.py

@@ -25,6 +25,7 @@ from sae_lens.config import (
     LanguageModelSAERunnerConfig,
 )
 from sae_lens.constants import DTYPE_MAP
+from sae_lens.pretokenize_runner import get_special_token_from_cfg
 from sae_lens.saes.sae import SAE, T_SAE_CONFIG, T_TRAINING_SAE_CONFIG
 from sae_lens.tokenization_and_batching import concat_and_batch_sequences
 from sae_lens.training.mixing_buffer import mixing_buffer
@@ -141,6 +142,8 @@ class ActivationsStore:
             dataset_trust_remote_code=cfg.dataset_trust_remote_code,
             seqpos_slice=cfg.seqpos_slice,
             exclude_special_tokens=exclude_special_tokens,
+            disable_concat_sequences=cfg.disable_concat_sequences,
+            sequence_separator_token=cfg.sequence_separator_token,
         )
 
     @classmethod
@@ -157,6 +160,8 @@ class ActivationsStore:
         train_batch_size_tokens: int = 4096,
         total_tokens: int = 10**9,
         device: str = "cpu",
+        disable_concat_sequences: bool = False,
+        sequence_separator_token: int | Literal["bos", "eos", "sep"] | None = "bos",
     ) -> ActivationsStore:
         if sae.cfg.metadata.hook_name is None:
             raise ValueError("hook_name is required")
@@ -184,6 +189,8 @@ class ActivationsStore:
             dtype=sae.cfg.dtype,
             device=torch.device(device),
             seqpos_slice=sae.cfg.metadata.seqpos_slice or (None,),
+            disable_concat_sequences=disable_concat_sequences,
+            sequence_separator_token=sequence_separator_token,
         )
 
     def __init__(
@@ -209,6 +216,8 @@ class ActivationsStore:
         dataset_trust_remote_code: bool | None = None,
         seqpos_slice: tuple[int | None, ...] = (None,),
         exclude_special_tokens: torch.Tensor | None = None,
+        disable_concat_sequences: bool = False,
+        sequence_separator_token: int | Literal["bos", "eos", "sep"] | None = "bos",
     ):
         self.model = model
         if model_kwargs is None:
@@ -252,6 +261,10 @@ class ActivationsStore:
         self.seqpos_slice = seqpos_slice
         self.training_context_size = len(range(context_size)[slice(*seqpos_slice)])
         self.exclude_special_tokens = exclude_special_tokens
+        self.disable_concat_sequences = disable_concat_sequences
+        self.sequence_separator_token: int | Literal["bos", "eos", "sep"] | None = (
+            sequence_separator_token
+        )
 
         self.n_dataset_processed = 0
 
@@ -361,14 +374,18 @@ class ActivationsStore:
         else:
             tokenizer = getattr(self.model, "tokenizer", None)
             bos_token_id = None if tokenizer is None else tokenizer.bos_token_id
+
             yield from concat_and_batch_sequences(
                 tokens_iterator=self._iterate_raw_dataset_tokens(),
                 context_size=self.context_size,
                 begin_batch_token_id=(bos_token_id if self.prepend_bos else None),
                 begin_sequence_token_id=None,
-                sequence_separator_token_id=(
-                    bos_token_id if self.prepend_bos else None
-                ),
+                sequence_separator_token_id=get_special_token_from_cfg(
+                    self.sequence_separator_token, tokenizer
+                )
+                if tokenizer is not None
+                else None,
+                disable_concat_sequences=self.disable_concat_sequences,
             )
 
     def load_cached_activation_dataset(self) -> Dataset | None:

sae_lens/training/sae_trainer.py

@@ -1,5 +1,4 @@
 import contextlib
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Callable, Generic, Protocol
 
@@ -38,13 +37,6 @@ def _update_sae_lens_training_version(sae: TrainingSAE[Any]) -> None:
     sae.cfg.sae_lens_training_version = str(__version__)
 
 
-@dataclass
-class TrainSAEOutput:
-    sae: TrainingSAE[Any]
-    checkpoint_path: str
-    log_feature_sparsities: torch.Tensor
-
-
 class SaveCheckpointFn(Protocol):
     def __call__(
         self,

{sae_lens-6.2.0.dist-info → sae_lens-6.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.2.0
+Version: 6.3.1
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch

{sae_lens-6.2.0.dist-info → sae_lens-6.3.1.dist-info}/RECORD

@@ -1,17 +1,17 @@
-sae_lens/__init__.py,sha256=ByxdNdLeg_pvK89IX1lHa6iHgs2ab-UulX55Y0hUhY4,3073
+sae_lens/__init__.py,sha256=8vvwKdk-cv0-h2R1ah18VSmNjcBHt7X9gV3A1LtrroM,3073
 sae_lens/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/analysis/hooked_sae_transformer.py,sha256=Eyg1Y2hVIHNuiiLOCTgzstOuW6iA-7hPHqaGR8y_vMs,13809
 sae_lens/analysis/neuronpedia_integration.py,sha256=MrENqc81Mc2SMbxGjbwHzpkGUCAFKSf0i4EdaUF2Oj4,18707
 sae_lens/cache_activations_runner.py,sha256=L5hhuU2-zPQr2S3L64GMKKLeMQfqXxwDl8NbuOtrybI,12567
-sae_lens/config.py,sha256=qMMx9KuiXTD5lG3g0VzaekWOnvdAzGFSq8j1n-GObEQ,26467
+sae_lens/config.py,sha256=6xATsLdg80mXnEsW12x-cvCbAu6SjnONqbRz2eEbqAU,27796
 sae_lens/constants.py,sha256=CSjmiZ-bhjQeVLyRvWxAjBokCgkfM8mnvd7-vxLIWTY,639
 sae_lens/evals.py,sha256=kQyrzczKaVD9rHwfFa_DxL_gMXDxsoIVHmsFIPIU2bY,38696
-sae_lens/llm_sae_training_runner.py,sha256=58XbDylw2fPOD7C-ZfSAjeNqJLXB05uHGTuiYVVbXXY,13354
+sae_lens/llm_sae_training_runner.py,sha256=exxNX_OEhdiUrlgmBP9bjX9DOf0HUcNQGO4unKeDjKM,13713
 sae_lens/load_model.py,sha256=dBB_9gO6kWyQ4sXHq7qB8T3YUlXm3PGwYcpR4UVW4QY,8633
 sae_lens/loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/loading/pretrained_sae_loaders.py,sha256=5XEU4uFFeGCePwqDwhlE7CrFGRSI0U9Cu-UQVa33Y1E,36432
 sae_lens/loading/pretrained_saes_directory.py,sha256=4Vn-Jex6SveD7EbxcSOBv8cx1gkPfUMLU1QOP-ww1ZE,3752
-sae_lens/pretokenize_runner.py,sha256=0nHQq3s_d80VS8iVK4-e6y_orAYVO8c4RrLGtIDfK_E,6885
+sae_lens/pretokenize_runner.py,sha256=w0f6SfZLAxbp5eAAKnet8RqUB_DKofZ9RGsoJwFnYbA,7058
 sae_lens/pretrained_saes.yaml,sha256=nhHW1auhyi4GHYrjUnHQqbNVhI5cMJv-HThzbzU1xG0,574145
 sae_lens/registry.py,sha256=nhy7BPSudSATqW4lo9H_k3Na7sfGHmAf9v-3wpnLL_o,1490
 sae_lens/saes/__init__.py,sha256=RYqE1qkMws-kwQLmBZFhA_VCa69zVtBjGPIy_UAk2pw,1159
@@ -21,18 +21,18 @@ sae_lens/saes/jumprelu_sae.py,sha256=3xkhBcCol2mEpIBLceymCpudocm2ypOjTeTXbpiXoA4
 sae_lens/saes/sae.py,sha256=McpF4pTh70r6SQUbHFm0YQ9X2c2qPULBUSd_YmnEk4Y,38284
 sae_lens/saes/standard_sae.py,sha256=9UqYyYtQuThYxXKNaDjYcyowpOx2-7cShG-TeUP6JCQ,5940
 sae_lens/saes/topk_sae.py,sha256=CXMBI6CFvI5829bOhoQ350VXR9d8uFHUDlULTIWHXoU,8686
-sae_lens/tokenization_and_batching.py,sha256=oUAscjy_LPOrOb8_Ty6eLAcZ0B3HB_wiWjWktgolhG0,4314
+sae_lens/tokenization_and_batching.py,sha256=now7caLbU3p-iGokNwmqZDyIvxYoXgnG1uklhgiLZN4,4656
 sae_lens/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/training/activation_scaler.py,sha256=seEE-2Qd2JMHxqgnsNWPt-DGtYGZxWPnOwCGuVNSOtI,1719
-sae_lens/training/activations_store.py,sha256=HBN3oEib3PlPUDJb_yVFabQp0JcN9rWbnUN1s2DBMAs,31933
+sae_lens/training/activations_store.py,sha256=x2Fwt5QY7M83v6Vf1CSa821j2_WKMw9oPu1cdlLblvg,32887
 sae_lens/training/mixing_buffer.py,sha256=vDpYG5ZE70szDvBsRKcNHEES3h_WTKJ16qDYk5jPOVA,2015
 sae_lens/training/optim.py,sha256=TiI9nbffzXNsI8WjcIsqa2uheW6suxqL_KDDmWXobWI,5312
-sae_lens/training/sae_trainer.py,sha256=2xcO-02OozFunob5vwoHud-hVMhVl9d28_F9gDCiL6o,15529
+sae_lens/training/sae_trainer.py,sha256=6HPf5wtmY1wMUTkLFRg9DujNMMXJkVMPdAhB2svvlkk,15368
 sae_lens/training/types.py,sha256=qSjmGzXf3MLalygG0psnVjmhX_mpLmL47MQtZfe7qxg,81
 sae_lens/training/upload_saes_to_huggingface.py,sha256=r_WzI1zLtGZ5TzAxuG3xa_8T09j3zXJrWd_vzPsPGkQ,4469
 sae_lens/tutorial/tsea.py,sha256=fd1am_XXsf2KMbByDapJo-2qlxduKaa62Z2qcQZ3QKU,18145
 sae_lens/util.py,sha256=mCwLAilGMVo8Scm7CIsCafU7GsfmBvCcjwmloI4Ly7Y,1718
-sae_lens-6.2.0.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-6.2.0.dist-info/METADATA,sha256=Fqsq0scF5Uia0YBmeZQwVi4m4DX16_Ck-cKokbuch7U,5555
-sae_lens-6.2.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-sae_lens-6.2.0.dist-info/RECORD,,
+sae_lens-6.3.1.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.3.1.dist-info/METADATA,sha256=d-dAwcr-WiSFkybEqtOdFxhnJJBX0xiFec8uvln3ztE,5555
+sae_lens-6.3.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sae_lens-6.3.1.dist-info/RECORD,,