text: import tensorflow as tf
  import os, sys
  sys.path.append(os.curdir+'/src')
  import model
  
  
  def top_k_logits(logits, k):
      if k == 0:
          # no truncation
          return logits
  
      def _top_k():
          values, _ = tf.nn.top_k(logits, k=k)
          min_values = values[:, -1, tf.newaxis]
          return tf.where(
              logits < min_values,
              tf.ones_like(logits, dtype=logits.dtype) * -1e10,
              logits,
          )
      return tf.cond(
         tf.equal(k, 0),
         lambda: logits,
         lambda: _top_k(),
      )
  
  
  def top_p_logits(logits, p):
      with tf.variable_scope('top_p_logits'):
          logits_sort = tf.sort(logits, direction='DESCENDING')
          probs_sort = tf.nn.softmax(logits_sort)
          probs_sums = tf.cumsum(probs_sort, axis=1, exclusive=True)
          logits_masked = tf.where(probs_sums < p, logits_sort, tf.ones_like(logits_sort)*1000) # [batchsize, vocab]
          min_logits = tf.reduce_min(logits_masked, axis=1, keepdims=True) # [batchsize, 1]
          return tf.where(
              logits < min_logits,
              tf.ones_like(logits, dtype=logits.dtype) * -1e10,
              logits,
          )
  
  
  def sample_sequence(*, hparams, length, start_token=None, batch_size=None, context=None, temperature=1, top_k=0, top_p=0.0):
      if start_token is None:
          assert context is not None, 'Specify exactly one of start_token and context!'
      else:
          assert context is None, 'Specify exactly one of start_token and context!'
          context = tf.fill([batch_size, 1], start_token)
  
      def step(hparams, tokens, past=None):
          lm_output = model.model(hparams=hparams, X=tokens, past=past, reuse=tf.AUTO_REUSE)
  
          logits = lm_output['logits'][:, :, :hparams.n_vocab]
          presents = lm_output['present']
          presents.set_shape(model.past_shape(hparams=hparams, batch_size=batch_size))
          return {
              'logits': logits,
              'presents': presents,
          }
  
      with tf.name_scope('sample_sequence'):
          # Don't feed the last context token -- leave that to the loop below
          # TODO: Would be slightly faster if we called step on the entire context,
          # rather than leaving the last token transformer calculation to the while loop.
          context_output = step(hparams, context[:, :-1])
  
          def body(past, prev, output):
              next_outputs = step(hparams, prev[:, tf.newaxis], past=past)
              logits = next_outputs['logits'][:, -1, :]  / tf.to_float(temperature)
              if top_p > 0.0:
                  logits = top_p_logits(logits, p=top_p)
              else:
                  logits = top_k_logits(logits, k=top_k)
              samples = tf.multinomial(logits, num_samples=1, output_dtype=tf.int32, seed=0)
              #samples = tf.expand_dims(tf.nn.top_k(tf.squeeze(logits), k=1)[-1], 0)
              return [
                  tf.concat([past, next_outputs['presents']], axis=-2),
                  tf.squeeze(samples, axis=[1]),
                  tf.concat([output, samples], axis=1),
              ]
  
          def cond(*args):
              return True
  
          _, _, tokens = tf.while_loop(
              cond=cond, body=body,
              maximum_iterations=length,
              loop_vars=[
                  context_output['presents'],
                  context[:, -1],
                  context,
              ],
              shape_invariants=[
                  tf.TensorShape(model.past_shape(hparams=hparams, batch_size=batch_size)),
                  tf.TensorShape([batch_size]),
                  tf.TensorShape([batch_size, None]),
              ],
              back_prop=False,
          )
  
          return tokens
  