allow to load only one datasets (#150)

PrimeIntellect-ai · Nov 10, 2024 · 259dfe7 · 259dfe7
1 parent eb9a023
commit 259dfe7
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/zeroband/data.py b/src/zeroband/data.py
@@ -373,7 +373,10 @@ def _load_datasets(
         _ds = ParquetDataset(files=ds_arg["data_files"], tokenizer=tokenizer)
         datasets.append(_ds)
 
-    ds = InterleaveDataset(datasets=datasets, probabilities=probabilities)
+    if len(datasets) > 1:
+        ds = InterleaveDataset(datasets=datasets, probabilities=probabilities)
+    else:
+        ds = datasets[0]
 
     logger.info(f"Loaded datasets ({split})")
     return ds