no truncation nucleus sampling batch logits shape aslist sortedlogits tf sortlogits direction descending axis1 cumulativeprobs tf cumsumtf nn softmaxsortedlogits axis1 axis1 indices tf stack tf range0 batch number of indices to include tf maximumtf reducesumtf castcumulativeprobs p tf int32 axis1 1 0 axis1 minvalues tf gatherndsortedlogits indices return tf where logits minvalues tf oneslikelogits 1e10 logits def samplesequence hparams length starttokennone batchsizenone contextnone temperature1 topk0 topp1 if starttoken is none assert context is not none specify exactly one of starttoken and context else assert context is none specify exactly one of starttoken and context context tf fillbatchsize 1 starttoken def stephparams tokens pastnone lmoutput model modelhparamshparams xtokens pastpast reusetf autoreuse logits lmoutput logits hparams nvocab presents lmoutput present presents setshapemodel pastshapehparamshparams batchsizebatchsize return logits logits presents presents with tf namescope samplesequence def bodypast prev output nextoutputs stephparams prev pastpast logits nextoutputs logits 1 tf tofloattemperature logits topklogitslogits ktopk logits topplogitslogits ptopp samples tf multinomiallogits numsamples1 outputdtypetf int32 return nextoutputs presents if past is none else tf concatpast nextoutputs presents axis2 samples tf concatoutput samples axis1 past prev output bodynone context context def condargs return true tokens tf whileloop condcond bodybody maximumiterationslength 1 loopvars past prev output shapeinvariants tf tensorshapemodel pastshapehparamshparams batchsizebatchsize tf tensorshapebatchsize none tf tensorshapebatchsize none backpropfalse return tokens no truncation nucleus sampling number of indices to include
import tensorflow as tf

import model

def top_k_logits(logits, k):
    if k == 0:
        
        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):
    
    batch, _ = logits.shape.as_list()
    sorted_logits = tf.sort(logits, direction='DESCENDING', axis=-1)
    cumulative_probs = tf.cumsum(tf.nn.softmax(sorted_logits, axis=-1), axis=-1)
    indices = tf.stack([
        tf.range(0, batch),
        
        tf.maximum(tf.reduce_sum(tf.cast(cumulative_probs <= p, tf.int32), axis=-1) - 1, 0),
    ], axis=-1)
    min_values = tf.gather_nd(sorted_logits, indices)
    return tf.where(
        logits < min_values,
        tf.ones_like(logits) * -1e10,
        logits,
    )


def sample_sequence(*, hparams, length, start_token=None, batch_size=None, context=None, temperature=1, top_k=0, top_p=1):
    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'):
        def body(past, prev, output):
            next_outputs = step(hparams, prev, past=past)
            logits = next_outputs['logits'][:, -1, :]  / tf.to_float(temperature)
            logits = top_k_logits(logits, k=top_k)
            logits = top_p_logits(logits, p=top_p)
            samples = tf.multinomial(logits, num_samples=1, output_dtype=tf.int32)
            return [
                next_outputs['presents'] if past is None else tf.concat([past, next_outputs['presents']], axis=-2),
                samples,
                tf.concat([output, samples], axis=1)
            ]

        past, prev, output = body(None, context, context)

        def cond(*args):
            return True

        _, _, tokens = tf.while_loop(
            cond=cond, body=body,
            maximum_iterations=length - 1,
            loop_vars=[
                past,
                prev,
                output
            ],
            shape_invariants=[
                tf.TensorShape(model.past_shape(hparams=hparams, batch_size=batch_size)),
                tf.TensorShape([batch_size, None]),
                tf.TensorShape([batch_size, None]),
            ],
            back_prop=False,
        )

        return tokens