sample from the huggingface implementation of openai gpt2

generate_from_gpt_2.py

# pip install pytorch-pretrained-bert>=0.6

from pytorch_pretrained_bert.tokenization_gpt2 import GPT2Tokenizer
from pytorch_pretrained_bert.modeling_gpt2 import GPT2LMHeadModel
import torch

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

# The end of text marker.
END_OF_TEXT = tokenizer.encoder["<|endoftext|>"]

SEED = "Twitter is"


def generate(seed: str = SEED, num_steps: int = 20) -> str:
    token_ids = tokenizer.encode(seed)

    # Last value of hidden states
    presents = None
    # Input ids
    inputs = torch.LongTensor([token_ids])

    for _ in range(num_steps):
        # Run model
        logits, presents = model.forward(inputs, past=presents)

        # Sample from logits
        d = torch.distributions.Categorical(logits=logits[0, -1])
        next_id = d.sample().item()

        if next_id == END_OF_TEXT:
            break

        token_ids.append(next_id)
        inputs = torch.LongTensor([[next_id]])

    # Decode
    return tokenizer.decode(token_ids)

print(generate(seed=SEED, num_steps=50))

add a temperature argument to add more diversity. def generate(seed: str = SEED, num_steps: int = 20, temperature: float = 1.0) -> str: and torch.distributions.Categorical(logits=logits[0, -1] / temperature)