(func, state, batch_size, min_length, max_length, pad_id,
bos_id, eos_id)
| 102 | |
| 103 | |
| 104 | def random_sample(func, state, batch_size, min_length, max_length, pad_id, |
| 105 | bos_id, eos_id): |
| 106 | init_seqs = tf.fill([batch_size, 1], bos_id) |
| 107 | init_scores = tf.zeros([batch_size]) |
| 108 | init_flags = tf.zeros([batch_size], tf.bool) |
| 109 | |
| 110 | state = SamplerState( |
| 111 | inputs=init_seqs, |
| 112 | state=state, |
| 113 | scores=init_scores, |
| 114 | flags=init_flags |
| 115 | ) |
| 116 | |
| 117 | max_step = tf.reduce_max(max_length) |
| 118 | |
| 119 | def _is_finished(t, s): |
| 120 | all_finished = tf.reduce_all(s.flags) |
| 121 | cond = tf.logical_and(tf.less(t, max_step), |
| 122 | tf.logical_not(all_finished)) |
| 123 | |
| 124 | return cond |
| 125 | |
| 126 | def _loop_fn(t, s): |
| 127 | outs = _sampling_step(t, func, s, min_length, max_length, pad_id, |
| 128 | eos_id) |
| 129 | return outs |
| 130 | |
| 131 | time = tf.constant(0, name="time") |
| 132 | shape_invariants = SamplerState( |
| 133 | inputs=tf.TensorShape([None, None]), |
| 134 | state=nest.map_structure(utils.infer_shape_invariants, state.state), |
| 135 | scores=tf.TensorShape([None]), |
| 136 | flags=tf.TensorShape([None]) |
| 137 | ) |
| 138 | outputs = tf.while_loop(_is_finished, _loop_fn, [time, state], |
| 139 | shape_invariants=[tf.TensorShape([]), |
| 140 | shape_invariants], |
| 141 | parallel_iterations=1, |
| 142 | back_prop=False) |
| 143 | |
| 144 | final_state = outputs[1] |
| 145 | final_seqs = final_state.inputs |
| 146 | final_scores = final_state.scores |
| 147 | |
| 148 | return final_seqs, final_scores |
| 149 | |
| 150 | |
| 151 | def create_sampling_graph(models, features, params): |
no test coverage detected