storybot

File path: storybot.py

from transformers import GPT2LMHeadModel, GPT2Tokenizer import torch

class StoryBot: def init(self, model_name: str = 'gpt2', max_length: int = 1000, temperature: float = 0.7, top_p: float = 0.9): # Load the pre-trained GPT-2 model and tokenizer self.tokenizer = GPT2Tokenizer.from_pretrained(model_name) self.model = GPT2LMHeadModel.from_pretrained(model_name) self.max_length = max_length self.temperature = temperature self.top_p = top_p

def __call__(self, user_prompt: str) -> str:
    # Tokenize user input
    inputs = self.tokenizer(user_prompt, return_tensors="pt", truncation=True)

    # Generate output
    story_ids = self.model.generate(
        inputs['input_ids'],
        max_length=self.max_length,
        temperature=self.temperature,
        top_p=self.top_p,
        do_sample=True,
        pad_token_id=self.tokenizer.eos_token_id
    )

    # Decode the generated story to a string
    story = self.tokenizer.decode(story_ids[0], skip_special_tokens=True)
    return story

Testing the bot

if name == "main": # Create the StoryBot instance story_bot = StoryBot(max_length=1200)

# Example user input
user_prompt = "Once upon a time, in a kingdom far away, there lived a brave knight on a quest for the legendary sword."

# Generate and print the story
story = story_bot(user_prompt)
print(story)