Fix dataset pickling

fast_llm/data/dataset/config.py

@@ -20,6 +20,7 @@ class SamplingConfig:
     num_samples: int
     seed: int
     cache_directory: pathlib.Path | None
+    # TODO: This prevents the sampling config from being pickled in multiprocessing.
     distributed: "Distributed"
     phase: PhaseType
 

fast_llm/data/dataset/gpt/fim.py

@@ -19,7 +19,7 @@ def __init__(
     ):
         self._config = config
         self._dataset = dataset
-        self._sampling_config = sampling_config
+        self._seed = sampling_config.seed
         self._tokenizer = sampling_config.tokenizer
         if self._tokenizer is None:
             raise ValueError("Fim requires a tokenizer")
@@ -35,9 +35,7 @@ def __len__(self) -> int:
         return len(self._dataset)
 
     def __getitem__(self, idx: int) -> np.ndarray:
-        sample = self._fim(
-            self._dataset[idx], np.random.RandomState(seed=(self._sampling_config.seed + idx) % MAX_SEED)
-        )
+        sample = self._fim(self._dataset[idx], np.random.RandomState(seed=(self._seed + idx) % MAX_SEED))
         return sample
 
     @property

fast_llm/data/dataset/gpt/random.py

@@ -23,14 +23,17 @@ def name(self) -> str:
 class GPTRandomSampledDataset(SampledDataset):
     def __init__(self, config: GPTSamplingConfig, name: str):
         self._name = name
-        self._config = config
+        self._seed = config.seed
+        self._sequence_length = config.sequence_length
+        self._vocab_size = config.vocab_size
+        self._num_samples = config.num_samples
 
     def __len__(self) -> int:
-        return self._config.num_samples
+        return self._num_samples
 
     def __getitem__(self, idx) -> np.ndarray:
-        return np.random.RandomState(self._config.seed + 48576439 + 74593 * idx).randint(
-            0, self._config.vocab_size, size=(self._config.sequence_length + 1,), dtype=np.int64
+        return np.random.RandomState(self._seed + 48576439 + 74593 * idx).randint(
+            0, self._vocab_size, size=(self._sequence_length + 1,), dtype=np.int64
         )
 
     @property

tests/test_dataset.py

@@ -391,7 +391,6 @@ def test_gpt_blended():
         GPTBlendedDatasetConfig,
     ).build_and_sample(get_sampling_config(8, sequence_length=5))
     Assert.eq(len(sampled), 8)
-    print(np.stack([sampled[i] for i in range(8)]).tolist())
     Assert.all_equal(
         np.stack([sampled[i] for i in range(8)]),
         np.array(GPT_BLENDED_EXPECTED_SAMPLES),
@@ -480,7 +479,6 @@ def test_gpt_blended_mixed():
         GPTBlendedDatasetConfig,
     ).build_and_sample(get_sampling_config(8, sequence_length=5))
     Assert.eq(len(sampled), 8)
-    print(np.stack([sampled[i] for i in range(8)]).tolist())
     Assert.all_equal(
         np.stack([sampled[i] for i in range(8)]),
         np.array(GPT_BLENDED_MIXED_EXPECTED_SAMPLES),