Move Croissant related functions to croissant_utils.py

PiperOrigin-RevId: 642256699

Author: fineguy

tensorflow_datasets/core/dataset_builders/croissant_builder.py

@@ -49,6 +49,7 @@
 from tensorflow_datasets.core.features import image_feature
 from tensorflow_datasets.core.features import sequence_feature
 from tensorflow_datasets.core.features import text_feature
+from tensorflow_datasets.core.utils import croissant_utils
 from tensorflow_datasets.core.utils import huggingface_utils
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils import version as version_utils
@@ -173,9 +174,7 @@ def __init__(
     if mapping is None:
       mapping = {}
     self.dataset = mlc.Dataset(jsonld, mapping=mapping)
-    self.name = huggingface_utils.get_tfds_name_from_croissant_dataset(
-        self.dataset
-    )
+    self.name = croissant_utils.get_tfds_dataset_name(self.dataset)
     self.metadata = self.dataset.metadata
 
     # In TFDS, version is a mandatory attribute, while in Croissant it is only a

tensorflow_datasets/core/utils/croissant_utils.py

@@ -0,0 +1,41 @@
+# coding=utf-8
+# Copyright 2024 The TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utility functions for croissant_builder."""
+
+from __future__ import annotations
+
+import typing
+
+from tensorflow_datasets.core.utils import huggingface_utils
+
+if typing.TYPE_CHECKING:
+  # pylint: disable=g-bad-import-order
+  import mlcroissant as mlc
+
+_HUGGINGFACE_URL_PREFIX = 'https://huggingface.co/datasets/'
+
+
+def get_dataset_name(dataset: mlc.Dataset) -> str:
+  """Returns dataset name of the given MLcroissant dataset."""
+  if (url := dataset.metadata.url) and url.startswith(_HUGGINGFACE_URL_PREFIX):
+    return url.removeprefix(_HUGGINGFACE_URL_PREFIX)
+  return dataset.metadata.name
+
+
+def get_tfds_dataset_name(dataset: mlc.Dataset) -> str:
+  """Returns TFDS compatible dataset name of the given MLcroissant dataset."""
+  dataset_name = get_dataset_name(dataset)
+  return huggingface_utils.convert_hf_name(dataset_name)

tensorflow_datasets/core/utils/croissant_utils_test.py

@@ -0,0 +1,36 @@
+# coding=utf-8
+# Copyright 2024 The TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import mlcroissant as mlc
+import pytest
+from tensorflow_datasets.core.utils import croissant_utils
+
+
+@pytest.mark.parametrize(
+    'croissant_name,croissant_url,tfds_name',
+    [
+        (
+            'Name+1',
+            'https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k',
+            'huggingfaceh4__ultrachat_200k',
+        ),
+        ('Name+1', 'bad_url', 'name_1'),
+        ('Name+1', None, 'name_1'),
+    ],
+)
+def test_get_tfds_dataset_name(croissant_name, croissant_url, tfds_name):
+  metadata = mlc.Metadata(name=croissant_name, url=croissant_url)
+  dataset = mlc.Dataset.from_metadata(metadata)
+  assert croissant_utils.get_tfds_dataset_name(dataset) == tfds_name

tensorflow_datasets/core/utils/huggingface_utils.py

@@ -19,7 +19,6 @@
 
 from collections.abc import Mapping, Sequence
 import datetime
-import typing
 from typing import Any, Type, TypeVar
 
 from etils import epath
@@ -33,10 +32,6 @@
 from tensorflow_datasets.core.utils.lazy_imports_utils import datasets as hf_datasets
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
 
-if typing.TYPE_CHECKING:
-  # pylint: disable=g-bad-import-order
-  import mlcroissant as mlc
-
 
 _HF_DTYPE_TO_NP_DTYPE = immutabledict.immutabledict({
     'bool': np.bool_,
@@ -249,16 +244,6 @@ def convert_hf_value(
   )
 
 
-def get_tfds_name_from_croissant_dataset(dataset: mlc.Dataset) -> str:
-  """Returns TFDS compatible dataset name of the given MLcroissant dataset."""
-  if (url := dataset.metadata.url) and url.startswith(
-      'https://huggingface.co/datasets/'
-  ):
-    url_suffix = url.removeprefix('https://huggingface.co/datasets/')
-    return convert_hf_name(url_suffix)
-  return convert_hf_name(dataset.metadata.name)
-
-
 def convert_hf_name(hf_name: str) -> str:
   """Converts Huggingface name to a TFDS compatible dataset name.
 

tensorflow_datasets/core/utils/huggingface_utils_test.py

@@ -22,7 +22,6 @@
 from tensorflow_datasets.core import lazy_imports_lib
 from tensorflow_datasets.core import registered
 from tensorflow_datasets.core.utils import huggingface_utils
-from tensorflow_datasets.core.utils.lazy_imports_utils import mlcroissant as mlc
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
 
 
@@ -217,29 +216,6 @@ def test_convert_value(hf_value, feature, expected_value):
   assert huggingface_utils.convert_hf_value(hf_value, feature) == expected_value
 
 
-@pytest.mark.parametrize(
-    'croissant_name,croissant_url,tfds_name',
-    [
-        (
-            'Name+1',
-            'https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k',
-            'huggingfaceh4__ultrachat_200k',
-        ),
-        ('Name+1', 'bad_url', 'name_1'),
-        ('Name+1', None, 'name_1'),
-    ],
-)
-def test_get_tfds_name_from_croissant_dataset(
-    croissant_name, croissant_url, tfds_name
-):
-  metadata = mlc.Metadata(name=croissant_name, url=croissant_url)
-  dataset = mlc.Dataset.from_metadata(metadata)
-  assert (
-      huggingface_utils.get_tfds_name_from_croissant_dataset(dataset)
-      == tfds_name
-  )
-
-
 @pytest.mark.parametrize(
     'hf_name,tfds_name',
     [