Add pre-download sweep support (#154)

Author: alan-cooney

sparse_autoencoder/train/sweep.py

@@ -155,10 +155,11 @@ def setup_source_data(hyperparameters: RuntimeHyperparameters) -> SourceDataset:
 
     if hyperparameters["source_data"]["pre_tokenized"]:
         return PreTokenizedDataset(
-            dataset_path=hyperparameters["source_data"]["dataset_path"],
             context_size=hyperparameters["source_data"]["context_size"],
             dataset_dir=dataset_dir,
             dataset_files=dataset_files,
+            dataset_path=hyperparameters["source_data"]["dataset_path"],
+            pre_download=hyperparameters["source_data"]["pre_download"],
         )
 
     if hyperparameters["source_data"]["tokenizer_name"] is None:
@@ -171,12 +172,13 @@ def setup_source_data(hyperparameters: RuntimeHyperparameters) -> SourceDataset:
     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=dataset_dir,
         dataset_files=dataset_files,
+        dataset_path=hyperparameters["source_data"]["dataset_path"],
         n_processes_preprocessing=4,
+        pre_download=hyperparameters["source_data"]["pre_download"],
+        tokenizer=tokenizer,
     )
 
 

sparse_autoencoder/train/sweep_config.py

@@ -177,9 +177,12 @@ class SourceDataHyperparameters(NestedParameter):
     dataset_dir: Parameter[str] | None = field(default=None)
     """Dataset directory (within the HF dataset)"""
 
-    dataset_files: Parameter[str] | None = field(default=None)
+    dataset_files: Parameter[list[str]] | None = field(default=None)
     """Dataset files (within the HF dataset)."""
 
+    pre_download: Parameter[bool] = field(default=Parameter(value=False))
+    """Whether to pre-download the dataset."""
+
     pre_tokenized: Parameter[bool] = field(default=Parameter(value=True))
     """If the dataset is pre-tokenized."""
 
@@ -209,8 +212,9 @@ class SourceDataRuntimeHyperparameters(TypedDict):
 
     context_size: int
     dataset_dir: str | None
-    dataset_files: str | None
+    dataset_files: list[str] | None
     dataset_path: str
+    pre_download: bool
     pre_tokenized: bool
     tokenizer_name: str | None
 

sparse_autoencoder/train/tests/test_sweep.py

@@ -53,6 +53,7 @@ def dummy_hyperparameters() -> RuntimeHyperparameters:
             "dataset_path": "NeelNanda/c4-code-tokenized-2b",
             "pre_tokenized": True,
             "tokenizer_name": None,
+            "pre_download": False,
         },
         "source_model": {
             "dtype": "float32",

sparse_autoencoder/train/utils/wandb_sweep_types.py

@@ -269,6 +269,7 @@ def __repr__(self) -> str:
     float,
     int,
     str,
+    list[str],
 )