Update app.py

app.py

@@ -2,84 +2,67 @@ import gradio as gr
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
-from langchain_community.llms import HuggingFaceEndpoint
-from langchain.chains import RetrievalQA
+from langchain_community.chat_models import ChatGroq
+from langchain.chains import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
 import os
 import shutil
 
-# Hugging Face API key (store in your Space's secrets for security)
-HUGGINGFACEHUB_API_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+# PDF upload folder
+UPLOAD_DIR = "pdf_uploads"
+if not os.path.exists(UPLOAD_DIR):
+    os.makedirs(UPLOAD_DIR)
 
-# Load HF Inference Endpoint (like mistralai/Mistral-7B-Instruct)
-llm = HuggingFaceEndpoint(
-    repo_id="mistralai/Mistral-7B-Instruct-v0.2",
-    temperature=0.2,
-    huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN
-)
+# Initialize LLaMA 3 (no API key directly in code)
+llm = ChatGroq(model_name="llama3-8b-8192")
 
-# Embeddings (Hugging Face miniLM for fast processing)
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
-# Temporary folder to store PDFs
-UPLOAD_DIR = "pdf_uploads"
-if not os.path.exists(UPLOAD_DIR):
-    os.makedirs(UPLOAD_DIR)
+qa_chain = None
 
-def process_pdf(file):
-    # Save PDF file
-    file_path = os.path.join(UPLOAD_DIR, file.name)
-    with open(file_path, "wb") as f:
-        f.write(file.read())
-    
-    # Load PDF text using langchain
+def process_pdf(file_path):
     loader = PyPDFLoader(file_path)
     pages = loader.load_and_split()
-    
-    # Create Chroma vector store (in-memory)
+
     vectordb = Chroma.from_documents(pages, embedding=embeddings)
-    retriever = vectordb.as_retriever()
-    
-    # Create RetrievalQA chain
-    qa_chain = RetrievalQA.from_chain_type(
+    qa_chain = ConversationalRetrievalChain.from_llm(
         llm=llm,
-        chain_type="stuff",
-        retriever=retriever,
-        return_source_documents=True
+        retriever=vectordb.as_retriever(),
+        memory=memory
     )
-    
-    # Return the QA chain to use in the chat
     return qa_chain
 
-# Global variable to hold QA chain for the session
-qa_chain = None
-
 def upload_pdf(file):
     global qa_chain
-    qa_chain = process_pdf(file)
+    file_path = os.path.join(UPLOAD_DIR, os.path.basename(file))
+    shutil.copy(file, file_path)
+    qa_chain = process_pdf(file_path)
     return "✅ PDF uploaded and processed! Ask me anything about it."
 
 def chatbot(user_message, history):
     if qa_chain is None:
         return "❌ Please upload a PDF first.", history
-    
-    response = qa_chain.run(user_message)
-    history.append((user_message, response))
+
+    response = qa_chain({"question": user_message, "chat_history": history})
+    answer = response["answer"]
+    history.append((user_message, answer))
     return "", history
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("<h1 style='text-align:center;'>📚 PDF Chatbot Assistant</h1>")
-    
+    gr.Markdown("<h1 style='text-align:center;'>📄 LLaMA 3 PDF Chatbot</h1>")
+
     with gr.Row():
-        pdf_upload = gr.File(label="Upload your PDF", file_types=[".pdf"])
+        pdf_upload = gr.File(label="Upload PDF", file_types=[".pdf"])
         upload_btn = gr.Button("Process PDF")
-    
+
     chatbot_ui = gr.Chatbot(height=400)
     user_input = gr.Textbox(label="Ask something about the PDF...", placeholder="Type your question here and hit Enter")
 
     upload_btn.click(upload_pdf, inputs=pdf_upload, outputs=chatbot_ui)
     user_input.submit(chatbot, [user_input, chatbot_ui], [user_input, chatbot_ui])
-    
-    gr.Markdown("<footer style='text-align:center; font-size:0.85rem; color:#64748b;'>Created by YourName - Powered by Hugging Face</footer>")
+
+    gr.Markdown("<footer style='text-align:center; font-size:0.85rem; color:#64748b;'>Built with LLaMA 3 + LangChain on Hugging Face Spaces 🚀</footer>")
 
 if __name__ == "__main__":
     demo.launch()