import gradio as gr
from transformers import pipeline

# Load the text generation model
text_generator = pipeline("text-generation", model="gpt2")

# Define the function for story generation
def generate_story(prompt, word_count):
    # Calculate the maximum length based on word count
    max_length = word_count + len(prompt.split())
    # Generate a story based on the user's prompt and word count
    generated_text = text_generator(prompt, max_length=max_length, num_return_sequences=1)[0]['generated_text']
    return generated_text

# Define example inputs for the Gradio interface
example_inputs = [
    ["Once upon a time, in a magical forest, there was a curious rabbit named Oliver.", 100],
    ["Amidst the hustle and bustle of a busy city, there lived a lonely street musician.", 150],
    ["On a distant planet, explorers discovered an ancient alien artifact buried in the sand.", 200],
    ["Hidden in the attic of an old house, a forgotten journal revealed a family secret.", 250],
    ["In a futuristic world, a brilliant scientist invented a time-traveling device.", 300],
    ["Deep in the ocean, an underwater explorer encountered a mysterious and ancient creature.", 350]
]

# Create a Gradio interface with examples and a word count slider
iface = gr.Interface(
    fn=generate_story,
    inputs=[
        gr.components.Textbox(label="Prompt"),
        gr.components.Slider(minimum=50, maximum=500, default=100, label="Word Count")
    ],
    outputs="text",
    title="Story Generator with Word Count",
    description="Enter a prompt and select the word count to generate a story.",
    examples=example_inputs
)

# Launch the interface
iface.launch()