Add PyGrain kd.data.Tfds

PiperOrigin-RevId: 732635205
google-research · Mar 2, 2025 · 939b2f6 · 939b2f6
1 parent cb65a07
commit 939b2f6
Showing 1 changed file with 17 additions and 11 deletions.
diff --git a/kauldron/data/__init__.py b/kauldron/data/__init__.py
@@ -37,14 +37,20 @@
 
   from kauldron.data.utils import BatchSize
 
-# ****************************************************************************
-# DO NOT ADD preprocessing ops here. Instead, add them to `kd.contrib.data`
-# ****************************************************************************
-
-# TODO(epot): Should migrate all users to use explicitly `kd.data.tf`
-from kauldron.data.transforms.base import Elements
-from kauldron.data.transforms.base import ElementWiseTransform
-from kauldron.data.transforms.base import TreeFlattenWithPath
-from kauldron.data.transforms.map_transforms import Gather
-from kauldron.data.transforms.map_transforms import Rearrange
-from kauldron.data.transforms.map_transforms import ValueRange
+  from kauldron.data.py.data_sources import DataSource
+  from kauldron.data.py.data_sources import Json
+  from kauldron.data.py.data_sources import HuggingFace
+  from kauldron.data.py.data_sources import Tfds
+  from kauldron.data.py.mixtures import Mix
+
+  # ****************************************************************************
+  # DO NOT ADD preprocessing ops here. Instead, add them to `kd.contrib.data`
+  # ****************************************************************************
+
+  # TODO(epot): Should migrate all users to use explicitly `kd.data.tf`
+  from kauldron.data.transforms.base import Elements
+  from kauldron.data.transforms.base import ElementWiseTransform
+  from kauldron.data.transforms.base import TreeFlattenWithPath
+  from kauldron.data.transforms.map_transforms import Gather
+  from kauldron.data.transforms.map_transforms import Rearrange
+  from kauldron.data.transforms.map_transforms import ValueRange