LangChainDemo/OPM_Retirement_Assistant
0
1import streamlit as st2from pypdf import PdfReader3import os4from pathlib import Path5from dotenv import load_dotenv6import pickle7import timeit8from PIL import Image9import zipfile10import datetime11import shutil12from collections import defaultdict13import pandas as pd14 15from langchain.embeddings import HuggingFaceEmbeddings16from langchain.document_loaders import PyPDFLoader17from langchain.text_splitter import RecursiveCharacterTextSplitter18from langchain.document_loaders import PyPDFLoader, DirectoryLoader19from langchain.memory import ConversationBufferMemory20from langchain.chains import ConversationalRetrievalChain21from langchain.prompts.prompt import PromptTemplate22from langchain.vectorstores import Chroma23from langchain.document_loaders import PyPDFDirectoryLoader24from langchain.retrievers import BM25Retriever, EnsembleRetriever25from langchain.document_loaders import UnstructuredHTMLLoader26from langchain.llms import OpenAI27from langchain.chat_models import ChatOpenAI28from langchain.agents.agent_toolkits import create_retriever_tool29from langchain.agents.agent_toolkits import create_conversational_retrieval_agent30from langchain.utilities import SerpAPIWrapper31from langchain.agents import Tool32from langchain.agents import load_tools33from langchain.chat_models import ChatOpenAI34from langchain.retrievers.multi_query import MultiQueryRetriever35from langchain.chains import RetrievalQA36from langchain.retrievers import ContextualCompressionRetriever37from langchain.retrievers.document_compressors import CohereRerank38 39import logging40 41 42load_dotenv()43 44 45current_timestamp = datetime.datetime.now()46timestamp_string = current_timestamp.strftime("%Y-%m-%d %H:%M:%S")47 48 49def build_llm():50 '''51 Loading OpenAI model52 '''53 # llm= OpenAI(temperature=0.2)54 llm= ChatOpenAI(temperature = 0)55 return llm56 57def build_embedding_model():58 '''59 Loading Sentence transformer model for text embedding60 '''61 embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2',62 model_kwargs={'device': 'cpu'})63 return embeddings64 65def unzip_opm():66 '''67 This function is used to unzip the documents file. This is required if there is no extisting vector database68 created and wanted to build from the scratch69 '''70 # Specify the path to your ZIP file71 zip_file_path = r'OPM_Files/OPM_Retirement_backup-20230902T130906Z-001.zip'72 73 # Get the directory where the ZIP file is located74 extract_path = os.path.dirname(zip_file_path)75 76 # Create a folder with the same name as the ZIP file (without the .zip extension)77 extract_folder = os.path.splitext(os.path.basename(zip_file_path))[0]78 extract_folder_path = os.path.join(extract_path, extract_folder)79 80 # Create the folder if it doesn't exist81 if not os.path.exists(extract_folder_path):82 os.makedirs(extract_folder_path)83 84 # Open the ZIP file for reading85 with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:86 # Extract all the contents into the created folder87 zip_ref.extractall(extract_folder_path)88 89 print(f'Unzipped {zip_file_path} to {extract_folder_path}')90 return extract_folder_path91 92 93 94 95 96 return 97 98def count_files_by_type(folder_path):99 '''100 Counting files by file type in the specified folder.101 This is required if there is no extisting vector database102 created and wanted to build from the scratch103 '''104 file_count_by_type = defaultdict(int)105 106 for root, _, files in os.walk(folder_path):107 for file in files:108 _, extension = os.path.splitext(file)109 file_count_by_type[extension] += 1110 111 return file_count_by_type112 113def generate_file_count_table(file_count_by_type):114 '''115 Generate a table files count file type.116 This is required if there is no extisting vector database117 created and wanted to build from the scratch118 '''119 data = {"File Type": [], "Number of Files": []}120 for extension, count in file_count_by_type.items():121 data["File Type"].append(extension)122 data["Number of Files"].append(count)123 124 df = pd.DataFrame(data)125 df = df.sort_values(by="Number of Files", ascending=False) # Sort by number of files126 return df127 128def move_files_to_folders(folder_path):129 '''130 Move files to respective folder. Example, PDF docs to PDFs folder, HTML docs to HTMLs folder.131 This is required if there is no extisting vector database132 created and wanted to build from the scratch133 '''134 for root, _, files in os.walk(folder_path):135 for file in files:136 _, extension = os.path.splitext(file)137 source_path = os.path.join(root, file)138 139 if extension == '.pdf':140 dest_folder = "PDFs"141 elif extension == '.html':142 dest_folder = "HTMLs"143 else:144 continue145 146 dest_path = os.path.join(dest_folder, file)147 os.makedirs(dest_folder, exist_ok=True)148 shutil.copy(source_path, dest_path)149 150 151 152def load_vectorstore(persist_directory, embeddings):153 '''154 This function will try first to load chroma database from the disk. If it does exist,155 It will do the following,156 1) Load the pdfs157 2) create text chunks158 3) Index it and store it in a Chroma DB159 4) Peform the same for HTML files160 5) Store the final chroma db in the disk.161 This is required if there is no extisting vector database162 created and wanted to build from the scratch163 '''164 if os.path.exists(persist_directory):165 print("Using existing vectore store for these documents.")166 vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)167 print("Chroma DB loaded from the disk")168 return vectorstore169 else:170 folder_path= unzip_opm()171 print("Vector store is not available. Creating new one.")172 file_count_by_type = count_files_by_type(folder_path)173 file_count_table = generate_file_count_table(file_count_by_type)174 print("File Count Table:")175 print(file_count_table)176 #move files into respective folders177 move_files_to_folders(folder_path)178 print("PDF and HTML files copied to separate folders.")179 180 # Load the pdf files from the pdffolder in order to create new chroma db181 pdf_folder_path= f"{folder_path}/PDFs" #pdf folder182 html_folder_path= f"{folder_path}/HTMLs" #html folder183 pdf_dir_loader = PyPDFDirectoryLoader(pdf_folder_path)184 pdf_pages = pdf_dir_loader.load()185 print("PDF files are loaded from the folder.")186 187 188 #Loading HTML files from the html folder in order to create new chroma db 189 HTML_docs_path_list = [os.path.join(html_folder_path, f) for f in os.listdir(html_folder_path) if os.path.isfile(os.path.join(html_folder_path, f))]190 191 html_loaders= []192 for html_file in HTML_docs_path_list:193 loader = UnstructuredHTMLLoader(html_file)194 html_loaders.append(loader)195 196 html_pages = []197 docs_cannot_load= []198 for loader in html_loaders:199 try:200 html_pages.extend(loader.load())201 except:202 print("Cannot load the file:", loader)203 docs_cannot_load.append(loader)204 print("HTML files are loaded from the folder.")205 # Create text chunks from the PDF docs206 text_splitter = RecursiveCharacterTextSplitter(207 # Set a really small chunk size, just to show.208 chunk_size = 1000,209 chunk_overlap = 200,210 length_function = len,211 is_separator_regex = False,212 )213 214 pdf_texts = text_splitter.transform_documents(pdf_pages)215 # Create text chunks from the HTML docs216 html_texts = text_splitter.transform_documents(html_pages)217 # Merging all the text chunks (HTML + PDF)218 all_texts= pdf_texts+html_texts219 print("PDF and HTML docs are split into chunks and created a final list representing all the chunks.")220 221 # Create embeddings for all the text chunks and store it in a Chroma DB222 vectorstore = Chroma.from_documents(all_texts,223 embeddings,224 persist_directory=persist_directory)225 vectorstore.persist()226 print("Chroma DB created and loaded")227 return vectorstore228 229 230def load_text_chunks(text_chunks_pkl_dir):231 '''232 We need to get all the text chunks as it is required for bm25 retriever incase we are using it for creating enemble retriever 233 Loading the pickle file that holds all the documents from the disk.234 If it does not exist, create new one.235 Text documents are required to create BM25 Retriever. But loading all the documents in236 every session will be a time consuming process. So we are storing all the docs in a pickle file237 and load the pickle file from the disk to overcome this problem.238 '''239 try:240 print("Text chunks are loading from the disk")241 with open(text_chunks_pkl_dir, 'rb') as file:242 cached_text_chunks = pickle.load(file)243 # Now, `cached_text_chunks` contains your cached data244 print("Text chunks are loaded from the disk")245 return cached_text_chunks246 except:247 print("Creating text chunks from the docs and caching it.")248 folder_path= unzip_opm()249 pdf_folder_path= f"{folder_path}/PDFs" #pdf folder250 html_folder_path= f"{folder_path}/HTMLs" #html folder251 pdf_dir_loader = PyPDFDirectoryLoader(pdf_folder_path)252 pdf_pages = pdf_dir_loader.load()253 HTML_docs_path_list = [os.path.join(html_folder_path, f) for f in os.listdir(html_folder_path) if os.path.isfile(os.path.join(html_folder_path, f))]254 255 html_loaders= []256 for html_file in HTML_docs_path_list:257 loader = UnstructuredHTMLLoader(html_file)258 html_loaders.append(loader)259 260 html_pages = []261 for loader in html_loaders:262 try:263 html_pages.extend(loader.load())264 except:265 print("Cannot load the file:", loader)266 all_texts= pdf_pages+html_pages267 # Cache the list to a file268 with open('text_chunks.pkl', 'wb') as file:269 pickle.dump(all_texts, file)270 print("Text chunks are created and cached")271 272def load_retriver(chroma_vectorstore):273 """Load cohere rerank method for retrieval"""274 # bm25_retriever = BM25Retriever.from_documents(text_chunks)275 # bm25_retriever.k = 2276 chroma_retriever = chroma_vectorstore.as_retriever(search_kwargs={"k": 3}) 277 # ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, chroma_retriever], weights=[0.3, 0.7])278 logging.basicConfig()279 logging.getLogger('langchain.retrievers.multi_query').setLevel(logging.INFO)280 multi_query_retriever = MultiQueryRetriever.from_llm(retriever=chroma_retriever,281 llm=ChatOpenAI(temperature=0))282 compressor = CohereRerank()283 compression_retriever = ContextualCompressionRetriever(284 base_compressor=compressor,285 base_retriever=multi_query_retriever)286 return compression_retriever287 288 289def load_retriver(chroma_vectorstore):290 """Load cohere rerank method for retrieval"""291 # bm25_retriever = BM25Retriever.from_documents(text_chunks)292 # bm25_retriever.k = 2293 chroma_retriever = chroma_vectorstore.as_retriever(search_kwargs={"k": 3}) 294 # ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, chroma_retriever], weights=[0.3, 0.7])295 logging.basicConfig()296 logging.getLogger('langchain.retrievers.multi_query').setLevel(logging.INFO)297 multi_query_retriever = MultiQueryRetriever.from_llm(retriever=chroma_retriever,298 llm=ChatOpenAI(temperature=0))299 compressor = CohereRerank()300 compression_retriever = ContextualCompressionRetriever(301 base_compressor=compressor,302 base_retriever=multi_query_retriever)303 return compression_retriever304 305 306def load_conversational_retrievel_chain(retriever, llm):307 '''308 Create RetrievalQA chain with memory309 '''310 # template = """You are a helpful assistant. You do not respond as 'User' or pretend to be 'User'. You only respond once as 'Assistant'.311 # Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.312 # Only include information found in the results and don't add any additional information.313 # Make sure the answer is correct and don't output false content.314 # If the text does not relate to the query, simply state 'Text Not Found in the Document'. Ignore outlier,315 # search results which has nothing to do with the question. Only answer what is asked.316 # The answer should be short and concise. Answer step-by-step.317 318 # {context}319 320 # {history}321 # Question: {question}322 # Helpful Answer:"""323 324 # prompt = PromptTemplate(input_variables=["history", "context", "question"], template=template)325 # memory = ConversationBufferMemory(input_key="question", memory_key="history")326 327 qa = RetrievalQA.from_chain_type(328 llm=llm,329 chain_type="stuff",330 retriever=retriever,331 return_source_documents=True,332 # chain_type_kwargs={"memory": memory},333 )334 return qa335 336