| 112 | keys_etc = empty_example.keys() |
| 113 | |
| 114 | def write_packed_example(partial, outputs): |
| 115 | new_partial = empty_example.copy() |
| 116 | new_outputs = {} |
| 117 | for k in keys_etc: |
| 118 | # use pad_id to pad inputs/targets, use 0 to pad *_position |
| 119 | pad_id_to_use = 0 if len(k.split("_")) > 1 else pad_id |
| 120 | new_outputs[k] = outputs[k].write( |
| 121 | outputs[k].size(), tf.pad(partial[k], [[0, key2length[k] - tf.size(partial[k])]], constant_values=pad_id_to_use) |
| 122 | ) |
| 123 | return new_partial, new_outputs |
| 124 | |
| 125 | def map_fn(x): |
| 126 | """Internal function to flat_map over. |