Update app.py

app.py

@@ -1,102 +1,146 @@
 import gradio as gr
 import torch
-import numpy as np
 import faiss
-import time
+import numpy as np
 import logging
+import time
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 from pypdf import PdfReader
 
-# ==========================
-# Logging
-# ==========================
-logging.basicConfig(level=logging.INFO)
+# =====================================================
+# LOGGING CONFIGURATION
+# =====================================================
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
 logger = logging.getLogger(__name__)
 
-# ==========================
-# Load Embedding Model
-# ==========================
+logger.info("Starting application...")
+
+# =====================================================
+# DEVICE CONFIG
+# =====================================================
+
+DEVICE = "cpu"
+torch.set_num_threads(4)
+
+# =====================================================
+# LOAD EMBEDDING MODEL
+# =====================================================
+
+logger.info("Loading embedding model...")
 embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+logger.info("Embedding model loaded.")
+
+# =====================================================
+# LOAD PHI-3 MODEL
+# =====================================================
 
-# ==========================
-# Load Phi-3 Mini (CPU Optimized)
-# ==========================
-model_name = "microsoft/Phi-3-mini-4k-instruct"
+MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+logger.info("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
+logger.info("Loading Phi-3 model (CPU optimized)...")
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
+    MODEL_NAME,
     torch_dtype=torch.float32,
     low_cpu_mem_usage=True
 )
 
-model.to("cpu")
+model.to(DEVICE)
 model.eval()
 
-# ==========================
-# Global Storage
-# ==========================
+logger.info("Model loaded successfully.")
+
+# =====================================================
+# GLOBAL STORAGE
+# =====================================================
+
 chunks = []
-index = None
+faiss_index = None
 
+# =====================================================
+# PDF PROCESSING
+# =====================================================
 
-# ==========================
-# Process PDF
-# ==========================
 def process_pdf(file):
-    global chunks, index
+    global chunks, faiss_index
+
+    logger.info("Processing PDF...")
 
     reader = PdfReader(file)
-    text = ""
+    full_text = ""
 
     for page in reader.pages:
-        content = page.extract_text()
-        if content:
-            text += content
+        text = page.extract_text()
+        if text:
+            full_text += text + "\n"
 
-    chunk_size = 350
-    chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
+    if not full_text.strip():
+        return "❌ Could not extract text from PDF."
+
+    # Chunking
+    chunk_size = 400
+    chunks = [
+        full_text[i:i+chunk_size]
+        for i in range(0, len(full_text), chunk_size)
+    ]
+
+    logger.info(f"Created {len(chunks)} chunks.")
+
+    # Embeddings
+    embeddings = embed_model.encode(chunks, convert_to_numpy=True)
 
-    embeddings = embed_model.encode(chunks)
     dimension = embeddings.shape[1]
+    faiss_index = faiss.IndexFlatL2(dimension)
+    faiss_index.add(embeddings)
 
-    index = faiss.IndexFlatL2(dimension)
-    index.add(np.array(embeddings))
+    logger.info("FAISS index built successfully.")
 
-    return "✅ PDF processed successfully. You can now start chatting."
+    return f"✅ PDF processed successfully ({len(chunks)} chunks created)."
 
+# =====================================================
+# CHAT FUNCTION
+# =====================================================
 
-# ==========================
-# Chat Function (RAG + Phi-3)
-# ==========================
-def chat_fn(message, history):
-    global chunks, index
+def generate_answer(message, history):
+    global chunks, faiss_index
 
-    if index is None:
+    if faiss_index is None:
         return "⚠ Please upload and process a PDF first."
 
+    logger.info(f"Received question: {message}")
+
     start_time = time.time()
 
-    # Embed query
-    query_embedding = embed_model.encode([message])
-    D, I = index.search(np.array(query_embedding), k=2)
+    # Step 1: Embed Query
+    query_embedding = embed_model.encode([message], convert_to_numpy=True)
+
+    # Step 2: Retrieve top 2 chunks
+    distances, indices = faiss_index.search(query_embedding, k=2)
+
+    retrieved_context = "\n\n".join(
+        [chunks[i] for i in indices[0]]
+    )
 
-    context = "\n".join([chunks[i] for i in I[0]])
+    logger.info("Retrieved relevant context.")
 
-    # Proper Phi-3 Instruct Template
-    prompt = f"""<|system|>
+    # Step 3: Create structured prompt
+    prompt = f"""
+<|system|>
 You are a professional AI assistant.
-Answer clearly, concisely and intelligently.
-Use structured explanation if helpful.
-Avoid repeating the question.
-If answer not found in context, say so.
+Provide clear, structured, intelligent answers.
+Keep answers concise but informative.
+If information is missing in context, say so.
 <|end|>
 
 <|user|>
 Context:
-{context}
+{retrieved_context}
 
 Question:
 {message}
@@ -110,53 +154,62 @@ Question:
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=120,
-            temperature=0.6,
+            max_new_tokens=150,
+            temperature=0.5,
             top_p=0.9,
+            repetition_penalty=1.1,
             do_sample=True,
-            repetition_penalty=1.15,
             use_cache=True
         )
 
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     answer = response.split("<|assistant|>")[-1].strip()
 
-    logger.info(f"Response time: {time.time() - start_time:.2f}s")
+    elapsed = time.time() - start_time
+    logger.info(f"Response generated in {elapsed:.2f} seconds.")
 
     return answer
 
+# =====================================================
+# GRADIO UI
+# =====================================================
 
-# ==========================
-# Beautiful Chat UI
-# ==========================
-with gr.Blocks(theme=gr.themes.Soft(), css="""
-#chatbot {height: 600px}
-""") as demo:
+with gr.Blocks() as demo:
 
-    gr.Markdown(
-        """
-        # 🤖 Smart RAG Assistant  
-        Powered by Phi-3 Mini + FAISS  
-        Upload a PDF and start chatting like ChatGPT.
-        """
-    )
+    gr.Markdown("# 🤖 Smart RAG Assistant")
+    gr.Markdown("Upload a PDF and chat intelligently using Phi-3 Mini.")
 
     with gr.Row():
+
         with gr.Column(scale=1):
             pdf_file = gr.File(label="Upload PDF")
             upload_btn = gr.Button("Process PDF")
             status = gr.Markdown()
 
         with gr.Column(scale=3):
-            chatbot = gr.ChatInterface(
-                fn=chat_fn,
-                chatbot=gr.Chatbot(elem_id="chatbot"),
-                textbox=gr.Textbox(placeholder="Ask something about the document...", container=False),
-                title="📘 Document Chat",
-                retry_btn="🔄 Retry",
-                clear_btn="🗑 Clear Chat"
+            chatbot = gr.Chatbot(height=600)
+            msg = gr.Textbox(
+                placeholder="Ask something about the document..."
             )
+            clear = gr.Button("Clear Chat")
+
+    upload_btn.click(
+        process_pdf,
+        inputs=pdf_file,
+        outputs=status
+    )
+
+    def respond(message, chat_history):
+        answer = generate_answer(message, chat_history)
+        chat_history.append((message, answer))
+        return "", chat_history
+
+    msg.submit(
+        respond,
+        inputs=[msg, chatbot],
+        outputs=[msg, chatbot]
+    )
 
-    upload_btn.click(process_pdf, inputs=pdf_file, outputs=status)
+    clear.click(lambda: [], None, chatbot)
 
-demo.launch()
+demo.launch(theme=gr.themes.Soft())