inputs: logits: [batch_size, vocab_size], the values of log logits threshold: int when using top_k, and a probability (0~1) when using top_p outputs: samples: [batch_size]
(self, logits, threshold, num_samples=1)
| 28 | pass |
| 29 | |
| 30 | def sample(self, logits, threshold, num_samples=1): |
| 31 | ''' |
| 32 | inputs: |
| 33 | logits: [batch_size, vocab_size], the values of log logits |
| 34 | threshold: int when using top_k, and a probability (0~1) when using top_p |
| 35 | |
| 36 | outputs: |
| 37 | samples: [batch_size] |
| 38 | ''' |
| 39 | |
| 40 | logits = self.sample_method(logits, threshold) |
| 41 | samples = tf.multinomial(logits, num_samples=num_samples, output_dtype=tf.int32) |
| 42 | samples = tf.reshape(samples, [-1]) |
| 43 | return samples |
| 44 | |
| 45 | def top_k_logits(self, logits, k): |
| 46 | if k == 0: |