import streamlit as st
from transformers import Conversation, pipeline

# Initialize the chat pipeline with the dolphin-2.6-mistral-7b model
chat_pipeline = pipeline("conversational", model="cognitivecomputations/dolphin-2.6-mistral-7b")

# Set page configuration for the Streamlit app
st.set_page_config(page_title="Dolphin Chatbot", page_icon=":robot_face:")
st.header("Dolphin 2.6 Mistral 7B Chatbot")

# Initialize the conversation object
if "conversation" not in st.session_state:
    st.session_state.conversation = Conversation()

# Function to get user input and generate a response
def load_answer(question):
    st.session_state.conversation.add_user_input(question)
    responses = chat_pipeline(st.session_state.conversation)
    # The latest response is the last element in the list
    return responses.generations[-1].text

# Function to display the input bar and detect user input
def get_input():
    return st.text_input("You:", key="input")

# Display input bar and wait for user input
user_input = get_input()

# Button to generate the response
submit = st.button('Generate')

# Actions to take when the 'Generate' button is clicked
if submit:
    if user_input:
        # Get the assistant's response
        response = load_answer(user_input)
        # Display the response
        st.subheader("Answer:")
        st.write(response)
    else:
        # If no user input, prompt the user to enter a question
        st.warning("Please enter a question for the chatbot.")