| 15 | import tensorflow as tf |
| 16 | |
| 17 | class Sampling(): |
| 18 | |
| 19 | def __init__(self, sample_method): |
| 20 | if sample_method == "top_k": |
| 21 | self.sample_method = self.top_k_logits |
| 22 | elif sample_method == "top_p": |
| 23 | self.sample_method = self.top_p_logits |
| 24 | else: |
| 25 | print("[ERROR] the sample method should be one of top_k and top_p") |
| 26 | exit(-1) |
| 27 | |
| 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: |
| 47 | return logits |
| 48 | else: |
| 49 | values, _ = tf.nn.top_k(logits, k=k) # [batch size, k] |
| 50 | min_values = values[:, -1, tf.newaxis] #[batch size, 1] |
| 51 | return tf.where( |
| 52 | logits < min_values, |
| 53 | tf.ones_like(logits, dtype=logits.dtype) * logits.dtype.min, |
| 54 | logits |
| 55 | ) |
| 56 | |
| 57 | def top_p_logits(self, logits, p): |
| 58 | sorted_logits = tf.sort(logits, direction='DESCENDING') |
| 59 | sorted_probs = tf.nn.softmax(sorted_logits) |
| 60 | probs_sums = tf.cumsum(sorted_probs, axis=1, exclusive=True) |
| 61 | logits_masked = tf.where( |
| 62 | probs_sums < p, |
| 63 | sorted_logits, |
| 64 | tf.ones_like(sorted_logits) * 1000 |
| 65 | ) # [batchsize, vocab] |
| 66 | min_logits = tf.reduce_min(logits_masked, axis=1, keepdims=True) # [batch size, 1] |
| 67 | return tf.where( |
| 68 | logits < min_logits, |
| 69 | tf.ones_like(logits, dtype=logits.dtype) * logits.dtype.min, |
| 70 | logits |
| 71 | ) |
| 72 | |
| 73 | |
| 74 | |