Support text datasets with sweep (#151)

Author: alan-cooney

sparse_autoencoder/train/sweep.py

@@ -6,6 +6,7 @@
 import torch
 from transformer_lens import HookedTransformer
 from transformer_lens.utils import get_act_name, get_device
+from transformers import AutoTokenizer
 import wandb
 
 from sparse_autoencoder import (
@@ -18,6 +19,8 @@
     PreTokenizedDataset,
     SparseAutoencoder,
 )
+from sparse_autoencoder.source_data.abstract_dataset import SourceDataset
+from sparse_autoencoder.source_data.text_dataset import TextDataset
 from sparse_autoencoder.train.sweep_config import (
     RuntimeHyperparameters,
     SweepConfig,
@@ -126,18 +129,41 @@ def setup_optimizer(
     )
 
 
-def setup_source_data(hyperparameters: RuntimeHyperparameters) -> PreTokenizedDataset:
+def setup_source_data(hyperparameters: RuntimeHyperparameters) -> SourceDataset:
     """Setup the source data for training.
 
     Args:
         hyperparameters: The hyperparameters dictionary.
 
     Returns:
-        PreTokenizedDataset: The initialized source data.
+        The initialized source dataset.
+
+    Raises:
+        ValueError: If the tokenizer name is not specified, but pre_tokenized is False.
     """
-    return PreTokenizedDataset(
+    if hyperparameters["source_data"]["pre_tokenized"]:
+        return PreTokenizedDataset(
+            dataset_path=hyperparameters["source_data"]["dataset_path"],
+            context_size=hyperparameters["source_data"]["context_size"],
+            dataset_dir=hyperparameters["source_data"]["dataset_dir"],
+            dataset_files=hyperparameters["source_data"]["dataset_files"],
+        )
+
+    if hyperparameters["source_data"]["tokenizer_name"] is None:
+        error_message = (
+            "If pre_tokenized is False, then tokenizer_name must be specified in the "
+            "hyperparameters."
+        )
+        raise ValueError(error_message)
+
+    tokenizer = AutoTokenizer.from_pretrained(hyperparameters["source_data"]["tokenizer_name"])
+
+    return TextDataset(
         dataset_path=hyperparameters["source_data"]["dataset_path"],
         context_size=hyperparameters["source_data"]["context_size"],
+        tokenizer=tokenizer,
+        dataset_dir=hyperparameters["source_data"]["dataset_dir"],
+        dataset_files=hyperparameters["source_data"]["dataset_files"],
     )
 
 
@@ -154,7 +180,7 @@ def run_training_pipeline(
     loss: LossReducer,
     optimizer: AdamWithReset,
     activation_resampler: ActivationResampler,
-    source_data: PreTokenizedDataset,
+    source_data: SourceDataset,
     run_name: str,
 ) -> None:
     """Run the training pipeline for the sparse autoencoder.

sparse_autoencoder/train/sweep_config.py

@@ -174,12 +174,45 @@ class SourceDataHyperparameters(NestedParameter):
     context_size: Parameter[int] = field(default=Parameter(DEFAULT_SOURCE_CONTEXT_SIZE))
     """Context size."""
 
+    dataset_dir: Parameter[str | None] = field(default=Parameter(None))
+    """Dataset directory (within the HF dataset)"""
+
+    dataset_files: Parameter[str | None] = field(default=Parameter(None))
+    """Dataset files (within the HF dataset)."""
+
+    pre_tokenized: Parameter[bool] = field(default=Parameter(value=True))
+    """If the dataset is pre-tokenized."""
+
+    tokenizer_name: Parameter[str | None] = field(default=Parameter(None))
+    """Tokenizer name.
+
+    Only set this if the dataset is not pre-tokenized.
+    """
+
+    def __post_init__(self) -> None:
+        """Post initialisation checks.
+
+        Raises:
+            ValueError: If there is an error in the source data hyperparameters.
+        """
+        if self.pre_tokenized.value is False and self.tokenizer_name.value is None:
+            error_message = "The tokenizer name must be specified, when `pre_tokenized` is False."
+            raise ValueError(error_message)
+
+        if self.pre_tokenized.value is True and self.tokenizer_name.value is not None:
+            error_message = "The tokenizer name must not be set, when `pre_tokenized` is True."
+            raise ValueError(error_message)
+
 
 class SourceDataRuntimeHyperparameters(TypedDict):
     """Source data runtime hyperparameters."""
 
-    dataset_path: str
     context_size: int
+    dataset_dir: str | None
+    dataset_files: str | None
+    dataset_path: str
+    pre_tokenized: bool
+    tokenizer_name: str | None
 
 
 @dataclass(frozen=True)

sparse_autoencoder/train/tests/test_sweep.py

@@ -46,7 +46,14 @@ def dummy_hyperparameters() -> RuntimeHyperparameters:
             "validation_number_activations": 1024,
         },
         "random_seed": 49,
-        "source_data": {"context_size": 128, "dataset_path": "NeelNanda/c4-code-tokenized-2b"},
+        "source_data": {
+            "context_size": 128,
+            "dataset_dir": None,
+            "dataset_files": None,
+            "dataset_path": "NeelNanda/c4-code-tokenized-2b",
+            "pre_tokenized": True,
+            "tokenizer_name": None,
+        },
         "source_model": {
             "dtype": "float32",
             "hook_dimension": 512,

sparse_autoencoder/train/utils/wandb_sweep_types.py

@@ -5,7 +5,7 @@
 from abc import ABC
 from dataclasses import asdict, dataclass, is_dataclass
 from enum import Enum, auto
-from typing import Any, Generic, TypeVar, final
+from typing import Any, Generic, TypeAlias, TypeVar, final
 
 from strenum import LowercaseStrEnum
 
@@ -264,7 +264,11 @@ def __repr__(self) -> str:
         return self.__str__()
 
 
-ParamType = TypeVar("ParamType", float, int, str)
+OptionalFloat: TypeAlias = float | None
+OptionalInt: TypeAlias = int | None
+OptionalStr: TypeAlias = str | None
+
+ParamType = TypeVar("ParamType", float, int, str, OptionalFloat, OptionalInt, OptionalStr)
 
 
 @dataclass(frozen=True)