import streamlit as st
from transformers import pipeline

# Load the model with caching
@st.cache(allow_output_mutation=True)
def load_model():
    model = pipeline('text-generation', model='gpt2')
    return model

story_generator = load_model()

st.title('Interactive Short-Story Creator')

# Genre selection
genre = st.selectbox(
    "Choose the genre of your story:",
    ('Fantasy', 'Sci-Fi', 'Mystery', 'Horror', 'Adventure', 'Romance')
)

# Additional inputs for a more tailored story
character = st.text_input("Enter a main character name:")
setting = st.text_input("Enter a setting (e.g., a magical forest, a spaceship):")
user_input = st.text_area("Start your story here...", "Once upon a time")

# Advanced configuration options
max_length = st.slider("Select the length of the story continuation:", 100, 1000, 500)
temperature = st.slider("Adjust creativity of the story:", 0.7, 1.5, 1.0)

if st.button('Generate Story Continuation'):
    with st.spinner('AI is writing the story...'):
        # Improved prompt construction with user inputs
        prompt = f"Genre: {genre}. Main character: {character}. Setting: {setting}. Begin the story with: '{user_input}'"
        
        try:
            generated_story = story_generator(prompt, max_length=max_length, temperature=temperature)[0]['generated_text']
            st.text_area("Here's the continuation of your story:", generated_story, height=250)
        except Exception as e:
            st.error(f"Error generating story: {str(e)}")

# Collecting user feedback
feedback = st.text_area("Feedback on the generated story:", "")
if st.button('Submit Feedback'):
    st.write("Thanks for your feedback!")