Fix conversion of HF datasets.

PiperOrigin-RevId: 623162414

Author: fineguy

tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder.py

@@ -116,7 +116,8 @@ def get_serialized_examples_iter():
         split=shard_spec.shard_split, run_post_process=False
     ):
       example = huggingface_utils.convert_hf_value(hf_value, features)
-      serialized_example = serializer.serialize_example(example)
+      encoded_example = features.encode_example(example)
+      serialized_example = serializer.serialize_example(encoded_example)
       num_bytes += len(serialized_example)
       yield serialized_example
 

tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder_test.py

@@ -16,9 +16,13 @@
 from unittest import mock
 
 import datasets as hf_datasets
+import numpy as np
 import pytest
+from tensorflow_datasets.core import lazy_imports_lib
 from tensorflow_datasets.core.dataset_builders import huggingface_dataset_builder
 
+PIL_Image = lazy_imports_lib.lazy_imports.PIL_Image
+
 
 class DummyHuggingfaceBuilder(hf_datasets.GeneratorBasedBuilder):
 
@@ -35,7 +39,11 @@ def _split_generators(self, dl_manager):
 
   def _generate_examples(self):
     for i in range(2):
-      yield i, {'feature': i}
+      yield i, {
+          'number': i,
+          'text': f'{i}',
+          'image': PIL_Image.new(mode='L', size=(4, 4)),
+      }
 
   def download_and_prepare(self, *args, **kwargs):
     # Disable downloads from GCS
@@ -87,8 +95,17 @@ def mock_huggingface_dataset_builder(
 def test_download_and_prepare(builder):
   builder.download_and_prepare()
   ds = builder.as_data_source()
+  expected_image = PIL_Image.new(mode='RGB', size=(4, 4))
   # Split names are sanitized, eg train.clean -> train_clean
-  assert list(ds['train_clean']) == [{'feature': 0}, {'feature': 1}]
+  for element, expected in zip(
+      ds['train_clean'],
+      [
+          {'number': 0, 'text': b'0', 'image': expected_image},
+          {'number': 1, 'text': b'1', 'image': expected_image},
+      ],
+  ):
+    for feature in ['number', 'text', 'image']:
+      assert np.array_equal(element[feature], expected[feature])
 
 
 def test_all_parameters_are_passed_down_to_hf(builder):
@@ -98,4 +115,8 @@ def test_all_parameters_are_passed_down_to_hf(builder):
 
 
 def test_hf_features(builder):
-  assert builder._hf_features() == {'feature': hf_datasets.Value('int64')}
+  assert builder._hf_features() == {
+      'number': hf_datasets.Value('int64'),
+      'text': hf_datasets.Value('string'),
+      'image': hf_datasets.Image(),
+  }