import streamlit as st
import transformers

# Replace "facebook/bart-base" with the desired LLM identifier from Hugging Face
model_name = "facebook/bart-base"
llm = transformers.pipeline("text-generation", model=model_name)


# Display the logo and title
st.image("logo.jpg", width=300)
st.title("Coach Virtual PRODI")

# Initialize a session state variable for history if it doesn't exist
if 'history' not in st.session_state:
    st.session_state['history'] = []

# Function to update the conversation history
def update_history(user_input, ai_response):
    st.session_state['history'].append(("User", user_input))
    st.session_state['history'].append(("AI", ai_response))

# Display the conversation history
for speaker, text in st.session_state['history']:
    if speaker == "User":
        st.text_input("Usuario", value=text, disabled=True)
    else:
        st.text_area("PRODI", value=text, height=75, disabled=True)

# Chat input for user prompt
user_input = st.chat_input("¿Cómo te puedo ayudar hoy?")
if user_input:
    with st.spinner("Generando respuesta..."):
        # Get the AI's response using the loaded llm object
        ai_response = llm(user_input, max_length=1000, do_sample=True, top_k=50, top_p=0.9)["generated_text"][0]
        # Update the conversation history
        update_history(user_input, ai_response)
        # Display the AI's response
        st.write(ai_response)