from transformers import pipeline
import gradio as gr

story_generator = pipeline("text-generation", model="tiiuae/falcon-7b-instruct")


# Define the story generation function
def generate_story(topic, age_range):
    prompt = f"Write a children’s story for ages {age_range} about {topic}. Use simple language and a friendly tone."
    story = story_generator(prompt, max_length=400, do_sample=True, temperature=0.8)
    return story[0]['generated_text']

# Launch Gradio interface
iface = gr.Interface(
    fn=generate_story,
    inputs=[
        gr.Textbox(label="Story Topic (e.g., dragons, friendship)"),
        gr.Textbox(label="Age Range (e.g., 5-7)")
    ],
    outputs="text",
    title="Children’s Book Generator",
    description="Enter a topic and age range to generate a friendly, age-appropriate children’s story!"
)

iface.launch()