import streamlit as st
from transformers import pipeline

# Create a text generation pipeline with the "gpt2" model
pipe = pipeline("text-generation", model="gpt2")

st.title("Poem Generator")
user_word = st.text_input("Enter a word:")

if st.button("Generate Poem"):
    if user_word:
        # Prompt the model with a structured poem request
        poem_prompt = f"Write a poem about '{user_word}'.\n"
        poem = pipe(poem_prompt, max_length=100, do_sample=True, num_return_sequences=1)[0]["generated_text"]
        st.markdown("**Poem:**")
        st.markdown(poem)
    else:
        st.warning("Please enter a word.")