Initial Draft

app.py

@@ -1,147 +1,249 @@
 import streamlit as st
-import tempfile
 import os
-import re
-import torch
+import requests
+import time
+import sys
 
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
-from langchain_community.document_loaders import PyPDFLoader
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
 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"""
+from langchain.prompts.prompt import PromptTemplate
+from langchain_community.llms import CTransformers
+from langchain.chains import RetrievalQA
 
-    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")
+# Upload pdf file into 'pdf-data' folder if it does not exist
+def fn_upload_pdf(mv_pdf_input_file, mv_processing_message):
+    """Upload pdf file into 'pdf-data' folder if it does not exist"""
 
-    return lv_pdf_content
+    lv_file_name = mv_pdf_input_file.name
 
-# 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"""
+    if not os.path.exists("pdf-data"):
+        os.makedirs("pdf-data")
+    
+    lv_temp_file_path = os.path.join("pdf-data",lv_file_name)
+    
+    if os.path.exists(lv_temp_file_path):
+        print("File already available")
+        fn_display_user_messages("File already available","Warning", mv_processing_message)
+    else:
+        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)
+        fn_display_user_messages("Step1: PDF uploaded successfully at -> " + lv_temp_file_path, "Info", mv_processing_message)
 
+# Create Vector DB of uploaded PDF
+def fn_create_vector_db(mv_pdf_input_file, mv_processing_message):
+    """Create Vector DB of uploaded PDF"""
+
+    lv_file_name = mv_pdf_input_file.name[:-4] + ".vectorstore"
+
+    if not os.path.exists("vectordb/fiaas"):
+        os.makedirs("vectordb/fiaas")
+    
+    lv_temp_file_path = os.path.join("vectordb/fiaas",lv_file_name)
     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")
+                                            model_name="sentence-transformers/all-mpnet-base-v2",
+                                            model_kwargs={'device': 'cpu'}
+                                        )
+    
+    if os.path.exists(lv_temp_file_path):
+        print("VectorDB already available for uploaded file")
+        fn_display_user_messages("VectorDB already available for uploaded file","Warning", mv_processing_message)
+
+        lv_vector_store = FAISS.load_local(lv_temp_file_path, lv_embeddings,allow_dangerous_deserialization=True)
+        return lv_vector_store
+    else:
+        lv_temp_pdf_file_path = os.path.join("pdf-data",mv_pdf_input_file.name)
+        
+        # -- Loading PDF Data
+        lv_pdf_loader = PyPDFLoader(lv_temp_pdf_file_path)
+        lv_pdf_content = lv_pdf_loader.load()
+        print("Step2: PDF content extracted")
+        fn_display_user_messages("Step2: PDF content extracted", "Info", mv_processing_message)
+
+        # -- Chunking PDF Data
+        lv_text_splitter = CharacterTextSplitter(
+                                                    separator="\n",
+                                                    chunk_size=300,
+                                                    chunk_overlap=30,
+                                                    length_function=len
+                                                )
+        lv_pdf_chunk_documents = lv_text_splitter.split_documents(lv_pdf_content)
+        print("Step3: PDF content chucked and document object created")
+        fn_display_user_messages("Step3: PDF content chucked and document object created", "Info", mv_processing_message)
+
+        # -- Creating FIASS Vector Store
+        lv_vector_store = FAISS.from_documents(lv_pdf_chunk_documents, lv_embeddings)
+        print("Step4: Vector store created")
+        fn_display_user_messages("Step4: Vector store created", "Info", mv_processing_message)
+        lv_vector_store.save_local(lv_temp_file_path)
 
-    return lv_vector_store
+        return lv_vector_store
+
+# Display user Error, Warning or Success Message
+def fn_display_user_messages(lv_text, lv_type, mv_processing_message):
+    """Display user Info, Error, Warning or Success Message"""
+    
+    if lv_type == "Success":
+        with mv_processing_message.container(): 
+            st.success(lv_text)
+    elif lv_type == "Error":
+        with mv_processing_message.container(): 
+            st.error(lv_text)
+    elif lv_type == "Warning":
+        with mv_processing_message.container(): 
+            st.warning(lv_text)
+    else:
+        with mv_processing_message.container(): 
+            st.info(lv_text)
+
+# Download TheBloke Models
+def fn_download_llm_models(mv_selected_model, mv_processing_message):
+    """Download TheBloke Models"""
+
+    lv_download_url = ""
+
+    print("Downloading TheBloke of "+mv_selected_model)
+    fn_display_user_messages("Downloading TheBloke of "+mv_selected_model, "Info", mv_processing_message)
+
+    if mv_selected_model == 'microsoft/phi-2':
+        lv_download_url = "https://huggingface.co/TheBloke/phi-2-GGUF/resolve/main/phi-2.Q2_K.gguf"
+    elif mv_selected_model == 'mistralai/Mistral-7B-Instruct-v0.2':
+        lv_download_url = "https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.2-GGUF/resolve/main/mistral-7b-instruct-v0.2.Q2_K.gguf"
+
+    if not os.path.exists("model"):
+        os.makedirs("model")
+    
+    lv_filename = os.path.basename(lv_download_url)
+    lv_temp_file_path = os.path.join("model",lv_filename)
+
+    if os.path.exists(lv_temp_file_path):
+        print("Model already available")
+        fn_display_user_messages("Model already available","Warning", mv_processing_message)
+    else:
+        lv_response = requests.get(lv_download_url, stream=True)
+        if lv_response.status_code == 200:
+            with open(lv_temp_file_path, 'wb') as f:
+                for chunk in lv_response.iter_content(chunk_size=1024):
+                    if chunk:
+                        f.write(chunk)
+            
+            print("Download completed")
+            fn_display_user_messages("Model download completed","Info", mv_processing_message)
+        else:
+            print(f"Model download completed {response.status_code}")
+            fn_display_user_messages(f"Model download completed {response.status_code}","Error", mv_processing_message)
 
 # 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 = []
