example.py

import transformers
from transformers import GPT2Tokenizer
from transformers import TFGPT2LMHeadModel

tokenizer =     GPT2Tokenizer.from_pretrained('gpt2-large')
model     = TFGPT2LMHeadModel.from_pretrained('gpt2-large')

print(sample_sequence('What does the fox say?', tokenizer, model, max_length=32, top_k=100, temperature=1.0, repetition_penalty=2.0))

sample_sequence.py

def sample_sequence(context, tokenizer, model, max_length=32, top_k=0, temperature=1.0, repetition_penalty=1.0):
    
    greedy      = temperature == 0.0
    temperature = (temperature if temperature > 0.0 else 1.0)
    tokens      = np.array([tokenizer.encode(context, add_special_tokens=False)], dtype=np.int64)
    
    for step in tqdm(range(max_length), leave=False):
        
        logits, states    = model(tokens[:, -1024:])
        next_token_logits = logits[:, -1, :].numpy()

        if temperature != 1.0: 
            next_token_logits /= temperature

        # repetition penalty from CTRL (https://arxiv.org/abs/1909.05858)
        if repetition_penalty != 1.0:
            for prev_token in set(tokens[0, -1024:].tolist()):
                next_token_logits[0, prev_token] /= repetition_penalty
                
        if top_k > 0:
            top_k_lower_bound = tf.math.top_k(next_token_logits, k=top_k, sorted=True)[0][0,-1].numpy()
            next_token_logits[next_token_logits < top_k_lower_bound] = -np.inf

        if greedy:
            next_token = tf.argmax(next_token_logits, axis=-1)
        else:
            next_token = tf.random.categorical(next_token_logits, 1)

        next_token = tf.reshape(next_token, (1, 1))
            
        tokens = np.concatenate([tokens, next_token], axis=-1)
        
    return tokenizer.decode(tokens[0], skip_special_tokens=True)