from transformers import GPT2LMHeadModel
from transformers import GPT2TokenizerFast
import torch

model = GPT2LMHeadModel.from_pretrained(
    "{HF_USERNAME}/{MODEL_NAME}"
)
tokenizer = GPT2TokenizerFast.from_pretrained(
    "{HF_USERNAME}/{MODEL_NAME}"
)

prompt = "Once upon a time there was a little cat"
inputs = tokenizer(prompt, return_tensors="pt")

outputs = model.generate(
    inputs["input_ids"],
    max_new_tokens     = 200,
    temperature        = 0.8,
    top_p              = 0.9,
    do_sample          = True,
    repetition_penalty = 1.2,
    pad_token_id       = tokenizer.eos_token_id
)

story = tokenizer.decode(
    outputs[0],
    skip_special_tokens = True
)
print(story)

Example Output

"Once upon a time there was a little cat called Mimi. She loved to play with her toys, but one day she got very sad because she couldn't find her favorite toy. They searched everywhere and finally found it under the bed! Mimi was so happy and hugged her mom tight."

Limitations

Generates children's stories only
Works best with story-style prompts
512 token context window

Abdourakib/tinystories-gpt2-124m

TinyStories GPT2 124M

Training Details

How To Use

Example Output

Limitations

Abdourakib/
tinystories-gpt2-124m