MachineLearningReply/q-and-a-tool

1from dotenv import load_dotenv2import pandas as pd3import streamlit as st4import streamlit_authenticator as stauth5from streamlit_modal import Modal6 7from utils import new_file, clear_memory, append_documentation_to_sidebar, load_authenticator_config, init_qa, \8    append_header9from haystack.document_stores.in_memory import InMemoryDocumentStore10from haystack import Document11 12load_dotenv()13 14OPENAI_MODELS = ['gpt-3.5-turbo',15                 "gpt-4",16                 "gpt-4-1106-preview"]17 18OPEN_MODELS = [19    'mistralai/Mistral-7B-Instruct-v0.1',20    'HuggingFaceH4/zephyr-7b-beta'21]22 23 24def reset_chat_memory():25    st.button(26        'Reset chat memory',27        key="reset-memory-button",28        on_click=clear_memory,29        help="Clear the conversational memory. Currently implemented to retain the 4 most recent messages.",30        disabled=False)31 32 33def manage_files(modal, document_store):34    open_modal = st.sidebar.button("Manage Files", use_container_width=True)35    if open_modal:36        modal.open()37 38    if modal.is_open():39        with modal.container():40            uploaded_file = st.file_uploader(41                "Upload a document in PDF format",42                type=("pdf",),43                on_change=new_file(),44                disabled=st.session_state['document_qa_model'] is None,45                label_visibility="collapsed",46                help="The document is used to answer your questions. The system will process the document and store it in a RAG to answer your questions.",47            )48            edited_df = st.data_editor(use_container_width=True, data=st.session_state['files'],49                                       num_rows='dynamic',50                                       column_order=['name', 'size', 'is_active'],51                                       column_config={'name': {'editable': False}, 'size': {'editable': False},52                                                      'is_active': {'editable': True, 'type': 'checkbox',53                                                                    'width': 100}}54                                       )55            st.session_state['files'] = pd.DataFrame(columns=['name', 'content', 'size', 'is_active'])56 57            if uploaded_file:58                st.session_state['file_uploaded'] = True59                st.session_state['files'] = pd.concat([st.session_state['files'], edited_df])60                with st.spinner('Processing the document content...'):61                    store_file_in_table(document_store, uploaded_file)62                    ingest_document(uploaded_file)63 64 65def ingest_document(uploaded_file):66    if not st.session_state['document_qa_model']:67        st.warning('Please select a model to start asking questions')68    else:69        try:70            st.session_state['document_qa_model'].ingest_pdf(uploaded_file)71            st.success('Document processed successfully')72        except Exception as e:73            st.error(f"Error processing the document: {e}")74            st.session_state['file_uploaded'] = False75 76 77def store_file_in_table(document_store, uploaded_file):78    pdf_content = uploaded_file.getvalue()79    st.session_state['pdf_content'] = pdf_content80    st.session_state.messages = []81    document = Document(content=pdf_content, meta={"name": uploaded_file.name})82    df = pd.DataFrame(st.session_state['files'])83    df['is_active'] = False84    st.session_state['files'] = pd.concat([df, pd.DataFrame(85        [{"name": uploaded_file.name, "content": pdf_content, "size": len(pdf_content),86          "is_active": True}])])87    document_store.write_documents([document])88 89 90def init_session_state():91    st.session_state.setdefault('files', pd.DataFrame(columns=['name', 'content', 'size', 'is_active']))92    st.session_state.setdefault('models', [])93    st.session_state.setdefault('api_keys', {})94    st.session_state.setdefault('current_selected_model', 'gpt-3.5-turbo')95    st.session_state.setdefault('current_api_key', '')96    st.session_state.setdefault('messages', [])97    st.session_state.setdefault('pdf_content', None)98    st.session_state.setdefault('memory', None)99    st.session_state.setdefault('pdf', None)100    st.session_state.setdefault('document_qa_model', None)101    st.session_state.setdefault('file_uploaded', False)102 103 104def set_page_config():105    st.set_page_config(106        page_title="AI Audit Assistant",107        page_icon=":shark:",108        initial_sidebar_state="expanded",109        layout="wide",110        menu_items={111            'Get Help': 'https://www.extremelycoolapp.com/help',112            'Report a bug': "https://www.extremelycoolapp.com/bug",113            'About': "# This is a header. This is an *extremely* cool app!"114        }115    )116 117 118def update_running_model(api_key, model):119    st.session_state['api_keys'][model] = api_key120    st.session_state['document_qa_model'] = init_qa(model, api_key)121 122 123def init_api_key_dict():124    st.session_state['models'] = OPENAI_MODELS + list(OPEN_MODELS) + ['local LLM']125    for model_name in OPENAI_MODELS:126        st.session_state['api_keys'][model_name] = None127 128 129def display_chat_messages(chat_box, chat_input):130    with chat_box:131        if chat_input:132            for message in st.session_state.messages:133                with st.chat_message(message["role"]):134                    st.markdown(message["content"], unsafe_allow_html=True)135 136            st.chat_message("user").markdown(chat_input)137            with st.chat_message("assistant"):138                # process user input and generate response139                response = st.session_state['document_qa_model'].inference(chat_input, st.session_state.messages)140 141                st.markdown(response)142                st.session_state.messages.append({"role": "user", "content": chat_input})143                st.session_state.messages.append({"role": "assistant", "content": response})144 145 146def setup_model_selection():147    model = st.selectbox(148        "Model:",149        options=st.session_state['models'],150        index=0,  # default to the first model in the list gpt-3.5-turbo151        placeholder="Select model",152        help="Select an LLM:"153    )154 155    if model:156        if model != st.session_state['current_selected_model']:157            st.session_state['current_selected_model'] = model158            if model == 'local LLM':159                st.session_state['document_qa_model'] = init_qa(model)160 161    api_key = st.sidebar.text_input("Enter LLM-authorization Key:", type="password",162                                    disabled=st.session_state['current_selected_model'] == 'local LLM')163    if api_key and api_key != st.session_state['current_api_key']:164        update_running_model(api_key, model)165        st.session_state['current_api_key'] = api_key166 167    return model168 169 170def setup_task_selection(model):171    # enable extractive and generative tasks if we're using a local LLM or an OpenAI model with an API key172    if model == 'local LLM' or st.session_state['api_keys'].get(model):173        task_options = ['Extractive', 'Generative']174    else:175        task_options = ['Extractive']176 177    task_selection = st.sidebar.radio('Select the task:', task_options)178 179    # TODO: Add the task selection logic here (initializing the model based on the task)180 181 182def setup_page_body():183    chat_box = st.container(height=350, border=False)184    chat_input = st.chat_input(185        placeholder="Upload a document to start asking questions...",186        disabled=not st.session_state['file_uploaded'],187    )188    if st.session_state['file_uploaded']:189        display_chat_messages(chat_box, chat_input)190 191 192class StreamlitApp:193    def __init__(self):194        self.authenticator_config = load_authenticator_config()195        self.document_store = InMemoryDocumentStore()196        set_page_config()197        self.authenticator = self.init_authenticator()198        init_session_state()199        init_api_key_dict()200 201    def init_authenticator(self):202        return stauth.Authenticate(203            self.authenticator_config['credentials'],204            self.authenticator_config['cookie']['name'],205            self.authenticator_config['cookie']['key'],206            self.authenticator_config['cookie']['expiry_days']207        )208 209    def setup_sidebar(self):210        with st.sidebar:211            st.sidebar.image("resources/ml_logo.png", use_column_width=True)212 213            # Sidebar for Task Selection214            st.sidebar.header('Options:')215            model = setup_model_selection()216            setup_task_selection(model)217            st.divider()218            self.authenticator.logout()219            reset_chat_memory()220            modal = Modal("Manage Files", key="demo-modal")221            manage_files(modal, self.document_store)222            st.divider()223            append_documentation_to_sidebar()224 225    def run(self):226        name, authentication_status, username = self.authenticator.login()227        if authentication_status:228            self.run_authenticated_app()229        elif st.session_state["authentication_status"] is False:230            st.error('Username/password is incorrect')231        elif st.session_state["authentication_status"] is None:232            st.warning('Please enter your username and password')233 234    def run_authenticated_app(self):235        self.setup_sidebar()236        append_header()237        setup_page_body()238 239 240app = StreamlitApp()241app.run()242