Jeff28/CipherReadPDF
1
1import gradio as gr2import os3 4from langchain_community.document_loaders import PyPDFLoader5from langchain.text_splitter import RecursiveCharacterTextSplitter6from langchain_community.vectorstores import Chroma7from langchain.chains import ConversationalRetrievalChain8from langchain_community.embeddings import HuggingFaceEmbeddings 9from langchain_community.llms import HuggingFacePipeline10from langchain.chains import ConversationChain11from langchain.memory import ConversationBufferMemory12from langchain_community.llms import HuggingFaceEndpoint13 14from pathlib import Path15import chromadb16from unidecode import unidecode17 18from transformers import AutoTokenizer19import transformers20import torch21import tqdm 22import accelerate23import re24 25 26 27# default_persist_directory = './chroma_HF/'28list_llm = ["mistralai/Mistral-7B-Instruct-v0.2", "mistralai/Mixtral-8x7B-Instruct-v0.1", "mistralai/Mistral-7B-Instruct-v0.1", \29 "google/gemma-7b-it","google/gemma-2b-it", \30 "HuggingFaceH4/zephyr-7b-beta", "HuggingFaceH4/zephyr-7b-gemma-v0.1", \31 "meta-llama/Llama-2-7b-chat-hf", "microsoft/phi-2", \32 "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "mosaicml/mpt-7b-instruct", "tiiuae/falcon-7b-instruct", \33 "google/flan-t5-xxl"34]35list_llm_simple = [os.path.basename(llm) for llm in list_llm]36 37# Load PDF document and create doc splits38def load_doc(list_file_path, chunk_size, chunk_overlap):39 # Processing for one document only40 #loader = PyPDFLoader(file_path)41 #pages = loader.load()42 loaders = [PyPDFLoader(x) for x in list_file_path]43 pages = []44 for loader in loaders:45 pages.extend(loader.load())46 # text_splitter = RecursiveCharacterTextSplitter(chunk_size = 600, chunk_overlap = 50)47 text_splitter = RecursiveCharacterTextSplitter(48 chunk_size = chunk_size, 49 chunk_overlap = chunk_overlap)50 doc_splits = text_splitter.split_documents(pages)51 return doc_splits52 53 54# Create vector database55def create_db(splits, collection_name):56 embedding = HuggingFaceEmbeddings()57 new_client = chromadb.EphemeralClient()58 vectordb = Chroma.from_documents(59 documents=splits,60 embedding=embedding,61 client=new_client,62 collection_name=collection_name,63 # persist_directory=default_persist_directory64 )65 return vectordb66 67 68# Load vector database69def load_db():70 embedding = HuggingFaceEmbeddings()71 vectordb = Chroma(72 # persist_directory=default_persist_directory, 73 embedding_function=embedding)74 return vectordb75 76 77# Initialize langchain LLM chain78def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):79 progress(0.1, desc="Initializing HF tokenizer...")80 # HuggingFacePipeline uses local model81 # Note: it will download model locally...82 # tokenizer=AutoTokenizer.from_pretrained(llm_model)83 # progress(0.5, desc="Initializing HF pipeline...")84 # pipeline=transformers.pipeline(85 # "text-generation",86 # model=llm_model,87 # tokenizer=tokenizer,88 # torch_dtype=torch.bfloat16,89 # trust_remote_code=True,90 # device_map="auto",91 # # max_length=1024,92 # max_new_tokens=max_tokens,93 # do_sample=True,94 # top_k=top_k,95 # num_return_sequences=1,96 # eos_token_id=tokenizer.eos_token_id97 # )98 # llm = HuggingFacePipeline(pipeline=pipeline, model_kwargs={'temperature': temperature})99 100 # HuggingFaceHub uses HF inference endpoints101 progress(0.5, desc="Initializing HF Hub...")102 # Use of trust_remote_code as model_kwargs103 # Warning: langchain issue104 # URL: https://github.com/langchain-ai/langchain/issues/6080105 if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1":106 llm = HuggingFaceEndpoint(107 repo_id=llm_model, 108 # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "load_in_8bit": True}109 temperature = temperature,110 max_new_tokens = max_tokens,111 top_k = top_k,112 load_in_8bit = True,113 )114 elif llm_model in ["HuggingFaceH4/zephyr-7b-gemma-v0.1","mosaicml/mpt-7b-instruct"]:115 raise gr.Error("LLM model is too large to be loaded automatically on free inference endpoint")116 llm = HuggingFaceEndpoint(117 repo_id=llm_model, 118 temperature = temperature,119 max_new_tokens = max_tokens,120 top_k = top_k,121 )122 elif llm_model == "microsoft/phi-2":123 # raise gr.Error("phi-2 model requires 'trust_remote_code=True', currently not supported by langchain HuggingFaceHub...")124 llm = HuggingFaceEndpoint(125 repo_id=llm_model, 126 # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "trust_remote_code": True, "torch_dtype": "auto"}127 temperature = temperature,128 max_new_tokens = max_tokens,129 top_k = top_k,130 trust_remote_code = True,131 torch_dtype = "auto",132 )133 elif llm_model == "TinyLlama/TinyLlama-1.1B-Chat-v1.0":134 llm = HuggingFaceEndpoint(135 repo_id=llm_model, 136 # model_kwargs={"temperature": temperature, "max_new_tokens": 250, "top_k": top_k}137 temperature = temperature,138 max_new_tokens = 250,139 top_k = top_k,140 )141 elif llm_model == "meta-llama/Llama-2-7b-chat-hf":142 raise gr.Error("Llama-2-7b-chat-hf model requires a Pro subscription...")143 llm = HuggingFaceEndpoint(144 repo_id=llm_model, 145 # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k}146 temperature = temperature,147 max_new_tokens = max_tokens,148 top_k = top_k,149 )150 else:151 llm = HuggingFaceEndpoint(152 repo_id=llm_model, 153 # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "trust_remote_code": True, "torch_dtype": "auto"}154 # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k}155 temperature = temperature,156 max_new_tokens = max_tokens,157 top_k = top_k,158 )159 160 progress(0.75, desc="Defining buffer memory...")161 memory = ConversationBufferMemory(162 memory_key="chat_history",163 output_key='answer',164 return_messages=True165 )166 # retriever=vector_db.as_retriever(search_type="similarity", search_kwargs={'k': 3})167 retriever=vector_db.as_retriever()168 progress(0.8, desc="Defining retrieval chain...")169 qa_chain = ConversationalRetrievalChain.from_llm(170 llm,171 retriever=retriever,172 chain_type="stuff", 173 memory=memory,174 # combine_docs_chain_kwargs={"prompt": your_prompt})175 return_source_documents=True,176 #return_generated_question=False,177 verbose=False,178 )179 progress(0.9, desc="Done!")180 return qa_chain181 182 183# Generate collection name for vector database184# - Use filepath as input, ensuring unicode text185def create_collection_name(filepath):186 # Extract filename without extension187 collection_name = Path(filepath).stem188 # Fix potential issues from naming convention189 ## Remove space190 collection_name = collection_name.replace(" ","-") 191 ## ASCII transliterations of Unicode text192 collection_name = unidecode(collection_name)193 ## Remove special characters194 #collection_name = re.findall("[\dA-Za-z]*", collection_name)[0]195 collection_name = re.sub('[^A-Za-z0-9]+', '-', collection_name)196 ## Limit length to 50 characters197 collection_name = collection_name[:50]198 ## Minimum length of 3 characters199 if len(collection_name) < 3:200 collection_name = collection_name + 'xyz'201 ## Enforce start and end as alphanumeric character202 if not collection_name[0].isalnum():203 collection_name = 'A' + collection_name[1:]204 if not collection_name[-1].isalnum():205 collection_name = collection_name[:-1] + 'Z'206 print('Filepath: ', filepath)207 print('Collection name: ', collection_name)208 return collection_name209 210 211# Initialize database212def initialize_database(list_file_obj, chunk_size, chunk_overlap, progress=gr.Progress()):213 # Create list of documents (when valid)214 list_file_path = [x.name for x in list_file_obj if x is not None]215 # Create collection_name for vector database216 progress(0.1, desc="Creating collection name...")217 collection_name = create_collection_name(list_file_path[0])218 progress(0.25, desc="Loading document...")219 # Load document and create splits220 doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)221 # Create or load vector database222 progress(0.5, desc="Generating vector database...")223 # global vector_db224 vector_db = create_db(doc_splits, collection_name)225 progress(0.9, desc="Done!")226 return vector_db, collection_name, "Complete!"227 228 229def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):230 # print("llm_option",llm_option)231 llm_name = list_llm[llm_option]232 print("llm_name: ",llm_name)233 qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, progress)234 return qa_chain, "Complete!"235 236 237def format_chat_history(message, chat_history):238 formatted_chat_history = []239 for user_message, bot_message in chat_history:240 formatted_chat_history.append(f"User: {user_message}")241 formatted_chat_history.append(f"Assistant: {bot_message}")242 return formatted_chat_history243 244 245def conversation(qa_chain, message, history):246 formatted_chat_history = format_chat_history(message, history)247 #print("formatted_chat_history",formatted_chat_history)248 249 # Generate response using QA chain250 response = qa_chain({"question": message, "chat_history": formatted_chat_history})251 response_answer = response["answer"]252 if response_answer.find("Helpful Answer:") != -1:253 response_answer = response_answer.split("Helpful Answer:")[-1]254 response_sources = response["source_documents"]255 response_source1 = response_sources[0].page_content.strip()256 response_source2 = response_sources[1].page_content.strip()257 response_source3 = response_sources[2].page_content.strip()258 # Langchain sources are zero-based259 response_source1_page = response_sources[0].metadata["page"] + 1260 response_source2_page = response_sources[1].metadata["page"] + 1261 response_source3_page = response_sources[2].metadata["page"] + 1262 # print ('chat response: ', response_answer)263 # print('DB source', response_sources)264 265 # Append user message and response to chat history266 new_history = history + [(message, response_answer)]267 # return gr.update(value=""), new_history, response_sources[0], response_sources[1] 268 return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page269 270 271def upload_file(file_obj):272 list_file_path = []273 for idx, file in enumerate(file_obj):274 file_path = file_obj.name275 list_file_path.append(file_path)276 # print(file_path)277 # initialize_database(file_path, progress)278 return list_file_path279 280 281def demo():282 with gr.Blocks(theme="base") as demo:283 vector_db = gr.State()284 qa_chain = gr.State()285 collection_name = gr.State()286 287 gr.Markdown(288 """<center><h2>PDF-based chatbot</center></h2>289 <h3>Ask any questions about your PDF documents</h3>""")290 gr.Markdown(291 """<b>Note:</b> CypherReadPDF is an advanced tool for exploring and analyzing PDF documents.\292 Using AI models like LLaMA, it allows users to upload PDFs, extract key information, and ask interactive questions about the content. 293 The user interface guides users through the analysis process, while conversational memory ensures answers consider past questions, with document references for clarity.<br>294 <br><b>Warning:</b> This space uses the free CPU Basic hardware from Hugging Face. Some steps and LLM models used below (free inference endpoints) can take some time to generate a reply.295 """)296 297 with gr.Tab("Step 1 - Upload PDF"):298 with gr.Row():299 document = gr.Files(height=100, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload your PDF documents (single or multiple)")300 # upload_btn = gr.UploadButton("Loading document...", height=100, file_count="multiple", file_types=["pdf"], scale=1)301 302 with gr.Tab("Step 2 - Process document"):303 with gr.Row():304 db_btn = gr.Radio(["ChromaDB"], label="Vector database type", value = "ChromaDB", type="index", info="Choose your vector database")305 with gr.Accordion("Advanced options - Document text splitter", open=False):306 with gr.Row():307 slider_chunk_size = gr.Slider(minimum = 100, maximum = 1000, value=600, step=20, label="Chunk size", info="Chunk size", interactive=True)308 with gr.Row():309 slider_chunk_overlap = gr.Slider(minimum = 10, maximum = 200, value=40, step=10, label="Chunk overlap", info="Chunk overlap", interactive=True)310 with gr.Row():311 db_progress = gr.Textbox(label="Vector database initialization", value="None")312 with gr.Row():313 db_btn = gr.Button("Generate vector database")314 315 with gr.Tab("Step 3 - Initialize QA chain"):316 with gr.Row():317 llm_btn = gr.Radio(list_llm_simple, \318 label="LLM models", value = list_llm_simple[0], type="index", info="Choose your LLM model")319 with gr.Accordion("Advanced options - LLM model", open=False):320 with gr.Row():321 slider_temperature = gr.Slider(minimum = 0.01, maximum = 1.0, value=0.7, step=0.1, label="Temperature", info="Model temperature", interactive=True)322 with gr.Row():323 slider_maxtokens = gr.Slider(minimum = 224, maximum = 4096, value=1024, step=32, label="Max Tokens", info="Model max tokens", interactive=True)324 with gr.Row():325 slider_topk = gr.Slider(minimum = 1, maximum = 10, value=3, step=1, label="top-k samples", info="Model top-k samples", interactive=True)326 with gr.Row():327 llm_progress = gr.Textbox(value="None",label="QA chain initialization")328 with gr.Row():329 qachain_btn = gr.Button("Initialize Question Answering chain")330 331 with gr.Tab("Step 4 - Chatbot"):332 chatbot = gr.Chatbot(height=300)333 with gr.Accordion("Advanced - Document references", open=False):334 with gr.Row():335 doc_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20)336 source1_page = gr.Number(label="Page", scale=1)337 with gr.Row():338 doc_source2 = gr.Textbox(label="Reference 2", lines=2, container=True, scale=20)339 source2_page = gr.Number(label="Page", scale=1)340 with gr.Row():341 doc_source3 = gr.Textbox(label="Reference 3", lines=2, container=True, scale=20)342 source3_page = gr.Number(label="Page", scale=1)343 with gr.Row():344 msg = gr.Textbox(placeholder="Type message (e.g. 'What is this document about?')", container=True)345 with gr.Row():346 submit_btn = gr.Button("Submit message")347 clear_btn = gr.ClearButton([msg, chatbot], value="Clear conversation")348 349 # Preprocessing events350 #upload_btn.upload(upload_file, inputs=[upload_btn], outputs=[document])351 db_btn.click(initialize_database, \352 inputs=[document, slider_chunk_size, slider_chunk_overlap], \353 outputs=[vector_db, collection_name, db_progress])354 qachain_btn.click(initialize_LLM, \355 inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], \356 outputs=[qa_chain, llm_progress]).then(lambda:[None,"",0,"",0,"",0], \357 inputs=None, \358 outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \359 queue=False)360 361 # Chatbot events362 msg.submit(conversation, \363 inputs=[qa_chain, msg, chatbot], \364 outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \365 queue=False)366 submit_btn.click(conversation, \367 inputs=[qa_chain, msg, chatbot], \368 outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \369 queue=False)370 clear_btn.click(lambda:[None,"",0,"",0,"",0], \371 inputs=None, \372 outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \373 queue=False)374 demo.queue().launch(debug=True)375 376 377if __name__ == "__main__":378 demo()