Skip to content

Commit ad13fe2

Browse files
author
The TensorFlow Datasets Authors
committed
Fix record_set_id in CroissantBuilder.
PiperOrigin-RevId: 631415903
1 parent a1f762a commit ad13fe2

File tree

1 file changed

+8
-5
lines changed

1 file changed

+8
-5
lines changed

tensorflow_datasets/core/dataset_builders/croissant_builder.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -181,12 +181,14 @@ def __init__(
181181

182182
if not record_set_ids:
183183
record_set_ids = [
184-
py_utils.make_valid_name(record_set.id)
185-
for record_set in self.metadata.record_sets
184+
record_set.id for record_set in self.metadata.record_sets
186185
]
186+
config_names = [
187+
py_utils.make_valid_name(record_set) for record_set in record_set_ids
188+
]
187189
self.BUILDER_CONFIGS: Sequence[dataset_builder.BuilderConfig] = [ # pylint: disable=invalid-name
188-
dataset_builder.BuilderConfig(name=record_set_id)
189-
for record_set_id in record_set_ids
190+
dataset_builder.BuilderConfig(name=config_name)
191+
for config_name in config_names
190192
]
191193

192194
self._disable_shuffling = disable_shuffling
@@ -248,6 +250,7 @@ def _split_generators(
248250
def _generate_examples(
249251
self,
250252
) -> split_builder_lib.SplitGenerator:
251-
records = self.dataset.records(self.builder_config.name)
253+
record_set = self.get_record_set(self.builder_config.name)
254+
records = self.dataset.records(record_set.id)
252255
for i, record in enumerate(records):
253256
yield i, record

0 commit comments

Comments
 (0)