Internal function to flat_map over. Consumes a batch of input examples and produces a variable number of output examples. Args: x: a single example Returns: a tf.data.Dataset
(x)
| 123 | return new_partial, new_outputs |
| 124 | |
| 125 | def map_fn(x): |
| 126 | """Internal function to flat_map over. |
| 127 | Consumes a batch of input examples and produces a variable number of output |
| 128 | examples. |
| 129 | Args: |
| 130 | x: a single example |
| 131 | Returns: |
| 132 | a tf.data.Dataset |
| 133 | """ |
| 134 | partial = empty_example.copy() |
| 135 | i = tf.zeros([], dtype=tf.int32) |
| 136 | dynamic_batch_size = tf.shape(x[keys[0]])[0] |
| 137 | outputs = {} |
| 138 | for k in keys: |
| 139 | outputs[k] = tf.TensorArray(tf.int32, size=0, dynamic_size=True, element_shape=[key2length[k]]) |
| 140 | outputs[k + "_position"] = tf.TensorArray(tf.int32, size=0, dynamic_size=True, element_shape=[key2length[k]]) |
| 141 | |
| 142 | def body_fn(i, partial, outputs): |
| 143 | """Body function for while_loop. |
| 144 | Args: |
| 145 | i: integer scalar |
| 146 | partial: dictionary of Tensor (partially-constructed example) |
| 147 | outputs: dictionary of TensorArray |
| 148 | Returns: |
| 149 | A triple containing the new values of the inputs. |
| 150 | """ |
| 151 | can_append = True |
| 152 | one_example = {} |
| 153 | for k in keys: |
| 154 | val = tf.cast(x[k][i], tf.int32) |
| 155 | # We consider only the valid tokens i.e., token_id != -1 |
| 156 | val = val[: tf.reduce_sum(tf.cast(tf.not_equal(val, -1), tf.int32))] |
| 157 | one_example[k] = val |
| 158 | for k in keys: |
| 159 | can_append = tf.logical_and( |
| 160 | can_append, tf.less_equal(tf.size(partial[k]) + tf.size(one_example[k]), key2length[k]) |
| 161 | ) |
| 162 | |
| 163 | def false_fn(): |
| 164 | return write_packed_example(partial, outputs) |
| 165 | |
| 166 | def true_fn(): |
| 167 | return partial, outputs |
| 168 | |
| 169 | partial, outputs = tf.cond(can_append, true_fn, false_fn) |
| 170 | new_partial = {} |
| 171 | for k in keys: |
| 172 | new_seq = one_example[k][: key2length[k]] |
| 173 | new_seq_len = tf.size(new_seq) |
| 174 | new_partial[k] = tf.concat([partial[k], new_seq], 0) |
| 175 | new_partial[k + "_position"] = tf.concat([partial[k + "_position"], tf.range(new_seq_len)], 0) |
| 176 | partial = new_partial |
| 177 | return i + 1, partial, outputs |
| 178 | |
| 179 | # For loop over all examples in the batch. |
| 180 | i, partial, outputs = tf.while_loop( |
| 181 | cond=lambda *_: True, |
| 182 | body=body_fn, |
nothing calls this directly
no test coverage detected