tiny-stories-19M-instruct

from transformers import AutoTokenizer, AutoModelForCausalLM

prompt = '''
Story: Once upon a time, there was a big, fat penguin named Puddles. Puddles loved to play with his friends on the ice.
One day, Puddles saw a big block of ice and decided to cut it. He used his sharp beak to cut the ice into small pieces.
Puddles and his friends had fun sliding on the ice pieces. They laughed and played until it was time to go home.
Puddles went to bed that night feeling happy and proud that he was able to cut the ice.

Summary:
'''

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")

model_id = 'SauravP97/tiny-stories-19M-instruct'
pretrained_model = AutoModelForCausalLM.from_pretrained(model_id)

inputs = tokenizer(prompt, return_tensors="pt")

# Generate
output_tokens = pretrained_model.generate(
    inputs.input_ids,
    max_new_tokens=100, 
    do_sample=True,
    temperature=1, 
    top_k=50,
    pad_token_id=tokenizer.eos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)

print(tokenizer.decode(output_tokens)[0])