Support loading datasetes saved via save_to_disk

axolotl-ai-cloud · Mar 22, 2024 · e75de2b · e75de2b
1 parent 4e69aa4
commit e75de2b
Showing 1 changed file with 11 additions and 8 deletions.
diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py
@@ -1,4 +1,5 @@
 """Module containing data utilities"""
+
 import functools
 import hashlib
 import logging
@@ -290,14 +291,16 @@ def for_d_in_datasets(dataset_configs):
             local_path = Path(config_dataset.path)
             if local_path.exists():
                 if local_path.is_dir():
-                    # TODO dirs with arrow or parquet files could be loaded with `load_from_disk`
-                    ds = load_dataset(
-                        config_dataset.path,
-                        name=config_dataset.name,
-                        data_files=config_dataset.data_files,
-                        streaming=False,
-                        split=None,
-                    )
+                    if config_dataset.data_files:
+                        ds = load_dataset(
+                            config_dataset.path,
+                            name=config_dataset.name,
+                            data_files=config_dataset.data_files,
+                            streaming=False,
+                            split=None,
+                        )
+                    else:
+                        ds = load_from_disk(config_dataset.path)
                 elif local_path.is_file():
                     ds_type = get_ds_type(config_dataset)