Support config for the hugginface loader

PiperOrigin-RevId: 729540335
google-research · Feb 21, 2025 · 22c8ac7 · 22c8ac7
1 parent 746dbc8
commit 22c8ac7
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/kauldron/data/py/data_sources.py b/kauldron/data/py/data_sources.py
@@ -70,17 +70,21 @@ def data_source(self) -> grain.RandomAccessDataSource:
 class HuggingFace(base.DataSourceBase):
   """HuggingFace loader."""
 
-  name: str
+  path: str
   _: dataclasses.KW_ONLY
+  config: str | None = None
   split: str
   data_dir: epath.PathLike | None = None
+  cache_dir: epath.PathLike | None = None
 
   @functools.cached_property
   def data_source(self) -> grain.RandomAccessDataSource:
     return datasets.load_dataset(
-        self.name,
+        self.path,
+        name=self.config,
         split=self.split,
         data_dir=self.data_dir,
+        cache_dir=self.cache_dir,
     )