Yield mini-batches of indices bucketed by size. Batches may contain sequences of different lengths. Args: indices (List[int]): ordered list of dataset indices num_tokens_fn (callable): function that returns the number of tokens at a given index max_t
(
indices, num_tokens_fn, max_tokens=None, max_sentences=None,
required_batch_size_multiple=1, distributed=False
)
| 87 | |
| 88 | |
| 89 | def batch_by_size( |
| 90 | indices, num_tokens_fn, max_tokens=None, max_sentences=None, |
| 91 | required_batch_size_multiple=1, distributed=False |
| 92 | ): |
| 93 | """ |
| 94 | Yield mini-batches of indices bucketed by size. Batches may contain |
| 95 | sequences of different lengths. |
| 96 | |
| 97 | Args: |
| 98 | indices (List[int]): ordered list of dataset indices |
| 99 | num_tokens_fn (callable): function that returns the number of tokens at |
| 100 | a given index |
| 101 | max_tokens (int, optional): max number of tokens in each batch |
| 102 | (default: None). |
| 103 | max_sentences (int, optional): max number of sentences in each |
| 104 | batch (default: None). |
| 105 | required_batch_size_multiple (int, optional): require batch size to |
| 106 | be a multiple of N (default: 1). |
| 107 | """ |
| 108 | max_tokens = max_tokens if max_tokens is not None else sys.maxsize |
| 109 | max_sentences = max_sentences if max_sentences is not None else sys.maxsize |
| 110 | bsz_mult = required_batch_size_multiple |
| 111 | |
| 112 | if isinstance(indices, types.GeneratorType): |
| 113 | indices = np.fromiter(indices, dtype=np.int64, count=-1) |
| 114 | |
| 115 | sample_len = 0 |
| 116 | sample_lens = [] |
| 117 | batch = [] |
| 118 | batches = [] |
| 119 | for i in range(len(indices)): |
| 120 | idx = indices[i] |
| 121 | num_tokens = num_tokens_fn(idx) |
| 122 | sample_lens.append(num_tokens) |
| 123 | sample_len = max(sample_len, num_tokens) |
| 124 | assert sample_len <= max_tokens, ( |
| 125 | "sentence at index {} of size {} exceeds max_tokens " |
| 126 | "limit of {}!".format(idx, sample_len, max_tokens) |
| 127 | ) |
| 128 | num_tokens = (len(batch) + 1) * sample_len |
| 129 | |
| 130 | if _is_batch_full(batch, num_tokens, max_tokens, max_sentences): |
| 131 | mod_len = max( |
| 132 | bsz_mult * (len(batch) // bsz_mult), |
| 133 | len(batch) % bsz_mult, |
| 134 | ) |
| 135 | batches.append(batch[:mod_len]) |
| 136 | batch = batch[mod_len:] |
| 137 | sample_lens = sample_lens[mod_len:] |
| 138 | sample_len = max(sample_lens) if len(sample_lens) > 0 else 0 |
| 139 | batch.append(idx) |
| 140 | if len(batch) > 0: |
| 141 | batches.append(batch) |
| 142 | return batches |
| 143 | |
| 144 | |
| 145 | def make_positions(tensor, padding_idx): |
nothing calls this directly
no test coverage detected