Update app.py

app.py

@@ -1,23 +1,24 @@
 import os
 import gradio as gr
-from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import fitz  # PyMuPDF
 import faiss
 import numpy as np
+from io import BytesIO
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-import fitz  # PyMuPDF
 from huggingface_hub import login
 
-# Authenticate with Hugging Face to access gated models
+# 1. Authenticate HuggingFace
 hf_token = os.environ.get("HUGGINGFACE_TOKEN")
-if hf_token is None:
-    raise ValueError("Please set the HUGGINGFACE_TOKEN environment variable")
+if not hf_token:
+    raise ValueError("⚠️ Please set the HUGGINGFACE_TOKEN environment variable.")
 login(token=hf_token)
 
-# Load embedding model
+# 2. Load embedding model
 embed_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
 
-# Load LLM model and tokenizer with 4bit quantization
+# 3. Load LLM (Mistral 7B Instruct with 4-bit quantization)
 model_id = "mistralai/Mistral-7B-Instruct-v0.1"
 tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=hf_token)
 model = AutoModelForCausalLM.from_pretrained(
@@ -28,67 +29,77 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 llm = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# Globals for FAISS index and document texts
+# 4. Globals
 index = None
 doc_texts = []
 
-# PDF/Text extraction
+# 5. Extract text from uploaded file
 def extract_text(file):
+    text = ""
+    file_bytes = file.read()
     if file.name.endswith(".pdf"):
-        text = ""
-        doc = fitz.open(file.name)
+        pdf_stream = BytesIO(file_bytes)
+        doc = fitz.open(stream=pdf_stream, filetype="pdf")
         for page in doc:
             text += page.get_text()
-        return text
     elif file.name.endswith(".txt"):
-        return file.read().decode("utf-8")
+        text = file_bytes.decode("utf-8")
     else:
-        return "❌ Invalid file type."
+        return "❌ Unsupported file type. Only PDF and TXT are allowed."
+    return text
 
-# File processing: chunk text, create embeddings, build FAISS index
+# 6. Process the file: split text, create embeddings, build FAISS index
 def process_file(file):
     global index, doc_texts
     text = extract_text(file)
     if text.startswith("❌"):
         return text
 
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
-    doc_texts = text_splitter.split_text(text)
-    embeddings = embed_model.encode(doc_texts)
+    # Split text
+    splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
+    doc_texts = splitter.split_text(text)
+
+    # Create embeddings
+    embeddings = embed_model.encode(doc_texts, convert_to_numpy=True)
 
+    # Build FAISS index
     dim = embeddings.shape[1]
     index = faiss.IndexFlatL2(dim)
-    index.add(np.array(embeddings))
+    index.add(embeddings)
 
     return "✅ File processed successfully. You can now ask questions!"
 
-# Generate answer using retrieved context and LLM
+# 7. Generate answer based on question + retrieved context
 def generate_answer(question):
     global index, doc_texts
-    if index is None or len(doc_texts) == 0:
+    if index is None or not doc_texts:
         return "⚠️ Please upload and process a file first."
-    
-    question_embedding = embed_model.encode([question])
-    _, I = index.search(np.array(question_embedding), k=3)
-    context = "\n".join([doc_texts[i] for i in I[0]])
 
-    prompt = f"""[System: You are a helpful assistant. Answer strictly based on the context.]
+    # Embed the question
+    question_emb = embed_model.encode([question], convert_to_numpy=True)
+    _, I = index.search(question_emb, k=3)
+
+    # Build context
+    context = "\n".join([doc_texts[i] for i in I[0]])
 
+    # Prompt
+    prompt = f"""[System: You are a helpful assistant. Answer strictly based on the context. Do not hallucinate.]
 Context:
 {context}
 
 Question: {question}
 Answer:"""
 
-    result = llm(prompt, max_new_tokens=300, do_sample=True, temperature=0.7)
-    return result[0]["generated_text"].split("Answer:")[-1].strip()
+    # Generate response
+    response = llm(prompt, max_new_tokens=300, do_sample=True, temperature=0.7)
+    return response[0]["generated_text"].split("Answer:")[-1].strip()
 
-# Gradio UI
-with gr.Blocks(title="RAG Chatbot") as demo:
-    gr.Markdown("## 📚 RAG Chatbot - Upload PDF/TXT and Ask Questions")
+# 8. Gradio UI
+with gr.Blocks(title="🧠 RAG Chatbot") as demo:
+    gr.Markdown("## 📚 Retrieval-Augmented Generation Chatbot\nUpload a `.pdf` or `.txt` and ask questions from the content.")
 
     with gr.Row():
-        file_input = gr.File(label="📁 Upload .pdf or .txt", file_types=[".pdf", ".txt"])
+        file_input = gr.File(label="📁 Upload PDF/TXT", file_types=[".pdf", ".txt"])
         upload_status = gr.Textbox(label="📥 Upload Status", interactive=False)
 
     with gr.Row():
@@ -98,4 +109,5 @@ with gr.Blocks(title="RAG Chatbot") as demo:
     file_input.change(fn=process_file, inputs=file_input, outputs=upload_status)
     question_box.submit(fn=generate_answer, inputs=question_box, outputs=answer_box)
 
+# 9. Launch the app
 demo.launch()