import gradio as gr import os from langchain_community.vectorstores import FAISS from langchain_community.document_loaders import PyPDFLoader from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.chains import ConversationalRetrievalChain from langchain_community.embeddings import HuggingFaceEmbeddings from langchain_community.llms import HuggingFacePipeline from langchain.memory import ConversationBufferMemory import torch from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline # List of local models (publicly accessible, no token required) list_llm = ["facebook/opt-350m", "distilbert/distilgpt2"] list_llm_simple = [os.path.basename(llm) for llm in list_llm] # Load and split PDF document def load_doc(list_file_path): loaders = [PyPDFLoader(x) for x in list_file_path] pages = [] for loader in loaders: pages.extend(loader.load()) text_splitter = RecursiveCharacterTextSplitter( chunk_size=1024, chunk_overlap=64 ) doc_splits = text_splitter.split_documents(pages) return doc_splits # Create vector database def create_db(splits): embeddings = HuggingFaceEmbeddings() vectordb = FAISS.from_documents(splits, embeddings) return vectordb # Initialize langchain LLM chain with local model def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()): # Load the model and tokenizer locally tokenizer = AutoTokenizer.from_pretrained(llm_model) model = AutoModelForCausalLM.from_pretrained( llm_model, device_map="auto", # Automatically use GPU if available torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, # Optimize for GPU or CPU trust_remote_code=True # Required for some models ) # Create a pipeline for text generation pipe = pipeline( "text-generation", model=model, tokenizer=tokenizer, max_new_tokens=max_tokens, temperature=temperature, top_k=top_k, do_sample=True, repetition_penalty=1.1, return_full_text=False ) # Wrap the pipeline in HuggingFacePipeline for LangChain llm = HuggingFacePipeline(pipeline=pipe) memory = ConversationBufferMemory( memory_key="chat_history", output_key="answer", return_messages=True ) retriever = vector_db.as_retriever() qa_chain = ConversationalRetrievalChain.from_llm( llm, retriever=retriever, chain_type="stuff", memory=memory, return_source_documents=True, verbose=False, ) return qa_chain # Initialize database def initialize_database(list_file_obj, progress=gr.Progress()): list_file_path = [x.name for x in list_file_obj if x is not None] doc_splits = load_doc(list_file_path) vector_db = create_db(doc_splits) return vector_db, "Database created!" # Initialize LLM def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()): 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, progress) return qa_chain, "QA chain initialized. Chatbot is ready!" 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): formatted_chat_history = format_chat_history(message, history) response = qa_chain.invoke({"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] response_sources = response["source_documents"] response_source1 = response_sources[0].page_content.strip() response_source2 = response_sources[1].page_content.strip() response_source3 = response_sources[2].page_content.strip() response_source1_page = response_sources[0].metadata["page"] + 1 response_source2_page = response_sources[1].metadata["page"] + 1 response_source3_page = response_sources[2].metadata["page"] + 1 new_history = history + [(message, response_answer)] return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page 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) return list_file_path def demo(): with gr.Blocks(theme=gr.themes.Default(primary_hue="red", secondary_hue="pink", neutral_hue="sky")) as demo: vector_db = gr.State() qa_chain = gr.State() gr.HTML("

RAG PDF chatbot

") gr.Markdown("""Query your PDF documents! This AI agent is designed to perform retrieval augmented generation (RAG) on PDF documents. This version runs locally and does not require an API token. \ Please do not upload confidential documents. """) with gr.Row(): with gr.Column(scale=86): gr.Markdown("Step 1 - Upload PDF documents and Initialize RAG pipeline") with gr.Row(): document = gr.Files(height=300, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload PDF documents") with gr.Row(): db_btn = gr.Button("Create vector database") with gr.Row(): db_progress = gr.Textbox(value="Not initialized", show_label=False) gr.Markdown("Select Large Language Model (LLM) and input parameters") with gr.Row(): llm_btn = gr.Radio(list_llm_simple, label="Available LLMs", value=list_llm_simple[0], type="index") with gr.Row(): with gr.Accordion("LLM input parameters", open=False): with gr.Row(): slider_temperature = gr.Slider(minimum=0.01, maximum=1.0, value=0.5, step=0.1, label="Temperature", info="Controls randomness in token generation", interactive=True) with gr.Row(): slider_maxtokens = gr.Slider(minimum=128, maximum=2048, value=512, step=128, label="Max New Tokens", info="Maximum number of tokens to be generated", interactive=True) with gr.Row(): slider_topk = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="top-k", info="Number of tokens to select the next token from", interactive=True) with gr.Row(): qachain_btn = gr.Button("Initialize Question Answering Chatbot") with gr.Row(): llm_progress = gr.Textbox(value="Not initialized", show_label=False) with gr.Column(scale=200): gr.Markdown("Step 2 - Chat with your Document") chatbot = gr.Chatbot(height=505) with gr.Accordion("Relevant context from the source document", open=False): with gr.Row(): doc_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20) source1_page = gr.Number(label="Page", scale=1) with gr.Row(): doc_source2 = gr.Textbox(label="Reference 2", lines=2, container=True, scale=20) source2_page = gr.Number(label="Page", scale=1) with gr.Row(): doc_source3 = gr.Textbox(label="Reference 3", lines=2, container=True, scale=20) source3_page = gr.Number(label="Page", scale=1) with gr.Row(): msg = gr.Textbox(placeholder="Ask a question", container=True) with gr.Row(): submit_btn = gr.Button("Submit") clear_btn = gr.ClearButton([msg, chatbot], value="Clear") # Preprocessing events db_btn.click(initialize_database, inputs=[document], outputs=[vector_db, db_progress]) qachain_btn.click(initialize_LLM, inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], outputs=[qa_chain, llm_progress]).then( lambda: [None, "", 0, "", 0, "", 0], inputs=None, outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], queue=False ) # Chatbot events msg.submit(conversation, inputs=[qa_chain, msg, chatbot], outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], queue=False) submit_btn.click(conversation, inputs=[qa_chain, msg, chatbot], outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], queue=False) clear_btn.click(lambda: [None, "", 0, "", 0, "", 0], inputs=None, outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], queue=False) demo.queue().launch(debug=True) if __name__ == "__main__": demo()