Update app.py

app.py

@@ -1,53 +1,50 @@
 import gradio as gr
 import fitz  # PyMuPDF
 import os
+import re
+from io import BytesIO
+from collections import Counter
+
 import numpy as np
 import matplotlib.pyplot as plt
 from wordcloud import WordCloud
-from io import BytesIO
-from collections import Counter
-import re
 
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from sentence_transformers import SentenceTransformer
 from transformers import pipeline
 from sklearn.metrics.pairwise import cosine_similarity
 
-# Load models
+# Load models once
 embed_model = SentenceTransformer('all-MiniLM-L6-v2')
-qa_pipeline = pipeline("question-answering", model="deepset/tinyroberta-squad2")  # ✅ lightweight model
+qa_pipeline = pipeline("question-answering", model="deepset/tinyroberta-squad2")
 
-# Globals
+# Globals to hold data
 all_chunks = []
 chunk_sources = []
 chunk_embeddings = None
 combined_text = ""
 
 def extract_text_from_pdfs(pdf_files):
-    """Extract text and page info from uploaded PDFs"""
-    global all_chunks, chunk_sources, combined_text
+    global combined_text, chunk_sources
+    combined_text = ""
     texts = []
     chunk_sources = []
-    combined_text = ""
 
     for file in pdf_files:
         doc = fitz.open(file.name)
-        for page_num, page in enumerate(doc):
+        for i, page in enumerate(doc):
             text = page.get_text()
             if text.strip():
-                texts.append((text, f"{os.path.basename(file.name)} - Page {page_num + 1}"))
+                texts.append((text, f"{os.path.basename(file.name)} - Page {i+1}"))
                 combined_text += " " + text
-
     return texts
 
 def split_and_embed(texts_with_sources):
-    """Split text into chunks and compute embeddings"""
     global all_chunks, chunk_sources, chunk_embeddings
     all_chunks = []
     chunk_sources = []
 
     splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-
     for text, source in texts_with_sources:
         docs = splitter.create_documents([text])
         for doc in docs:
@@ -60,80 +57,67 @@ def split_and_embed(texts_with_sources):
         chunk_embeddings = None
 
 def generate_wordcloud():
-    """Generate a word cloud from combined PDF text"""
     global combined_text
     if not combined_text.strip():
         return None
 
     cleaned = re.sub(r"[^a-zA-Z\s]", "", combined_text.lower())
     word_freq = Counter(cleaned.split())
-
     wc = WordCloud(width=800, height=400, background_color="white").generate_from_frequencies(word_freq)
 
     fig, ax = plt.subplots()
     ax.imshow(wc, interpolation='bilinear')
     ax.axis("off")
 
-    buffer = BytesIO()
-    plt.savefig(buffer, format="png")
-    buffer.seek(0)
-    return buffer
+    buf = BytesIO()
+    plt.savefig(buf, format="png")
+    plt.close(fig)
+    buf.seek(0)
+    return buf
 
 def answer_question(question):
-    """Retrieve top chunks, answer question, and show confidence"""
     global all_chunks, chunk_sources, chunk_embeddings
-    print("📥 Question received:", question)
-
     if not all_chunks or chunk_embeddings is None:
-        print("⚠️ PDF not processed or empty.")
-        return "Please upload and process some PDFs first.", None
+        return "Please upload and index PDFs first."
 
     q_emb = embed_model.encode([question], convert_to_numpy=True)
     sims = cosine_similarity(q_emb, chunk_embeddings)[0]
     top_k_idx = sims.argsort()[::-1][:3]
 
-    selected_chunks = [all_chunks[i] for i in top_k_idx]
-    selected_sources = [chunk_sources[i] for i in top_k_idx]
-    context = "\n\n".join(selected_chunks)
-
+    context = "\n\n".join([all_chunks[i] for i in top_k_idx])
     if not context.strip():
-        print("⚠️ Empty context from chunks.")
-        return "Could not extract relevant content from the PDFs.", None
+        return "No relevant content found in PDFs."
 
     try:
-        answer_dict = qa_pipeline(question=question, context=context)
-        answer = answer_dict.get("answer", "No answer found.")
-    except Exception as e:
-        print("❌ Error from QA model:", e)
-        return "Model failed to generate an answer.", None
-
-    avg_conf = np.mean([sims[i] for i in top_k_idx]) * 100
-    source_info = "\n".join([f"- {src}" for src in selected_sources])
-    result = f"**Answer**: {answer}\n\n**Sources**:\n{source_info}\n\n**Confidence Score**: {avg_conf:.2f}%"
-    print("✅ Answer generated.")
-    return result, None
-
-# Gradio UI
+        result = qa_pipeline(question=question, context=context)
+        answer = result.get("answer", "No answer found.")
+    except Exception:
+        return "Error generating answer from the model."
+
+    sources = "\n".join(set(chunk_sources[i] for i in top_k_idx))
+    confidence = np.mean([sims[i] for i in top_k_idx]) * 100
+    return f"**Answer:** {answer}\n\n**Sources:**\n{sources}\n\n**Confidence:** {confidence:.2f}%"
+
 with gr.Blocks() as demo:
-    gr.Markdown("# 📚 Enhanced RAG PDF Chatbot")
-    gr.Markdown("Upload PDFs → Preview Keywords → Ask Questions → Get Answers with Confidence & Sources")
+    gr.Markdown("# PDF Chatbot")
+    gr.Markdown("Upload PDFs, extract text, then ask questions.")
 
     with gr.Row():
-        pdf_input = gr.File(file_types=[".pdf"], file_count="multiple", label="Upload PDFs")
-        load_button = gr.Button("Extract & Index")
-        cloud_output = gr.Image(label="Keyword Preview (Word Cloud)")
+        pdf_input = gr.File(file_types=[".pdf"], file_count="multiple")
+        extract_btn = gr.Button("Extract & Index")
+        wc_img = gr.Image(label="Word Cloud")
 
     with gr.Row():
-        question_input = gr.Textbox(lines=2, placeholder="Ask your question here...", label="Question")
-        ask_button = gr.Button("Get Answer")
-        answer_output = gr.Markdown()
+        question_input = gr.Textbox(lines=2, placeholder="Ask your question here...")
+        ask_btn = gr.Button("Get Answer")
+        answer_out = gr.Markdown()
 
-    def load_and_index(files):
+    def extract_and_show_wordcloud(files):
         texts = extract_text_from_pdfs(files)
         split_and_embed(texts)
         return generate_wordcloud()
 
-    load_button.click(fn=load_and_index, inputs=[pdf_input], outputs=[cloud_output])
-    ask_button.click(fn=answer_question, inputs=[question_input], outputs=[answer_output, cloud_output])
+    extract_btn.click(extract_and_show_wordcloud, inputs=[pdf_input], outputs=[wc_img])
+    ask_btn.click(answer_question, inputs=[question_input], outputs=[answer_out])
 
 demo.launch()