Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import gradio as gr
-import fitz  # PyMuPDF
+import fitz
 import faiss
 import numpy as np
 from sentence_transformers import SentenceTransformer
@@ -8,31 +8,21 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from huggingface_hub import login
 
-# Load Hugging Face Token from environment
 hf_token = os.environ.get("HUGGINGFACE_TOKEN")
 if not hf_token:
-    raise ValueError("⚠️ Please set the HUGGINGFACE_TOKEN environment variable.")
+    raise ValueError("Hugging Face token not found.")
 login(token=hf_token)
 
-# Load embedding model
 embed_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
 
-# Load small, fast LLM (great for CPU)
 model_id = "tiiuae/falcon-rw-1b"
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    device_map={"": "cpu"},
-    torch_dtype="auto",
-    token=hf_token
-)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map={"": "cpu"}, torch_dtype="auto", token=hf_token)
 llm = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# Globals
 index = None
 doc_texts = []
 
-# Extract text from PDF or TXT (handle Hugging Face Spaces file upload)
 def extract_text(file):
     text = ""
     file_path = file.name if hasattr(file, 'name') else file
@@ -44,14 +34,13 @@ def extract_text(file):
         with open(file_path, "r", encoding="utf-8") as f:
             text = f.read()
     else:
-        return "❌ Unsupported file type."
+        return "Unsupported file type."
     return text
 
-# Process file and build FAISS index
 def process_file(file):
     global index, doc_texts
     text = extract_text(file)
-    if text.startswith("❌"):
+    if text.startswith("Unsupported"):
         return text
 
     splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
@@ -62,40 +51,39 @@ def process_file(file):
     index = faiss.IndexFlatL2(dim)
     index.add(embeddings)
 
-    return "✅ File processed successfully! Ask your question below."
+    return "Document processed successfully. You can now ask questions."
 
-# Generate answer
 def generate_answer(question):
     global index, doc_texts
     if index is None or not doc_texts:
-        return "⚠️ Please upload and process a document first."
+        return "Please upload and process a document first."
 
     question_emb = embed_model.encode([question], convert_to_numpy=True)
     _, I = index.search(question_emb, k=3)
     context = "\n".join([doc_texts[i] for i in I[0]])
 
-    prompt = f"""[System: You are a helpful assistant. Answer based on the context.]
-
-Context:
-{context}
-
-Question: {question}
-Answer:"""
+    prompt = (
+        f"You are an intelligent assistant. Use the context below to answer the user's question clearly, "
+        f"politely, and completely. Do not just extract text — give a helpful response.\n\n"
+        f"Context:\n{context}\n\n"
+        f"User's Question: {question}\n\n"
+        f"Answer:"
+    )
 
     result = llm(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
     return result[0]["generated_text"].split("Answer:")[-1].strip()
 
-# Gradio UI
-with gr.Blocks(title="RAG Chatbot (CPU-Optimized)") as demo:
-    gr.Markdown("## 📚 Upload PDF/TXT and Ask Questions (Fast CPU RAG Bot)")
+with gr.Blocks(title="Document Q&A Assistant") as demo:
+    gr.Markdown("<h1 style='text-align: center;'>📄 Document AI Assistant</h1>")
+    gr.Markdown("Upload a PDF or TXT file, and ask questions about its content. The assistant will provide answers using the document as context.")
 
     with gr.Row():
-        file_input = gr.File(label="📁 Upload PDF or TXT", file_types=[".pdf", ".txt"])
-        upload_output = gr.Textbox(label="Upload Status", interactive=False)
+        file_input = gr.File(label="Upload PDF or TXT", file_types=[".pdf", ".txt"])
+        upload_output = gr.Textbox(label="Upload Status")
 
     with gr.Row():
-        question_input = gr.Textbox(label="❓ Ask a Question", placeholder="E.g. What is the document about?")
-        answer_output = gr.Textbox(label="💬 Answer", interactive=False)
+        question_input = gr.Textbox(label="Ask a Question", placeholder="What is this document about?")
+        answer_output = gr.Textbox(label="Answer")
 
     file_input.change(fn=process_file, inputs=file_input, outputs=upload_output)
     question_input.submit(fn=generate_answer, inputs=question_input, outputs=answer_output)