Upload 2 files

app.py

@@ -0,0 +1,144 @@
+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_community.embeddings import HuggingFaceEmbeddings
+from langchain.chains import ConversationalRetrievalChain
+from langchain_community.llms import HuggingFaceEndpoint
+from langchain.memory import ConversationBufferMemory
+
+api_token = os.getenv("HF_TOKEN")
+
+list_llm = ["meta-llama/Meta-Llama-3-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.2"]
+list_llm_simple = [os.path.basename(llm) for llm in list_llm]
+
+def load_and_process_docs(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
+    )
+    return text_splitter.split_documents(pages)
+
+def create_vector_db(splits):
+    embeddings = HuggingFaceEmbeddings()
+    return FAISS.from_documents(splits, embeddings)
+
+def initialize_qa_chain(llm_model, vector_db, temperature=0.5, max_tokens=4096, top_k=3):
+    llm = HuggingFaceEndpoint(
+        repo_id=llm_model,
+        huggingfacehub_api_token=api_token,
+        temperature=temperature,
+        max_new_tokens=max_tokens,
+        top_k=top_k,
+    )
+    
+    memory = ConversationBufferMemory(
+        memory_key="chat_history",
+        output_key='answer',
+        return_messages=True
+    )
+
+    return ConversationalRetrievalChain.from_llm(
+        llm,
+        retriever=vector_db.as_retriever(),
+        chain_type="stuff",
+        memory=memory,
+        return_source_documents=True,
+        verbose=False,
+    )
+
+def format_response_with_citations(response_text, sources):
+    formatted_response = response_text
+    for idx, source in enumerate(sources, 1):
+        citation_marker = f"[{idx}]"
+        formatted_response += f"\n\n{citation_marker} (Page {source.metadata['page'] + 1}): {source.page_content.strip()}"
+    return formatted_response
+
+def chat(qa_chain, message, history):
+    formatted_history = []
+    for user_msg, bot_msg in history:
+        formatted_history.append(f"User: {user_msg}")
+        formatted_history.append(f"Assistant: {bot_msg}")
+
+    response = qa_chain.invoke({
+        "question": message,
+        "chat_history": formatted_history
+    })
+    
+    answer = response["answer"]
+    if "Helpful Answer:" in answer:
+        answer = answer.split("Helpful Answer:")[-1]
+    
+    formatted_response = format_response_with_citations(
+        answer,
+        response["source_documents"][:3]
+    )
+    
+    return qa_chain, "", history + [(message, formatted_response)]
+
+def demo():
+    with gr.Blocks(theme=gr.themes.Default(primary_hue="red")) as demo:
+        qa_chain = gr.State()
+        
+        gr.HTML("<center><h1>RAG PDF Chatbot</h1></center>")
+        gr.Markdown("""Query your PDF documents with citation support. 
+        **Please do not upload confidential documents.**""")
+        
+        with gr.Row():
+            with gr.Column():
+                document = gr.Files(
+                    height=100,
+                    file_count="multiple",
+                    file_types=["pdf"],
+                    label="Upload PDF Documents"
+                )
+                llm_choice = gr.Radio(
+                    list_llm_simple,
+                    label="Select Language Model",
+                    value=list_llm_simple[0],
+                    type="index"
+                )
+                
+        with gr.Column():
+            chatbot = gr.Chatbot(height=500)
+            msg = gr.Textbox(
+                placeholder="Ask a question about your documents",
+                container=True
+            )
+            with gr.Row():
+                submit_btn = gr.Button("Submit")
+                clear_btn = gr.ClearButton([msg, chatbot])
+
+        def initialize_system(files, llm_idx):
+            if not files:
+                return None
+            file_paths = [f.name for f in files]
+            splits = load_and_process_docs(file_paths)
+            vector_db = create_vector_db(splits)
+            return initialize_qa_chain(list_llm[llm_idx], vector_db)
+
+        # Auto-initialize when files are uploaded and model is selected
+        document.change(
+            initialize_system,
+            inputs=[document, llm_choice],
+            outputs=[qa_chain]
+        )
+        llm_choice.change(
+            initialize_system,
+            inputs=[document, llm_choice],
+            outputs=[qa_chain]
+        )
+
+        # Chat interactions
+        msg.submit(chat, inputs=[qa_chain, msg, chatbot], outputs=[qa_chain, msg, chatbot])
+        submit_btn.click(chat, inputs=[qa_chain, msg, chatbot], outputs=[qa_chain, msg, chatbot])
+
+    return demo.queue()
+
+if __name__ == "__main__":
+    demo().launch(debug=True)

requirements.txt

@@ -0,0 +1,9 @@
+torch
+transformers
+sentence-transformers
+langchain
+langchain-community
+tqdm
+accelerate
+pypdf
+faiss-gpu