import streamlit as st
from transformers import Conversation, pipeline

st.set_page_config(page_title="Dolphin Chatbot", page_icon=":robot_face:")
st.header("Dolphin 2.6 Mistral 7B Chatbot")

# Initialize the conversation pipeline with dolphin-2.6-mistral-7b model
conversational_pipeline = pipeline('conversational', model='google/flan-t5-large')

if "conversation" not in st.session_state:
    st.session_state.conversation = Conversation()

def load_answer(question):
    new_user_input = Conversation(text=question)
    st.session_state.conversation.add_user_input(question)
    assistant_answer = conversational_pipeline(st.session_state.conversation)
    st.session_state.conversation = assistant_answer
    return assistant_answer.text

def get_text():
    input_text = st.text_input("You: ", key="input_text")
    return input_text

user_input = get_text()
submit = st.button('Generate')

if submit:
    response = load_answer(user_input)
    st.subheader("Answer:")
    st.write(response)