Returns list of _ShardSpec instances, corresponding to shards to write. Args: num_examples: int, number of examples in split. total_size: int (bytes), sum of example sizes. bucket_lengths: list of ints, number of examples in each bucket. filename_template: template to format shard
(
num_examples: int,
total_size: int,
bucket_lengths: Sequence[int],
filename_template: naming.ShardedFileTemplate,
shard_config: shard_utils.ShardConfig,
)
| 114 | |
| 115 | |
| 116 | def _get_shard_specs( |
| 117 | num_examples: int, |
| 118 | total_size: int, |
| 119 | bucket_lengths: Sequence[int], |
| 120 | filename_template: naming.ShardedFileTemplate, |
| 121 | shard_config: shard_utils.ShardConfig, |
| 122 | ) -> Sequence[_ShardSpec]: |
| 123 | """Returns list of _ShardSpec instances, corresponding to shards to write. |
| 124 | |
| 125 | Args: |
| 126 | num_examples: int, number of examples in split. |
| 127 | total_size: int (bytes), sum of example sizes. |
| 128 | bucket_lengths: list of ints, number of examples in each bucket. |
| 129 | filename_template: template to format sharded filenames. |
| 130 | shard_config: the configuration for creating shards. |
| 131 | """ |
| 132 | num_shards = shard_config.get_number_shards(total_size, num_examples) |
| 133 | shard_boundaries = shard_utils.get_shard_boundaries(num_examples, num_shards) |
| 134 | shard_specs = [] |
| 135 | bucket_indexes = [str(i) for i in range(len(bucket_lengths))] |
| 136 | from_ = 0 |
| 137 | for shard_index, to in enumerate(shard_boundaries): |
| 138 | # Read the bucket indexes |
| 139 | file_instructions = shard_utils.get_file_instructions( |
| 140 | from_, to, bucket_indexes, bucket_lengths |
| 141 | ) |
| 142 | shard_path = filename_template.sharded_filepath( |
| 143 | shard_index=shard_index, num_shards=num_shards |
| 144 | ) |
| 145 | index_path = _get_index_path(os.fspath(shard_path)) |
| 146 | shard_specs.append( |
| 147 | _ShardSpec( |
| 148 | shard_index=shard_index, |
| 149 | path=os.fspath(shard_path), |
| 150 | index_path=index_path, |
| 151 | examples_number=to - from_, |
| 152 | file_instructions=file_instructions, |
| 153 | ) |
| 154 | ) |
| 155 | from_ = to |
| 156 | return shard_specs |
| 157 | |
| 158 | |
| 159 | def _write_index_file( |
no test coverage detected