1  import tensorflow as tf
   2  
   3  import model
   4  
   5  def top_k_logits(logits, k):
   6      if k == 0:
   7          # no truncation
   8          return logits
   9  
  10      def _top_k():
  11          values, _ = tf.nn.top_k(logits, k=k)
  12          min_values = values[:, -1, tf.newaxis]
  13          return tf.where(
  14              logits < min_values,
  15              tf.ones_like(logits, dtype=logits.dtype) * -1e10,
  16              logits,
  17          )
  18      return tf.cond(
  19         tf.equal(k, 0),
  20         lambda: logits,
  21         lambda: _top_k(),
  22      )
  23  
  24  
  25  def top_p_logits(logits, p):
  26      """Nucleus sampling"""
  27      batch, _ = logits.shape.as_list()
  28      sorted_logits = tf.sort(logits, direction='DESCENDING', axis=-1)
  29      cumulative_probs = tf.cumsum(tf.nn.softmax(sorted_logits, axis=-1), axis=-1)
  30      indices = tf.stack([
  31          tf.range(0, batch),
  32          # number of indices to include
  33          tf.maximum(tf.reduce_sum(tf.cast(cumulative_probs <= p, tf.int32), axis=-1) - 1, 0),
  34      ], axis=-1)
  35      min_values = tf.gather_nd(sorted_logits, indices)
  36      return tf.where(
  37          logits < min_values,
  38          tf.ones_like(logits) * -1e10,
  39          logits,
  40      )
  41  
  42  
  43  def sample_sequence(*, hparams, length, start_token=None, batch_size=None, context=None, temperature=1, top_k=0, top_p=1):
  44      if start_token is None:
  45          assert context is not None, 'Specify exactly one of start_token and context!'
  46      else:
  47          assert context is None, 'Specify exactly one of start_token and context!'
  48          context = tf.fill([batch_size, 1], start_token)
  49  
  50      def step(hparams, tokens, past=None):
  51          lm_output = model.model(hparams=hparams, X=tokens, past=past, reuse=tf.AUTO_REUSE)
  52  
  53          logits = lm_output['logits'][:, :, :hparams.n_vocab]
  54          presents = lm_output['present']
  55          presents.set_shape(model.past_shape(hparams=hparams, batch_size=batch_size))
  56          return {
  57              'logits': logits,
  58              'presents': presents,
  59          }
  60  
  61      with tf.name_scope('sample_sequence'):
  62          def body(past, prev, output):
  63              next_outputs = step(hparams, prev, past=past)
  64              logits = next_outputs['logits'][:, -1, :]  / tf.to_float(temperature)
  65              logits = top_k_logits(logits, k=top_k)
  66              logits = top_p_logits(logits, p=top_p)
  67              samples = tf.multinomial(logits, num_samples=1, output_dtype=tf.int32)
  68              return [
  69                  next_outputs['presents'] if past is None else tf.concat([past, next_outputs['presents']], axis=-2),
  70                  samples,
  71                  tf.concat([output, samples], axis=1)
  72              ]
  73  
  74          past, prev, output = body(None, context, context)
  75  
  76          def cond(*args):
  77              return True
  78  
  79          _, _, tokens = tf.while_loop(
  80              cond=cond, body=body,
  81              maximum_iterations=length - 1,
  82              loop_vars=[
  83                  past,
  84                  prev,
  85                  output
  86              ],
  87              shape_invariants=[
  88                  tf.TensorShape(model.past_shape(hparams=hparams, batch_size=batch_size)),
  89                  tf.TensorShape([batch_size, None]),
  90                  tf.TensorShape([batch_size, None]),
  91              ],
  92              back_prop=False,
  93          )
  94  
  95          return tokens
  96