Effectively writes examples to the shards.
(self)
| 378 | self._shuffler.add(key, serialized_example) |
| 379 | |
| 380 | def finalize(self) -> tuple[list[int], int]: |
| 381 | """Effectively writes examples to the shards.""" |
| 382 | if self._shuffler.num_examples == 0: |
| 383 | raise AssertionError("No examples were yielded.") |
| 384 | |
| 385 | shard_specs = _get_shard_specs( |
| 386 | num_examples=self._shuffler.num_examples, |
| 387 | total_size=self._shuffler.size, |
| 388 | bucket_lengths=self._shuffler.bucket_lengths, |
| 389 | filename_template=self._filename_template, |
| 390 | shard_config=self._shard_config, |
| 391 | ) |
| 392 | filename = self._filename_template.sharded_filepaths_pattern() |
| 393 | # Here we just loop over the examples, and don't use the instructions, just |
| 394 | # the final number of examples in every shard. Instructions could be used to |
| 395 | # parallelize, but one would need to be careful not to sort buckets twice. |
| 396 | examples_generator = iter( |
| 397 | utils.tqdm( |
| 398 | self._shuffler, |
| 399 | desc=f"Shuffling {filename}...", |
| 400 | total=self._shuffler.num_examples, |
| 401 | unit=" examples", |
| 402 | leave=False, |
| 403 | mininterval=1.0, |
| 404 | ) |
| 405 | ) |
| 406 | try: |
| 407 | for shard_spec in shard_specs: |
| 408 | iterator = itertools.islice( |
| 409 | examples_generator, 0, shard_spec.examples_number |
| 410 | ) |
| 411 | record_keys = self._example_writer.write(shard_spec.path, iterator) |
| 412 | |
| 413 | # No shard keys returned (e.g: TFRecord format), index cannot be |
| 414 | # created. |
| 415 | if not record_keys: |
| 416 | continue |
| 417 | |
| 418 | # Number of `shard_keys` received should match the number of examples |
| 419 | # written in this shard. |
| 420 | if len(record_keys) != int(shard_spec.examples_number): |
| 421 | raise RuntimeError( |
| 422 | f"Length of example `keys` ({len(record_keys)}) does not match " |
| 423 | f"`shard_spec.examples_number: (`{shard_spec.examples_number})" |
| 424 | ) |
| 425 | _write_index_file(shard_spec.index_path, record_keys) |
| 426 | |
| 427 | except shuffle.DuplicatedKeysError as err: |
| 428 | _raise_error_for_duplicated_keys( |
| 429 | err.item1, err.item2, self._serializer.example_specs |
| 430 | ) |
| 431 | |
| 432 | # Finalize the iterator to clear-up TQDM |
| 433 | try: |
| 434 | val = next(examples_generator) |
| 435 | except StopIteration: |
| 436 | pass |
| 437 | else: |