PR feedback

Author: nikg4

src/oumi/builders/data.py

@@ -24,7 +24,7 @@
 )
 from oumi.datasets.trl_dpo_preprocessor import trl_dpo_chat_preprocessor_fn
 from oumi.datasets.ultrachat_200k import trl_sft_ultrachat_200k_preprocessor_fn
-from oumi.utils.hf_datasets_utils import is_disk_cached_hf_dataset
+from oumi.utils.hf_datasets_utils import is_cached_to_disk_hf_dataset
 from oumi.utils.logging import logger
 
 DatasetType = TypeVar("DatasetType", datasets.Dataset, datasets.IterableDataset)
@@ -371,7 +371,7 @@ def _load_dataset(
             return dataset.to_hf()
 
     dataset_name_or_path: Path = Path(dataset_params.dataset_name)
-    if is_disk_cached_hf_dataset(dataset_name_or_path):
+    if is_cached_to_disk_hf_dataset(dataset_name_or_path):
         return datasets.Dataset.load_from_disk(dataset_name_or_path)
     else:
         return datasets.load_dataset(

src/oumi/core/datasets/base_dataset.py

@@ -9,7 +9,7 @@
 
 from oumi.core.tokenizers import BaseTokenizer
 from oumi.core.types.turn import Conversation
-from oumi.utils.hf_datasets_utils import is_disk_cached_hf_dataset
+from oumi.utils.hf_datasets_utils import is_cached_to_disk_hf_dataset
 from oumi.utils.logging import logger
 
 
@@ -135,7 +135,7 @@ def _load_data(self) -> pd.DataFrame:
                 and dataset_path.is_file()
             ):
                 result = self._load_parquet_dataset(self.dataset_name_or_path)
-            elif is_disk_cached_hf_dataset(self.dataset_name_or_path):
+            elif is_cached_to_disk_hf_dataset(self.dataset_name_or_path):
                 result = self._load_dataset_from_disk(self.dataset_name_or_path)
             else:
                 raise ValueError(

src/oumi/utils/hf_datasets_utils.py

@@ -4,7 +4,7 @@
 from oumi.utils.logging import logger
 
 
-def is_disk_cached_hf_dataset(dataset_name_or_path: Union[str, Path]) -> bool:
+def is_cached_to_disk_hf_dataset(dataset_name_or_path: Union[str, Path]) -> bool:
     """Detects whether a dataset was saved using `dataset.save_to_disk()`.
 
     Such datasets should be loaded using `datasets.Daataset.load_from_disk()`

tests/utils/test_hf_datasets_utils.py

@@ -3,7 +3,7 @@
 
 import datasets
 
-from oumi.utils.hf_datasets_utils import is_disk_cached_hf_dataset
+from oumi.utils.hf_datasets_utils import is_cached_to_disk_hf_dataset
 
 
 def test_is_saved_to_disk_hf_dataset():
@@ -12,16 +12,16 @@ def test_is_saved_to_disk_hf_dataset():
             {"pokemon": ["bulbasaur", "squirtle"], "type": ["grass", "water"]}
         )
         ds_dir = Path(output_temp_dir) / "toy_dataset"
-        assert not is_disk_cached_hf_dataset(ds_dir)
+        assert not is_cached_to_disk_hf_dataset(ds_dir)
 
         ds_dir.mkdir(parents=True, exist_ok=True)
-        assert not is_disk_cached_hf_dataset(ds_dir)
+        assert not is_cached_to_disk_hf_dataset(ds_dir)
 
         ds.save_to_disk(ds_dir, num_shards=2)
-        assert is_disk_cached_hf_dataset(ds_dir)
+        assert is_cached_to_disk_hf_dataset(ds_dir)
 
         for filename in ("dataset_info.json", "state.json"):
             sub_path: Path = Path(ds_dir) / filename
             assert sub_path.exists() and sub_path.is_file()
             sub_path.unlink()
-            assert not is_disk_cached_hf_dataset(ds_dir)
+            assert not is_cached_to_disk_hf_dataset(ds_dir)