Update app.py

app.py

@@ -1,98 +1,87 @@
 import gradio as gr
 from PyPDF2 import PdfReader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.chains.question_answering import load_qa_chain
 from langchain_community.llms import HuggingFacePipeline
-from langchain.chains import ConversationalRetrievalChain
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# -------------------------------
-# PDF TEXT LOADER
-# -------------------------------
-def load_pdfs(pdf_files):
+
+# ----------------------------
+# Lazy load model & embeddings
+# ----------------------------
+def load_llm():
+    model_id = "google/flan-t5-base"  # lightweight model for Hugging Face Spaces
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+    pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+    llm = HuggingFacePipeline(pipeline=pipe)
+    return llm
+
+
+def load_embeddings():
+    return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+
+# ----------------------------
+# Process PDF
+# ----------------------------
+def process_pdf(pdf_file):
+    pdf_reader = PdfReader(pdf_file.name)
     text = ""
-    for pdf in pdf_files:
-        pdf_reader = PdfReader(pdf.name)   # use .name for gradio file objects
-        for page in pdf_reader.pages:
-            page_text = page.extract_text()
-            if page_text:
-                text += page_text
-    return text
-
-# -------------------------------
-# BUILD VECTORSTORE
-# -------------------------------
-def build_vectorstore(text):
-    splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,
-        chunk_overlap=200,
-        length_function=len
-    )
+    for page in pdf_reader.pages:
+        text += page.extract_text() or ""
+
+    if not text.strip():
+        return None, "❌ No extractable text found in PDF!"
+
+    # Split text
+    splitter = CharacterTextSplitter(separator="\n", chunk_size=1000, chunk_overlap=200, length_function=len)
     chunks = splitter.split_text(text)
 
-    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-    vectorstore = FAISS.from_texts(chunks, embeddings)
-    return vectorstore
-
-# -------------------------------
-# SETUP QA CHAIN
-# -------------------------------
-def build_conversation_chain(vectorstore):
-    llm_pipeline = pipeline(
-        "text2text-generation",
-        model="google/flan-t5-base",  # lightweight, fast model
-        tokenizer="google/flan-t5-base",
-        max_new_tokens=256
-    )
-    llm = HuggingFacePipeline(pipeline=llm_pipeline)
-
-    qa_chain = ConversationalRetrievalChain.from_llm(
-        llm=llm,
-        retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
-        return_source_documents=False
-    )
-    return qa_chain
-
-# -------------------------------
-# GRADIO INTERFACE
-# -------------------------------
-conversation_chain = None
-chat_history = []
-
-def process_pdfs(pdf_files):
-    global conversation_chain, chat_history
-    chat_history = []  # reset history
-    text = load_pdfs(pdf_files)
-    vs = build_vectorstore(text)
-    conversation_chain = build_conversation_chain(vs)
-    return "✅ PDFs processed successfully. You can now ask questions!"
-
-def chat(message, history):
-    global conversation_chain, chat_history
-    if not conversation_chain:
-        return "⚠️ Please upload and process PDFs first."
-    
-    response = conversation_chain({"question": message, "chat_history": chat_history})
-    answer = response["answer"]
-    chat_history.append((message, answer))
-    return answer
+    # Create embeddings + FAISS index
+    embeddings = load_embeddings()
+    knowledge_base = FAISS.from_texts(chunks, embeddings)
+
+    return knowledge_base, "✅ PDF processed successfully!"
 
 
+# ----------------------------
+# Chat Function
+# ----------------------------
+def chat_with_pdf(pdf_file, query, history=[]):
+    if pdf_file is None:
+        return history + [["User: " + query, "⚠️ Please upload a PDF first!"]]
+
+    # Process PDF
+    knowledge_base, msg = process_pdf(pdf_file)
+    if knowledge_base is None:
+        return history + [["System", msg]]
+
+    # Run LLM QA Chain
+    llm = load_llm()
+    chain = load_qa_chain(llm, chain_type="stuff")
+
+    docs = knowledge_base.similarity_search(query, k=3)
+    answer = chain.run(input_documents=docs, question=query)
+
+    history.append(["User: " + query, "Bot: " + answer])
+    return history
+
+
+# ----------------------------
+# Gradio UI
+# ----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 📚 Multiple PDF Chatbot")
-    
-    with gr.Row():
-        pdf_input = gr.File(file_types=[".pdf"], file_types_display="PDF Files", file_types_visible=True, file_types_select_multiple=True, label="Upload PDFs", type="file", file_types_accept_multiple=True)
-        process_btn = gr.Button("Process PDFs")
-    
-    output_status = gr.Textbox(label="Status", interactive=False)
+    gr.Markdown("## 📄 Multiple PDF Chatbot (LangChain + Hugging Face)")
 
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Your Question")
-    send_btn = gr.Button("Ask")
+    with gr.Row():
+        pdf_file = gr.File(label="Upload PDF", file_types=[".pdf"])
+        query = gr.Textbox(label="Ask a question about the PDF")
+    chatbot = gr.Chatbot(label="Conversation")
 
-    process_btn.click(process_pdfs, inputs=pdf_input, outputs=output_status)
-    send_btn.click(chat, inputs=[msg, chatbot], outputs=chatbot)
+    btn = gr.Button("Ask")
+    btn.click(fn=chat_with_pdf, inputs=[pdf_file, query, chatbot], outputs=chatbot)
 
 demo.launch()