Update app.py

app.py

@@ -1,71 +1,83 @@
 import os
+import time
 import torch
 import gradio as gr
-import faiss
-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
+from transformers import AutoTokenizer
+from auto_gptq import AutoGPTQForCausalLM
+from sentence_transformers import SentenceTransformer
+from langchain_community.vectorstores import FAISS
 
-# 🔐 Authenticate with Hugging Face using token stored in Secrets
-hf_token = os.getenv("HUGGINGFACE_TOKEN")
+# Load HF token and login
+hf_token = os.environ.get("HUGGINGFACE_TOKEN")
 if not hf_token:
-    raise ValueError("❌ HUGGINGFACE_TOKEN not set in environment variables.")
-login(token=hf_token)
-
-# 🔍 Load model and tokenizer
-model_id = "TheBloke/Mistral-7B-Instruct-v0.1-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
+    raise ValueError("Please set the HUGGINGFACE_TOKEN environment variable")
 
-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()
+login(token=hf_token)
 
-# 🔧 Gradio UI
+# Load tokenizer and quantized model
+model_id = "TheBloke/mistral-7B-GPTQ"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
+
+print("Loading quantized model...")
+start = time.time()
+model = AutoGPTQForCausalLM.from_quantized(
+    model_id,
+    use_safetensors=True,
+    device=device,
+    use_triton=True,
+    quantize_config=None,
+)
+print(f"Model loaded in {time.time() - start:.2f} seconds on {device}")
+
+# Load embedding model for FAISS vector store
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
+
+# Sample documents to build vector index (can replace with your own)
+texts = [
+    "Hello world",
+    "Mistral 7B is a powerful language model",
+    "Langchain and FAISS make vector search easy",
+    "This is a test document for vector search",
+]
+embeddings = embedder.encode(texts)
+
+faiss_index = FAISS.from_embeddings(embeddings, texts)
+
+# Generate text from prompt
+def generate_text(prompt, max_length=128):
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_length=max_length)
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return decoded
+
+# Search docs with vector similarity
+def search_docs(query):
+    query_emb = embedder.encode([query])
+    results = faiss_index.similarity_search_by_vector(query_emb[0], k=3)
+    return "\n\n".join(results)
+
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 📄 Document Q&A using Mistral-GPTQ")
-
-    with gr.Row():
-        pdf_file = gr.File(label="Upload PDF", type="filepath")
-        upload_btn = gr.Button("Process PDF")
+    gr.Markdown("# Mistral GPTQ + FAISS Vector Search Demo")
 
-    status = gr.Textbox(label="Status", interactive=False)
+    with gr.Tab("Text Generation"):
+        prompt_input = gr.Textbox(label="Enter prompt", lines=3)
+        generate_btn = gr.Button("Generate")
+        output_text = gr.Textbox(label="Output", lines=6)
 
-    with gr.Row():
-        user_question = gr.Textbox(label="Ask a Question")
-        ask_btn = gr.Button("Get Answer")
+        generate_btn.click(fn=generate_text, inputs=prompt_input, outputs=output_text)
 
-    answer = gr.Textbox(label="Answer", lines=10)
+    with gr.Tab("Vector Search"):
+        query_input = gr.Textbox(label="Enter search query", lines=2)
+        search_btn = gr.Button("Search")
+        search_output = gr.Textbox(label="Search Results", lines=6)
 
-    upload_btn.click(process_pdf, inputs=pdf_file, outputs=status)
-    ask_btn.click(query_answer, inputs=user_question, outputs=answer)
+        search_btn.click(fn=search_docs, inputs=query_input, outputs=search_output)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()