Text2Story / app.py
mkoot007's picture
Update app.py
406f121
raw
history blame
1.14 kB
import streamlit as st
from transformers import pipeline, set_seed
# Create a text2text-generation pipeline with the "google/flan-t5-base" model
text_generator = pipeline("text2text-generation", model="google/flan-t5-base")
st.title("Story Generator")
title = st.text_input("Title of the Story:")
character = st.text_input("Character (Hero) Name:")
era = st.selectbox("Era of the Story:", ["Medieval", "Modern", "Future"])
genre = st.selectbox("Genre:", ["Action", "Novelistic", "Love"])
ending = st.selectbox("Ending:", ["Happy", "Sad"])
if st.button("Generate Story"):
if title and character:
# Prepare a structured prompt for the story
prompt = f"Write a {genre} story titled '{title}' set in the {era} era. The main character, {character}, faces a {genre} adventure. In the end, the story has a {ending} ending."
# Generate the story
set_seed(42) # Set a seed for reproducibility
story = text_generator(prompt, max_length=300, do_sample=True)[0]["generated_text"]
st.markdown(story)
else:
st.warning("Please provide a title and character name.")