Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import fitz  # PyMuPDF
+import fitz
 import re
 import faiss
 import torch
@@ -10,15 +10,12 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
 # ===============================
-# MODEL LOADING (ONCE)
+# MODEL LOADING
 # ===============================
 
-# Embedding model (lightweight & fast)
 embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
 
-# Open-source LLM (CPU friendly)
 LLM_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-
 tokenizer = AutoTokenizer.from_pretrained(LLM_NAME)
 llm = AutoModelForCausalLM.from_pretrained(
     LLM_NAME,
@@ -29,7 +26,7 @@ llm.eval()
 
 
 # ===============================
-# PDF PROCESSING FUNCTIONS
+# PDF PROCESSING
 # ===============================
 
 def extract_text_from_pdf(pdf_path):
@@ -47,48 +44,41 @@ def clean_text(text):
 def chunk_text(text, chunk_size=500, overlap=50):
     chunks = []
     start = 0
-    text_length = len(text)
-
-    while start < text_length:
+    while start < len(text):
         end = start + chunk_size
         chunks.append(text[start:end])
         start = end - overlap
-
     return chunks
 
 
 # ===============================
-# VECTOR DATABASE (FAISS)
+# VECTOR DB (FAISS)
 # ===============================
 
 def build_faiss_index(chunks):
     embeddings = embedding_model.encode(chunks)
     embeddings = np.array(embeddings).astype("float32")
-
     index = faiss.IndexFlatL2(embeddings.shape[1])
     index.add(embeddings)
-
     return index, chunks
 
 
 def retrieve_relevant_chunks(query, index, chunks, top_k=3):
     query_embedding = embedding_model.encode([query]).astype("float32")
     _, indices = index.search(query_embedding, top_k)
-
     return [chunks[i] for i in indices[0]]
 
 
 # ===============================
-# ANSWER GENERATION (LLM)
+# LLM ANSWER
 # ===============================
 
 def generate_answer(question, context_chunks):
     context = "\n\n".join(context_chunks)
 
     prompt = f"""
-You are an AI assistant.
 Answer the question strictly using the given context.
-If the answer is not found, reply:
+If the answer is not found, say:
 "Information not found in the document."
 
 Context:
@@ -109,72 +99,61 @@ Answer:
             temperature=0.2
         )
 
-    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
-    return decoded_output.split("Answer:")[-1].strip()
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+    return decoded.split("Answer:")[-1].strip()
 
 
 # ===============================
-# MAIN RAG PIPELINE
+# MAIN PIPELINE
 # ===============================
 
 def pdf_rag_chat(pdf_file, question):
     if pdf_file is None or question.strip() == "":
         return "Please upload a PDF and enter a question."
 
-    # 1. Extract & preprocess text
-    raw_text = extract_text_from_pdf(pdf_file.name)
-    cleaned_text = clean_text(raw_text)
+    text = extract_text_from_pdf(pdf_file.name)
+    text = clean_text(text)
 
-    # 2. Chunking
-    chunks = chunk_text(cleaned_text)
-
-    # 3. Vector DB
+    chunks = chunk_text(text)
     index, chunks = build_faiss_index(chunks)
+    context = retrieve_relevant_chunks(question, index, chunks)
 
-    # 4. Retrieval
-    relevant_chunks = retrieve_relevant_chunks(question, index, chunks)
-
-    # 5. LLM Answer
-    return generate_answer(question, relevant_chunks)
+    return generate_answer(question, context)
 
 
 # ===============================
-# GRADIO UI (PRODUCTION READY)
+# GRADIO UI (GRADIO 6 SAFE)
 # ===============================
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+with gr.Blocks() as demo:
 
     gr.Markdown("""
     # 📄 PDF RAG Chatbot (Open-Source AI)
 
-    Upload a **PDF document** and ask questions based **only on its content**.  
-    This system uses a **Retrieval Augmented Generation (RAG)** architecture with  
-    **open-source Hugging Face models**, running entirely on **free CPU**.
-
-    ---
+    Upload a **PDF** and ask questions based **only on its content**.  
+    Built using **Retrieval Augmented Generation (RAG)** and  
+    **open-source Hugging Face models**, running on **free CPU**.
     """)
 
     with gr.Row():
         with gr.Column(scale=1):
             pdf_input = gr.File(
                 label="📤 Upload PDF",
-                file_types=[".pdf"],
-                file_count="single"
+                file_types=[".pdf"]
             )
 
             question_input = gr.Textbox(
-                label="❓ Ask your question",
+                label="❓ Ask a question",
                 placeholder="e.g. What is the objective of the project?",
                 lines=2
             )
 
-            submit_btn = gr.Button("🔍 Get Answer", variant="primary")
+            submit_btn = gr.Button("🔍 Get Answer")
 
         with gr.Column(scale=2):
             answer_output = gr.Textbox(
                 label="📌 Answer",
-                lines=10,
-                show_copy_button=True
+                lines=10
             )
 
     submit_btn.click(
@@ -185,14 +164,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     gr.Markdown("""
     ---
-    ### ⚙️ System Information
-    - **LLM:** TinyLlama (Open-Source, Hugging Face)
-    - **Embeddings:** Sentence Transformers
-    - **Vector Store:** FAISS
-    - **Deployment:** Hugging Face Spaces (Free CPU)
-
-    ---
-    © **Simranpreet Kaur**  
+    **© Simranpreet Kaur**  
     **NIELIT Ropar | AIML Six Months Training | 2026**
     """)