Update app.py

app.py

@@ -11,6 +11,9 @@ import numpy as np
 from PIL import Image
 from transformers import AutoProcessor, VisionEncoderDecoderModel, Gemma3nForConditionalGeneration, pipeline
 import torch
+from sentence_transformers import SentenceTransformer
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
 import os
 import tempfile
 import uuid
@@ -312,6 +315,15 @@ except Exception as e:
     dolphin_model = None
     model_status = f"❌ Model failed to load: {str(e)}"
 
+# Initialize embedding model for RAG
+try:
+    print("Loading embedding model...")
+    embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+    print("✅ Embedding model loaded successfully")
+except Exception as e:
+    print(f"❌ Error loading embedding model: {e}")
+    embedding_model = None
+
 # Initialize chatbot model
 try:
     import os
@@ -350,12 +362,57 @@ except Exception as e:
 # Global state for managing tabs
 processed_markdown = ""
 show_results_tab = False
+document_chunks = []
+document_embeddings = None
+embedding_model = None
 # chatbot_model is initialized above
 
 
+def chunk_document(text, chunk_size=500, overlap=50):
+    """Split document into overlapping chunks for RAG"""
+    words = text.split()
+    chunks = []
+    
+    for i in range(0, len(words), chunk_size - overlap):
+        chunk = ' '.join(words[i:i + chunk_size])
+        if chunk.strip():
+            chunks.append(chunk)
+    
+    return chunks
+
+def create_embeddings(chunks):
+    """Create embeddings for document chunks"""
+    if embedding_model is None:
+        return None
+    
+    try:
+        embeddings = embedding_model.encode(chunks)
+        return embeddings
+    except Exception as e:
+        print(f"Error creating embeddings: {e}")
+        return None
+
+def retrieve_relevant_chunks(question, chunks, embeddings, top_k=3):
+    """Retrieve most relevant chunks for a question"""
+    if embedding_model is None or embeddings is None:
+        return chunks[:3]  # Fallback to first 3 chunks
+    
+    try:
+        question_embedding = embedding_model.encode([question])
+        similarities = cosine_similarity(question_embedding, embeddings)[0]
+        
+        # Get top-k most similar chunks
+        top_indices = np.argsort(similarities)[-top_k:][::-1]
+        relevant_chunks = [chunks[i] for i in top_indices]
+        
+        return relevant_chunks
+    except Exception as e:
+        print(f"Error retrieving chunks: {e}")
+        return chunks[:3]  # Fallback
+
 def process_uploaded_pdf(pdf_file, progress=gr.Progress()):
     """Main processing function for uploaded PDF"""
-    global processed_markdown, show_results_tab
+    global processed_markdown, show_results_tab, document_chunks, document_embeddings
     
     if dolphin_model is None:
         return "❌ Model not loaded", gr.Tabs(visible=False)
@@ -368,6 +425,13 @@ def process_uploaded_pdf(pdf_file, progress=gr.Progress()):
         
         if status == "processing_complete":
             processed_markdown = combined_markdown
+            
+            # Create chunks and embeddings for RAG
+            print("Creating document chunks for RAG...")
+            document_chunks = chunk_document(processed_markdown)
+            document_embeddings = create_embeddings(document_chunks)
+            print(f"Created {len(document_chunks)} chunks")
+            
             show_results_tab = True
             return "✅ PDF processed successfully! Check the 'Document' tab above.", gr.Tabs(visible=True)
         else:
@@ -446,11 +510,13 @@ with gr.Blocks(
         # Home Tab
         with gr.TabItem("🏠 Home", id="home"):
             chatbot_status = "✅ Chatbot ready" if chatbot_model else "❌ Chatbot not loaded"
+            embedding_status = "✅ RAG ready" if embedding_model else "❌ RAG not loaded"
             gr.Markdown(
                 "# Scholar Express\n"
                 "### Upload a research paper to get a web-friendly version, an AI chatbot, and a podcast summary. Because of our reliance on Generative AI, some errors are inevitable.\n"
                 f"**PDF Processing:** {model_status}\n"
-                f"**Chatbot:** {chatbot_status}"
+                f"**Chatbot:** {chatbot_status}\n"
+                f"**RAG System:** {embedding_status}"
             )
             
             with gr.Column(elem_classes="upload-container"):
@@ -521,7 +587,7 @@ with gr.Blocks(
                 send_btn = gr.Button("Send", variant="primary", scale=1)
             
             gr.Markdown(
-                "*Ask questions about your processed document. The AI will use the document content to provide accurate answers.*",
+                "*Ask questions about your processed document. The AI uses RAG (Retrieval-Augmented Generation) to find relevant sections and provide accurate answers.*",
                 elem_id="chat-notice"
             )
     
@@ -562,18 +628,23 @@ with gr.Blocks(
             return history + [[message, "❌ Please process a PDF document first before asking questions."]]
         
         try:
-            # Create context with the processed document
-            context = f"Document content:\n{processed_markdown[:3000]}..." if len(processed_markdown) > 3000 else f"Document content:\n{processed_markdown}"
+            # Use RAG to get relevant chunks instead of full document
+            if document_chunks and len(document_chunks) > 0:
+                relevant_chunks = retrieve_relevant_chunks(message, document_chunks, document_embeddings)
+                context = "\n\n".join(relevant_chunks)
+            else:
+                # Fallback to truncated document if RAG fails
+                context = processed_markdown[:1500] + "..." if len(processed_markdown) > 1500 else processed_markdown
             
             # Create chat messages
             messages = [
                 {
                     "role": "system",
-                    "content": [{"type": "text", "text": "You are a helpful assistant that answers questions about documents. Use the provided document content to answer questions accurately."}]
+                    "content": [{"type": "text", "text": "You are a helpful assistant that answers questions about documents. Use the provided document content to answer questions accurately and concisely."}]
                 },
                 {
                     "role": "user",
-                    "content": [{"type": "text", "text": f"{context}\n\nQuestion: {message}"}]
+                    "content": [{"type": "text", "text": f"Document content:\n{context}\n\nQuestion: {message}"}]
                 }
             ]