import streamlit as st 
from langchain_community.llms import HuggingFaceEndpoint
from langchain.prompts import PromptTemplate
#from transformers import pipeline, set_seed

#st.set_page_config(page_title="Poetry Generaion Demo", page_icon=":robot:")
#st.header("Hey, I am your poetry generator")

st.title("❤️ Personalized Poetry Writing Tool ❤️")



theme = st.text_input("Enter the theme of the poem (eg. Rain, Sunshine, Tomato, etc.)")
style = st.text_input("Enter the style for the poem (eg. Shakespear, Little Girl, Professional, Student, etc.)")
keywords = st.text_input("Enter the keywords to be used in poem (comma-separated)")

llm = HuggingFaceEndpoint(
    repo_id="microsoft/Phi-3-mini-4k-instruct",
    task="text-generation",
    max_new_tokens=512,
    do_sample=False,
    repetition_penalty=1.03,
    temperature=0.6,
)


def generate_poem(theme, style, keywords):
    """Generates a poem based on user input."""

    prompt = f"Write a poem about {theme} in the style of {style} using these keywords: {keywords}"
    #output = generator(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
    output = llm.invoke(prompt, max_length=100, num_return_sequences=1)
    return output

submit = st.button("Enter to generate the Poem")

if submit:
    st.subheader("Here is the poem: ")
    st.write(generate_poem(theme, style, keywords))