Update app.py

app.py

@@ -1,125 +1,71 @@
 import os
+import torch
 import gradio as gr
-import fitz  # PyMuPDF
 import faiss
-import numpy as np
-from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer
-from auto_gptq import AutoGPTQForCausalLM
+from transformers import AutoTokenizer, pipeline
+from langchain_community.vectorstores import FAISS
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from huggingface_hub import login
 
-# Authenticate
-hf_token = os.environ.get("HUGGINGFACE_TOKEN")
+# 🔐 Authenticate with Hugging Face using token stored in Secrets
+hf_token = os.getenv("HUGGINGFACE_TOKEN")
 if not hf_token:
-    raise ValueError("Hugging Face token not found.")
+    raise ValueError("❌ HUGGINGFACE_TOKEN not set in environment variables.")
 login(token=hf_token)
 
-# Load embedding model
-embed_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
-
-# Load 4-bit quantized Mistral model
+# 🔍 Load model and tokenizer
 model_id = "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ"
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
-model = AutoGPTQForCausalLM.from_quantized(
-    model_id,
-    use_safetensors=True,
-    trust_remote_code=True,
-    device_map="auto"
-)
-
-# Internal state
-index = None
-doc_texts = []
-
-# PDF/TXT text extraction
-def extract_text(file):
-    try:
-        text = ""
-        file_path = file.name if hasattr(file, 'name') else file
-        if file_path.endswith(".pdf"):
-            with fitz.open(file_path) as doc:
-                for page in doc:
-                    text += page.get_text()
-        elif file_path.endswith(".txt"):
-            with open(file_path, "r", encoding="utf-8") as f:
-                text = f.read()
-        else:
-            return "❌ Unsupported file type."
-        return text
-    except Exception as e:
-        return f"❌ Error extracting text: {e}"
-
-# Preprocess and embed
-def process_file(file):
-    global index, doc_texts
-    try:
-        text = extract_text(file)
-        if text.startswith("❌"):
-            return text
-
-        text = text[:15000]  # Limit size
-        splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
-        doc_texts = splitter.split_text(text)
-
-        if not doc_texts:
-            return "❌ Document could not be split."
-
-        embeddings = embed_model.encode(doc_texts, convert_to_numpy=True)
-        dim = embeddings.shape[1]
-        index = faiss.IndexFlatL2(dim)
-        index.add(embeddings)
-
-        return "✅ Document processed. Ask your question below."
-    except Exception as e:
-        return f"❌ Error processing file: {e}"
-
-# Generate answer using context
-def generate_answer(question):
-    global index, doc_texts
-    try:
-        if index is None or not doc_texts:
-            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"You are a helpful assistant. Use the context below to answer clearly.\n\n"
-            f"Context:\n{context}\n\n"
-            f"Question: {question}\n\n"
-            f"Answer:"
-        )
-
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        output = model.generate(
-            **inputs,
-            max_new_tokens=150,
-            do_sample=True,
-            temperature=0.7,
-            top_k=50,
-            top_p=0.95
-        )
-        answer = tokenizer.decode(output[0], skip_special_tokens=True)
-        return answer.split("Answer:")[-1].strip()
-    except Exception as e:
-        return f"❌ Error generating answer: {e}"
-
-# Gradio UI
-with gr.Blocks(title="📄 Document Q&A (Mistral 4-bit)") as demo:
-    gr.Markdown("<h1 style='text-align: center;'>📄 Document Q&A with Mistral 4-bit</h1>")
-    gr.Markdown("Upload a PDF or TXT and ask questions. Powered by Mistral-7B GPTQ.")
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=True)
+pipe = pipeline("text-generation", model=model_id, tokenizer=tokenizer, 
+                torch_dtype=torch.float16, device_map="auto", use_auth_token=True)
+
+# 🔎 Sentence transformer for embeddings
+embed_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+# Global store for vector DB
+db = None
+
+def process_pdf(pdf_path):
+    """Load, chunk, embed and index PDF into FAISS."""
+    loader = PyPDFLoader(pdf_path)
+    pages = loader.load()
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    docs = text_splitter.split_documents(pages)
+
+    global db
+    db = FAISS.from_documents(docs, embed_model)
+    return "✅ PDF processed successfully. Ask your questions now."
+
+def query_answer(question):
+    if not db:
+        return "⚠️ Please upload and process a PDF first."
+
+    docs = db.similarity_search(question, k=3)
+    context = "\n".join([doc.page_content for doc in docs])
+    prompt = f"[INST] You are a helpful assistant. Use the context below to answer the question:\n\nContext:\n{context}\n\nQuestion: {question}\n\nAnswer: [/INST]"
+
+    result = pipe(prompt, max_new_tokens=256, do_sample=True, top_k=5)[0]["generated_text"]
+    return result.replace(prompt, "").strip()
+
+# 🔧 Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# 📄 Document Q&A using Mistral-GPTQ")
 
     with gr.Row():
-        file_input = gr.File(label="Upload Document", file_types=[".pdf", ".txt"])
-        upload_output = gr.Textbox(label="Upload Status")
+        pdf_file = gr.File(label="Upload PDF", type="filepath")
+        upload_btn = gr.Button("Process PDF")
+
+    status = gr.Textbox(label="Status", interactive=False)
 
     with gr.Row():
-        question_input = gr.Textbox(label="Ask a Question", placeholder="e.g. What is this document about?")
-        answer_output = gr.Textbox(label="Answer")
+        user_question = gr.Textbox(label="Ask a Question")
+        ask_btn = gr.Button("Get Answer")
+
+    answer = gr.Textbox(label="Answer", lines=10)
 
-    file_input.change(fn=process_file, inputs=file_input, outputs=upload_output)
-    question_input.submit(fn=generate_answer, inputs=question_input, outputs=answer_output)
+    upload_btn.click(process_pdf, inputs=pdf_file, outputs=status)
+    ask_btn.click(query_answer, inputs=user_question, outputs=answer)
 
-demo.launch(show_error=True)
+demo.launch()