Update app.py

app.py

@@ -1,31 +1,29 @@
 import gradio as gr
 import torch
-import time
-import logging
 import numpy as np
 import faiss
+import time
+import logging
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 from pypdf import PdfReader
 
-# ==============================
-# Logging Setup
-# ==============================
+# ==========================
+# Logging
+# ==========================
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# ==============================
+# ==========================
 # Load Embedding Model
-# ==============================
-logger.info("Loading embedding model...")
+# ==========================
 embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-# ==============================
+# ==========================
 # Load Phi-3 Mini (CPU Optimized)
-# ==============================
+# ==========================
 model_name = "microsoft/Phi-3-mini-4k-instruct"
 
-logger.info("Loading Phi-3-mini model...")
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 model = AutoModelForCausalLM.from_pretrained(
@@ -37,22 +35,20 @@ model = AutoModelForCausalLM.from_pretrained(
 model.to("cpu")
 model.eval()
 
-logger.info("Model loaded successfully.")
-
-# ==============================
+# ==========================
 # Global Storage
-# ==============================
+# ==========================
 chunks = []
 index = None
 
-# ==============================
-# PDF Processing
-# ==============================
-def process_pdf(pdf_file):
+
+# ==========================
+# Process PDF
+# ==========================
+def process_pdf(file):
     global chunks, index
 
-    logger.info("Processing PDF...")
-    reader = PdfReader(pdf_file)
+    reader = PdfReader(file)
     text = ""
 
     for page in reader.pages:
@@ -60,49 +56,42 @@ def process_pdf(pdf_file):
         if content:
             text += content
 
-    # Smaller chunks = faster generation
     chunk_size = 350
     chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
 
-    logger.info(f"Total chunks created: {len(chunks)}")
-
     embeddings = embed_model.encode(chunks)
     dimension = embeddings.shape[1]
 
     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 "✅ PDF processed successfully!"
 
-# ==============================
-# Smart RAG Question Answering
-# ==============================
-def ask_question(query):
+# ==========================
+# Chat Function (RAG + Phi-3)
+# ==========================
+def chat_fn(message, history):
     global chunks, index
 
     if index is None:
-        return "Please upload and process a PDF first."
+        return "⚠ Please upload and process a PDF first."
 
-    start_total = time.time()
-    logger.info("Received question.")
+    start_time = time.time()
 
     # Embed query
-    query_embedding = embed_model.encode([query])
-
-    # Retrieve top 2 relevant chunks
+    query_embedding = embed_model.encode([message])
     D, I = index.search(np.array(query_embedding), k=2)
 
     context = "\n".join([chunks[i] for i in I[0]])
 
-    # Phi-3 Instruct Template (CRITICAL)
+    # Proper Phi-3 Instruct Template
     prompt = f"""<|system|>
-You are an expert AI assistant.
-Answer clearly, accurately, and concisely.
-Use structured explanation when helpful.
+You are a professional AI assistant.
+Answer clearly, concisely and intelligently.
+Use structured explanation if helpful.
 Avoid repeating the question.
-If answer not in context, say so.
+If answer not found in context, say so.
 <|end|>
 
 <|user|>
@@ -110,7 +99,7 @@ Context:
 {context}
 
 Question:
-{query}
+{message}
 <|end|>
 
 <|assistant|>
@@ -118,11 +107,7 @@ Question:
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
-    logger.info(f"Prompt token length: {len(inputs['input_ids'][0])}")
-
     with torch.no_grad():
-        start_gen = time.time()
-
         outputs = model.generate(
             **inputs,
             max_new_tokens=120,
@@ -133,32 +118,45 @@ Question:
             use_cache=True
         )
 
-        logger.info(f"Generation time: {time.time() - start_gen:.2f}s")
-
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Remove prompt from response
     answer = response.split("<|assistant|>")[-1].strip()
 
-    logger.info(f"Total response time: {time.time() - start_total:.2f}s")
+    logger.info(f"Response time: {time.time() - start_time:.2f}s")
 
     return answer
 
 
-# ==============================
-# Gradio UI
-# ==============================
-with gr.Blocks() as demo:
-    gr.Markdown("# 📚 Optimized RAG with Phi-3-mini")
-
-    pdf_input = gr.File(label="Upload PDF")
-    upload_btn = gr.Button("Process PDF")
-    status = gr.Textbox(label="Status")
-
-    question = gr.Textbox(label="Ask a question")
-    answer = gr.Textbox(label="Answer")
-
-    upload_btn.click(process_pdf, inputs=pdf_input, outputs=status)
-    question.submit(ask_question, inputs=question, outputs=answer)
+# ==========================
+# Beautiful Chat UI
+# ==========================
+with gr.Blocks(theme=gr.themes.Soft(), css="""
+#chatbot {height: 600px}
+""") as demo:
+
+    gr.Markdown(
+        """
+        # 🤖 Smart RAG Assistant  
+        Powered by Phi-3 Mini + FAISS  
+        Upload a PDF and start chatting like ChatGPT.
+        """
+    )
+
+    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"
+            )
+
+    upload_btn.click(process_pdf, inputs=pdf_file, outputs=status)
 
 demo.launch()