import gradio as gr import os from dotenv import load_dotenv from langchain_community.document_loaders import PyPDFLoader from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_community.vectorstores import Chroma from langchain.chains import ConversationalRetrievalChain from langchain_community.embeddings import HuggingFaceEmbeddings from langchain_community.llms import HuggingFacePipeline from langchain.chains import ConversationChain from langchain.memory import ConversationBufferMemory from langchain.llms import HuggingFaceHub # from doctr.models import ocr_predictor # from doctr.io import DocumentFile from pathlib import Path import chromadb # Later Packages from getpass import getpass import weasyprint import matplotlib.pyplot as plt from langchain.document_loaders import PyPDFDirectoryLoader load_dotenv() # model = ocr_predictor(pretrained = True) huggingfacehub_api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN") openai_key = os.getenv("OPEN_API_KEY") # default_persist_directory = './chroma_HF/' list_llm = ["mistralai/Mixtral-8x7B-Instruct-v0.1", "mistralai/Mistral-7B-Instruct-v0.1", \ "google/gemma-7b-it","google/gemma-2b-it", \ "HuggingFaceH4/zephyr-7b-beta", \ "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "tiiuae/falcon-7b-instruct", \ "google/flan-t5-xxl" ] list_llm_simple = [os.path.basename(llm) for llm in list_llm] #Extract text data from doctr reaponse def extract_value_from_response(response): value = '' for page in response.pages: for block in page.blocks: for line in block.lines: for word in line.words: value += " "+word.value return value # Craete PDf from URL def create_pdf_from_url(url): pdf = weasyprint.HTML(url).write_pdf() output_dir = "pdfDir" if not os.path.exists(output_dir): os.makedirs(output_dir) file_path = os.path.join(output_dir,'url_pdf.pdf') with open(file_path,'wb') as f: f.write(pdf) return file_path # Load PDF document and create doc splits def load_doc(list_file_path, chunk_size, chunk_overlap): # Processing for one document only loaders = [PyPDFLoader(x) for x in list_file_path] pages = [] for loader in loaders: pages.extend(loader.load()) text_splitter = RecursiveCharacterTextSplitter( chunk_size = chunk_size, chunk_overlap = chunk_overlap) doc_splits = text_splitter.split_documents(pages) # if len(doc_splits) == 0: # doc = DocumentFile.from_pdf(list_file_path[0]) # result = model(doc) # response = extract_value_from_response(result) # doc_splits = text_splitter.split_documents(response) return doc_splits # Create vector database def create_db(splits, collection_name): embedding = HuggingFaceEmbeddings() new_client = chromadb.EphemeralClient() vectordb = Chroma.from_documents( documents = splits, embedding = embedding, client = new_client, collection_name = collection_name, # persist_directory=default_persist_directory ) return vectordb # Load vector database def load_db(): embedding = HuggingFaceEmbeddings() vectordb = Chroma( embedding_function = embedding) return vectordb # Initialize langchain LLM chain def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()): progress(0.1, desc="Initializing HF tokenizer...") # HuggingFaceHub uses HF inference endpoints progress(0.5, desc="Initializing HF Hub...") # Use of trust_remote_code as model_kwargs # Warning: langchain issue # URL: https://github.com/langchain-ai/langchain/issues/6080 if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1": llm = HuggingFaceHub( repo_id=llm_model, model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "load_in_8bit": True} ) elif llm_model == "TinyLlama/TinyLlama-1.1B-Chat-v1.0": llm = HuggingFaceHub( repo_id=llm_model, model_kwargs={"temperature": temperature, "max_new_tokens": 250, "top_k": top_k} ) else: llm = HuggingFaceHub( repo_id=llm_model, model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k} ) progress(0.75, desc="Defining buffer memory...") memory = ConversationBufferMemory( memory_key="chat_history", output_key='answer', return_messages=True ) progress(0.8, desc="Defining retrieval chain...") retriever = vector_db.as_retriever() qa_chain = ConversationalRetrievalChain.from_llm( llm, retriever = retriever, chain_type = "stuff", memory = memory, # combine_docs_chain_kwargs={"prompt": your_prompt}) return_source_documents=True, #return_generated_question=False, verbose = False, ) progress(0.9, desc="Done!") return qa_chain # Initialize database def initialize_database(list_file_obj, chunk_size, chunk_overlap, vector_db, url, progress = gr.Progress()): if url != "": file_path = create_pdf_from_url(url) list_file_obj = [] list_file_obj.append(file_path) list_file_path = list_file_obj else: # Create list of documents (when valid) list_file_path = [x.name for x in list_file_obj if x is not None] # Create collection_name for vector database progress(0.1, desc="Creating collection name...") collection_name = Path(list_file_path[0]).stem # Fix potential issues from naming convention ## Remove space collection_name = collection_name.replace(" ","-") ## Limit lenght to 50 characters collection_name = collection_name[:50] ## Enforce start and end as alphanumeric character if not collection_name[0].isalnum(): collection_name[0] = 'A' if not collection_name[-1].isalnum(): collection_name[-1] = 'Z' # print('list_file_path: ', list_file_path) print('Collection name: ', collection_name) progress(0.25, desc="Loading document...") # Load document and create splits doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap) # Create or load vector database progress(0.7, desc="Generating vector database...") # global vector_db vector_db = create_db(doc_splits, collection_name) return vector_db, collection_name, gr.update(value = ""), "Complete!" def re_initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db): llm_name = list_llm[llm_option] print("llm_name: ",llm_name) qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db) return qa_chain def format_chat_history(message, chat_history): formatted_chat_history = [] for user_message, bot_message in chat_history: formatted_chat_history.append(f"User: {user_message}") formatted_chat_history.append(f"Assistant: {bot_message}") return formatted_chat_history def conversation(qa_chain, message, history, llm_option): formatted_chat_history = format_chat_history(message, history) # Generate response using QA chain response = qa_chain({"question": message, "chat_history": formatted_chat_history}) response_answer = response["answer"] if response_answer.find("Helpful Answer:") != -1: response_answer = response_answer.split("Helpful Answer:")[-1] new_history = history + [(message, response_answer)] return qa_chain, gr.update(value = ""), new_history def upload_file(file_obj): list_file_path = [] for idx, file in enumerate(file_obj): file_path = file_obj.name list_file_path.append(file_path) # print(file_path) return list_file_path def demo(): with gr.Blocks(theme = "base") as demo: vector_db = gr.State() qa_chain = gr.State() collection_name = gr.State() gr.Markdown( '''