import streamlit as st import os import requests import time import sys from langchain_community.document_loaders import PyPDFLoader from langchain.text_splitter import CharacterTextSplitter from langchain_community.embeddings import HuggingFaceEmbeddings from langchain.vectorstores.faiss import FAISS from langchain.prompts.prompt import PromptTemplate from langchain_community.llms import CTransformers from langchain.chains import RetrievalQA # 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""" lv_file_name = mv_pdf_input_file.name 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/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 # 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_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 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") fn_display_user_messages("Step4: Generating LLM response","Info", mv_processing_message) 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} ) lv_response = lv_qa_chain({"query": mv_user_question}) 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) 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_vector_db = col3.selectbox('Select Vector DB', ['FAISS']) st.text("") # -- Reading PDF File col1, col2, col3 = st.columns(3) mv_pdf_input_file = col2.file_uploader("Choose a PDF file:", type=["pdf"]) # -- Display Processing Details st.text("") col1, col2, col3 = st.columns(3) mv_processing_message = col2.empty() st.text("") # -- Downloading Model Files fn_download_llm_models(mv_selected_model, mv_processing_message) # -- Processing PDF if (mv_pdf_input_file is not None): # -- 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) # -- 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 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"]) # -- Validate Data # -- Get Web Response # Calling Main Function if __name__ == '__main__': main()