Initial Draft

app.py

@@ -0,0 +1,170 @@
+import streamlit as st
+import tempfile
+import os
+import re
+import torch
+from threading import Thread
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
+from langchain_community.document_loaders import PyPDFLoader
+from langchain.vectorstores.faiss import FAISS
+from langchain.chains import ConversationalRetrievalChain
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
+
+# Function return langchain document object of PDF pages
+def fn_read_pdf(lv_temp_file_path, mv_processing_message):
+    """Returns langchain document object of PDF pages"""
+
+    lv_pdf_loader = PyPDFLoader(lv_temp_file_path)
+    lv_pdf_content = lv_pdf_loader.load()
+    print("Step2: PDF content extracted")
+    mv_processing_message.text("Step2: PDF content extracted")
+
+    return lv_pdf_content
+
+# Function return FAISS Vector store
+def fn_create_faiss_vector_store(lv_pdf_content, mv_processing_message):
+    """Returns FAISS vector store index of PDF Content"""
+
+    lv_embeddings = HuggingFaceEmbeddings(
+                                            model_name="sentence-transformers/msmarco-distilbert-base-v4",
+                                            model_kwargs={'device': 'cpu'},
+                                            encode_kwargs={'normalize_embeddings': False}
+                                         )
+    lv_vector_store = FAISS.from_documents(lv_pdf_content, lv_embeddings)
+    print("Step3: Vector store created")
+    mv_processing_message.text("Step3: Vector store created")
+
+    return lv_vector_store
+
+# Function return QA Response using Vector Store
+def fn_generate_QnA_response(mv_selected_model, mv_user_question, lv_vector_store, mv_processing_message):
+    """Returns QA Response using Vector Store"""
+
+    lv_chat_history = []
+
+    if "chat_history" not in st.session_state:
+        st.session_state.chat_history = []
+    else:
+        lv_chat_history = st.session_state.chat_history
+    
+    print("Step4: Generating LLM response")
+    mv_processing_message.text("Step4: Generating LLM response")
+
+    lv_tokenizer = AutoTokenizer.from_pretrained(mv_selected_model, trust_remote_code=True)
+    lv_model = AutoModelForCausalLM.from_pretrained(
+                                                        mv_selected_model, 
+                                                        torch_dtype="auto", 
+                                                        device_map="cpu", 
+                                                        trust_remote_code=True
+                                                   )
+    # lv_streamer = TextIteratorStreamer(
+    #                                         tokenizer=lv_tokenizer, 
+    #                                         skip_prompt=True, 
+    #                                         skip_special_tokens=True, 
+    #                                         timeout=300.0
+    #                                   )
+    lv_ms_phi2_pipeline = pipeline(
+                                    "text-generation", tokenizer=lv_tokenizer, model=lv_model,
+                                    device_map="cpu", max_new_tokens=4000, return_full_text=True
+                                  )
+    lv_hf_phi2_pipeline = HuggingFacePipeline(pipeline=lv_ms_phi2_pipeline)
+    lv_chain = ConversationalRetrievalChain.from_llm(lv_hf_phi2_pipeline, lv_vector_store.as_retriever(), return_source_documents=True)
+    lv_response = lv_chain({"question": mv_user_question, 'chat_history': lv_chat_history})
+    
+    lv_chat_history += [(mv_user_question, lv_response["answer"])]
+    st.session_state.chat_history = lv_chat_history
+
+    print("Step5: LLM response generated")
+    mv_processing_message.text("Step5: LLM response generated")
+
+    return lv_response['answer']
+
+
+# Main Function
+def main():
+    
+    # -- Streamlit Settings
+    st.set_page_config(layout='wide')
+    
+    # -- Initialize chat history
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+
+    col1, col2, col3 = st.columns(3)
+    col2.title("Chat with your PDF")
+    st.text("")
+    
+    col1, col2, col3 = st.columns(3)
+    mv_selected_model=col3.selectbox('Select Model',['microsoft/phi-2'])
+    st.text("")
+    st.text("")
+    st.text("")
+    col1, col2, col3 = st.columns(3)
+
+    # -- Reading PDF File
+    mv_pdf_input_file = col2.file_uploader("Choose a PDF file:", type=["pdf"])
+    
+    if 'mv_temp_file_storage_dir' not in st.session_state:
+        mv_temp_file_storage_dir = tempfile.mkdtemp()
+        st.session_state.mv_temp_file_storage_dir = mv_temp_file_storage_dir
+    else:
+        mv_temp_file_storage_dir = st.session_state.mv_temp_file_storage_dir
+
+    mv_processing_message = col2.empty()
+    st.text("")
+    st.text("")
+    st.text("")
+    st.text("")
+    st.text("")
+    st.text("")
+
+    mv_vector_storage_dir = "/workspace/knowledge-base/01-ML/01-dev/adhoc/Talk2PDF/vector_store"
+    
+    if (mv_pdf_input_file is not None):
+        mv_file_name = mv_pdf_input_file.name
+        # mv_vectorstore_file_name = os.path.join(mv_vector_storage_dir, mv_file_name[:-4] + ".vectorstore")
+        # mv_metadata_file_name = os.path.join(mv_vector_storage_dir, mv_file_name[:-4] + ".metadata")
+        
+        if 'lv_vector_store' not in st.session_state:
+            # -- Storing Uploaded PDF locally
+            lv_temp_file_path = os.path.join(mv_temp_file_storage_dir,mv_file_name)
+            with open(lv_temp_file_path,"wb") as lv_file:
+                lv_file.write(mv_pdf_input_file.getbuffer())
+            print("Step1: PDF uploaded successfully at -> " + lv_temp_file_path)    
+            mv_processing_message.text("Step1: PDF uploaded successfully at -> " + lv_temp_file_path)
+        
+            # -- Extracting PDF Text
+            lv_pdf_content = fn_read_pdf(lv_temp_file_path, mv_processing_message)
+
+            # -- Creating FAISS Vector Store
+            lv_vector_store = fn_create_faiss_vector_store(lv_pdf_content, mv_processing_message)
+            st.session_state.lv_vector_store = lv_vector_store
+        else:
+            lv_vector_store = st.session_state.lv_vector_store
+
+        # -- Taking input question and generate answer
+        col1, col2, col3 = st.columns(3)
+        lv_chat_history = col2.chat_message
+
+        if mv_user_question := col2.chat_input("Chat on PDF Data"):
+           # -- Add user message to chat history
+           st.session_state.messages.append({"role": "user", "content": mv_user_question})
+
+           # -- Generating LLM response
+           lv_response = fn_generate_QnA_response(mv_selected_model, mv_user_question, lv_vector_store, mv_processing_message)
+
+           # -- Adding assistant response to chat history
+           st.session_state.messages.append({"role": "assistant", "content": lv_response})
+        
+        # -- Display chat messages from history on app rerun
+        for message in st.session_state.messages:
+            with lv_chat_history(message["role"]):
+                st.markdown(message["content"])
+
+
+
+# Calling Main Function
+if __name__ == '__main__':
+    main()