+    lv_model_path = ""
+    lv_model_type = ""
+    lv_template   = """Instruction:
+                    You are an AI assistant for answering questions about the provided context.
+                    You are given the following extracted parts of a long document and a question. Provide a detailed answer.
+                    If you don't know the answer, just say "Hmm, I'm not sure." Don't try to make up an answer.
+                    =======
+                    {context}
+                    =======
+                    Question: {question}
+                    Output:\n"""
+    lv_qa_prompt = PromptTemplate(
+                                template=lv_template,
+                                input_variables=["question", "context"]
+                              )
 
-    if "chat_history" not in st.session_state:
-        st.session_state.chat_history = []
-    else:
-        lv_chat_history = st.session_state.chat_history
+    if mv_selected_model == 'microsoft/phi-2':
+        lv_model_path = "model/phi-2.Q2_K.gguf"
+        lv_model_type = "pi"
+    elif mv_selected_model == 'mistralai/Mistral-7B-Instruct-v0.2':
+        lv_model_path = "model/mistral-7b-instruct-v0.2.Q2_K.gguf"
+        lv_model_type = "mistral"
+
+    print("Model Absolute location -" +lv_model_path)
     
     print("Step4: Generating LLM response")
-    mv_processing_message.text("Step4: Generating LLM response")
+    fn_display_user_messages("Step4: Generating LLM response","Info", mv_processing_message)
 
-    lv_tokenizer = AutoTokenizer.from_pretrained(
-                                                        mv_selected_model, 
-                                                        model_max_length=2048,
-                                                        trust_remote_code=True
-                                                )
-    lv_model = AutoModelForCausalLM.from_pretrained(
-                                                        mv_selected_model,  
-                                                        device_map="cpu", 
-                                                        trust_remote_code=True
-                                                   )
-    lv_ms_phi2_pipeline = pipeline(
-                                    "text-generation", tokenizer=lv_tokenizer, model=lv_model,
-                                    pad_token_id=lv_tokenizer.eos_token_id, eos_token_id=lv_tokenizer.eos_token_id,
-                                    device_map="cpu", max_new_tokens=2048, 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
+    lv_model = CTransformers(
+                                model="model/mistral-7b-instruct-v0.2.Q2_K.gguf",
+                                model_type=lv_model_type,
+                                max_new_tokens=2048,
+                                temperature=0.00
+                            )
+    lv_retriever = lv_vector_store.as_retriever(search_kwargs={'k': 2})
+    lv_qa_chain = RetrievalQA.from_chain_type(  llm=lv_model,
+                                                chain_type='stuff',
+                                                retriever=lv_retriever,
+                                                return_source_documents=True,
+                                                chain_type_kwargs={'prompt': lv_qa_prompt}
+                                              )
 
-    print("Step5: LLM response generated")
-    mv_processing_message.text("Step5: LLM response generated")
+    lv_response = lv_qa_chain({"query": mv_user_question})
 
-    return lv_response['answer']
+    print("Step5: LLM response generated")
+    fn_display_user_messages("Step5: LLM response generated","Info", mv_processing_message)
 
+    return lv_response['result']
 
 # Main Function
 def main():
     
     # -- Streamlit Settings
     st.set_page_config(layout='wide')
-    
+    col1, col2, col3 = st.columns(3)
+    col2.title("Chat with PDF")
+    st.text("")
+
     # -- Initialize chat history
     if "messages" not in st.session_state:
         st.session_state.messages = []
 
+    # -- Display Supported Models
     col1, col2, col3 = st.columns(3)
-    col2.title("Chat with your PDF")
-    st.text("")
+    mv_selected_model = col3.selectbox('Select Model',
+                                        [
+                                            'microsoft/phi-2',
+                                            'mistralai/Mistral-7B-Instruct-v0.2'
+                                        ]
+                                      )
     
+    # -- Display Supported Vector Stores
     col1, col2, col3 = st.columns(3)
-    mv_selected_model=col3.selectbox('Select Model',['microsoft/phi-2'])
-    st.text("")
+    mv_selected_vector_db = col3.selectbox('Select Vector DB', ['FAISS'])
     st.text("")
-    st.text("")
-    col1, col2, col3 = st.columns(3)
 
     # -- Reading PDF File
+    col1, col2, col3 = st.columns(3)
     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("")
+    # -- Display Processing Details
     st.text("")
+    col1, col2, col3 = st.columns(3)
+    mv_processing_message = col2.empty()
     st.text("")
 
-    mv_vector_storage_dir = "/workspace/knowledge-base/01-ML/01-dev/adhoc/Talk2PDF/vector_store"
-    
+    # -- Downloading Model Files
+    fn_download_llm_models(mv_selected_model, mv_processing_message)
+
+    # -- Processing PDF
     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
+        # -- Upload PDF
+        fn_upload_pdf(mv_pdf_input_file, mv_processing_message)
+
+        # -- Create Vector Index
+        lv_vector_store = fn_create_vector_db(mv_pdf_input_file, mv_processing_message)
 
-        # -- Taking input question and generate answer
+        # -- Perform RAG
         col1, col2, col3 = st.columns(3)
+        st.text("")
         lv_chat_history = col2.chat_message
+        st.text("")
 
         if mv_user_question := col2.chat_input("Chat on PDF Data"):
            # -- Add user message to chat history
@@ -153,11 +255,14 @@ def main():
            # -- 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"])
+           # -- 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"])
+
+        # -- Validate Data
 
+        # -- Get Web Response
 
 # Calling Main Function
 if __name__ == '__main__':