import streamlit as st
from transformers import pipeline

# Initialize the GPT-2 model for text generation using Hugging Face pipeline
generator = pipeline("text-generation", model="gpt2")

# Initialize session state to keep track of conversation history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Title and introduction of the app
st.title("RozigGPT - AI Chatbot")
st.markdown("""
Welcome to RozigGPT! 🤖💬
Chat with this AI assistant and get instant responses.
""")

# Function to generate the bot's response
def get_bot_response(user_input):
    # Generate text using GPT-2 model (you can increase max_length for more detailed answers)
    response = generator(user_input, max_length=150, num_return_sequences=1, no_repeat_ngram_size=2)
    return response[0]['generated_text']

# Input field for the user to type their message
user_input = st.text_input("You: ", "")

# Button to clear chat history
if st.button("Clear Chat"):
    st.session_state.messages = []

# Process user input
if user_input:
    # Add user message to session state
    st.session_state.messages.append(f"You: {user_input}")
    
    # Get the bot's response
    bot_response = get_bot_response(user_input)
    st.session_state.messages.append(f"RozigGPT: {bot_response}")

# Display conversation history
for message in st.session_state.messages:
    if message.startswith("You:"):
        # Display user messages on the left side
        st.markdown(f"<p style='text-align: left; color: blue;'>{message}</p>", unsafe_allow_html=True)
    else:
        # Display bot messages on the right side
        st.markdown(f"<p style='text-align: right; color: green;'>{message}</p>", unsafe_allow_html=True)