Update app.py

app.py

@@ -1,177 +1,153 @@
-import os
-import json
-import fitz  # PyMuPDF
-import re
-from tqdm import tqdm
-from docx import Document
+# ✅ SmartManuals-AI: Hugging Face Space App (RAM Safe, Multi-model, No Preview)
+
+import os, json, fitz, torch, chromadb, docx
+import gradio as gr
 from PIL import Image
-import pytesseract
-import io
-import torch
-import chromadb
+from nltk.tokenize import sent_tokenize
 from sentence_transformers import SentenceTransformer, util
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import gradio as gr
+from tqdm import tqdm
 
 # ---------------------------
-# 📁 Configuration
+# ⚙️ Constants
 # ---------------------------
-MANUALS_FOLDER = "./Manuals"
+MANUALS_DIR = "Manuals"
 CHROMA_PATH = "./chroma_store"
+CHUNKS_JSONL = "manual_chunks.jsonl"
 COLLECTION_NAME = "manual_chunks"
+HF_TOKEN = os.environ.get("HF_TOKEN")
 CHUNK_SIZE = 750
 CHUNK_OVERLAP = 100
-MAX_CONTEXT_CHUNKS = 3
-HF_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
-HF_TOKEN = os.environ.get("HF_TOKEN")
+TOP_K = 3
+
+MODEL_OPTIONS = {
+    "LLaMA 3.1 (8B)": "meta-llama/Llama-3.1-8B-Instruct",
+    "Mistral 7B": "mistralai/Mistral-7B-Instruct-v0.3",
+    "Gemma 7B": "google/gemma-7b-it",
+    "Qwen3 7B": "Qwen/Qwen1.5-7B-Chat"
+}
+
+# ---------------------------
+# 📄 Extract Text from PDFs and DOCX
+# ---------------------------
+def extract_text_from_pdf(path):
+    text = ""
+    try:
+        doc = fitz.open(path)
+        for page in doc:
+            page_text = page.get_text()
+            text += page_text + "\n"
+        doc.close()
+    except Exception as e:
+        print(f"❌ PDF Error in {path}: {e}")
+    return text
+
+def extract_text_from_docx(path):
+    try:
+        doc = docx.Document(path)
+        return "\n".join(p.text for p in doc.paragraphs if p.text.strip())
+    except Exception as e:
+        print(f"❌ DOCX Error in {path}: {e}")
+        return ""
 
 # ---------------------------
-# 🧹 Helpers
+# 🧹 Clean + Chunk
 # ---------------------------
 def clean(text):
-    lines = text.splitlines()
-    return "\n".join(line.strip() for line in lines if line.strip())
+    return "\n".join([line.strip() for line in text.splitlines() if line.strip()])
 
 def split_sentences(text):
-    return re.split(r'(?<=[.!?])\s+', text.strip())
+    return sent_tokenize(text)
 
-def chunk_sentences(sentences, max_len=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
+def chunk_text(sentences, size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
     chunks, chunk, length = [], [], 0
     for sent in sentences:
-        tokens = len(sent.split())
-        if length + tokens > max_len and chunk:
-            chunks.append(" ".join(chunk))
-            chunk = chunk[-overlap:] if overlap else []
-            length = sum(len(s.split()) for s in chunk)
+        n = len(sent.split())
+        if length + n > size:
+            if chunk:
+                chunks.append(" ".join(chunk))
+                chunk = chunk[-overlap:]
+                length = sum(len(s.split()) for s in chunk)
         chunk.append(sent)
-        length += tokens
+        length += n
     if chunk:
         chunks.append(" ".join(chunk))
     return chunks
 
-def extract_text_from_pdf(path):
-    doc = fitz.open(path)
-    full_text = []
-    for page in doc:
-        text = page.get_text().strip()
-        if not text:
-            try:
-                pix = page.get_pixmap(dpi=300)
-                img_data = pix.tobytes("png")
-                img = Image.open(io.BytesIO(img_data))
-                text = pytesseract.image_to_string(img).strip()
-            except Exception:
-                text = ""
-        full_text.append(text)
-    return "\n".join(full_text)
-
-def extract_text_from_docx(path):
-    doc = Document(path)
-    return "\n".join([para.text for para in doc.paragraphs if para.text.strip()])
-
-def extract_metadata(filename):
-    name = filename.lower()
-    model = next((m for m in ["se3hd", "se3", "se4", "symbio", "explore", "integrity x", "integrity sl", "everest", "engage", "inspire", "discover", "95t", "95x", "95c", "95r", "97c"] if m in name), "unknown")
-    if "om" in name or "owner" in name:
-        doc_type = "owner manual"
-    elif "sm" in name or "service" in name:
-        doc_type = "service manual"
-    elif "assembly" in name:
-        doc_type = "assembly instructions"
-    elif "alert" in name:
-        doc_type = "installer alert"
-    elif "parts" in name:
-        doc_type = "parts manual"
-    elif "bulletin" in name:
-        doc_type = "service bulletin"
-    else:
-        doc_type = "unknown"
-    return model, doc_type
-
 # ---------------------------
-# 🚀 Build ChromaDB at Startup
+# 📦 Embed and Store in Chroma
 # ---------------------------
 def embed_all():
+    print("🔍 Scanning manuals and embedding...")
+    os.makedirs(CHROMA_PATH, exist_ok=True)
+    embedder = SentenceTransformer("all-MiniLM-L6-v2")
     client = chromadb.PersistentClient(path=CHROMA_PATH)
     if COLLECTION_NAME in [c.name for c in client.list_collections()]:
         client.delete_collection(COLLECTION_NAME)
     collection = client.create_collection(COLLECTION_NAME)
 
-    embedder = SentenceTransformer("all-MiniLM-L6-v2")
-    records = []
-
-    for fname in os.listdir(MANUALS_FOLDER):
-        path = os.path.join(MANUALS_FOLDER, fname)
-        if not fname.lower().endswith((".pdf", ".docx")):
-            continue
+    all_chunks = []
+    files = [f for f in os.listdir(MANUALS_DIR) if f.lower().endswith((".pdf", ".docx"))]
+    for fname in tqdm(files):
+        path = os.path.join(MANUALS_DIR, fname)
         text = extract_text_from_pdf(path) if fname.endswith(".pdf") else extract_text_from_docx(path)
-        sents = split_sentences(clean(text))
-        chunks = chunk_sentences(sents)
-        model, doc_type = extract_metadata(fname)
+        text = clean(text)
+        sents = split_sentences(text)
+        chunks = chunk_text(sents)
         for i, chunk in enumerate(chunks):
-            records.append({
+            all_chunks.append({
                 "id": f"{fname}::chunk_{i+1}",
                 "text": chunk,
-                "metadata": {"source_file": fname, "model": model, "doc_type": doc_type}
+                "metadata": {"source": fname}
             })
 
-    for i in range(0, len(records), 16):
-        batch = records[i:i+16]
-        texts = [r["text"] for r in batch]
-        ids = [r["id"] for r in batch]
-        metas = [r["metadata"] for r in batch]
-        embeddings = embedder.encode(texts).tolist()
-        collection.add(documents=texts, ids=ids, metadatas=metas, embeddings=embeddings)
+    # Batch embed and store
+    for i in range(0, len(all_chunks), 16):
+        batch = all_chunks[i:i+16]
+        docs = [c["text"] for c in batch]
+        ids = [c["id"] for c in batch]
+        metas = [c["metadata"] for c in batch]
+        embs = embedder.encode(docs).tolist()
+        collection.add(documents=docs, ids=ids, metadatas=metas, embeddings=embs)
 
+    print(f"✅ Embedded {len(all_chunks)} chunks.")
     return collection, embedder
 
 # ---------------------------
-# 💬 Load HF Model
+# 🔍 RAG Search & LLM Answer
 # ---------------------------
-llm_pipe = None
-if HF_TOKEN:
-    tokenizer = AutoTokenizer.from_pretrained(HF_MODEL, token=HF_TOKEN)
-    model = AutoModelForCausalLM.from_pretrained(HF_MODEL, token=HF_TOKEN, torch_dtype=torch.float32)
-    llm_pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
+def ask(query, model_key):
+    model_id = MODEL_OPTIONS[model_key]
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
+        model = AutoModelForCausalLM.from_pretrained(model_id, token=HF_TOKEN, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+        model.to("cuda" if torch.cuda.is_available() else "cpu")
+        gen = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
+    except Exception as e:
+        return f"❌ Model loading failed: {e}"
+
+    results = db.query(query_texts=[query], n_results=TOP_K)
+    chunks = results["documents"][0]
+    context = "\n\n".join(chunks)
+    prompt = f"Answer this using only the context below.\n\nContext:\n{context}\n\nQuestion: {query}\nAnswer:"
+
+    try:
+        res = gen(prompt, max_new_tokens=300, do_sample=False)[0]['generated_text']
+        return res.split("Answer:", 1)[-1].strip()
+    except Exception as e:
+        return f"❌ LLM failed: {e}"
 
 # ---------------------------
-# 🔎 RAG Function
-# ---------------------------
-def run_query(question):
-    if not question.strip():
-        return "Please enter a question."
-    if not db or not embedder:
-        return "Chroma or embedder not ready."
-
-    q_embed = embedder.encode(question).tolist()
-    res = db.query(query_embeddings=[q_embed], n_results=MAX_CONTEXT_CHUNKS)
-    contexts = res["documents"][0]
-    prompt = """
-You are a technical assistant.
-Answer only using the context below.
-Say 'I don't know' if not found.
-
-"""
-    context_text = "\n\n".join(contexts)
-    final_prompt = prompt + f"Context:\n{context_text}\n\nQuestion: {question}\nAnswer:"
-    if llm_pipe:
-        result = llm_pipe(final_prompt, max_new_tokens=300)[0]['generated_text']
-        return result.split("Answer:")[-1].strip()
-    return "Model not loaded."
-
-# ---------------------------
-# 🧠 Init embeddings once
+# ▶️ UI
 # ---------------------------
 db, embedder = embed_all()
-
-# ---------------------------
-# 🎛️ Gradio Interface
-# ---------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 SmartManuals-AI: Ask Technical Questions about Your Manuals")
-    question = gr.Textbox(placeholder="e.g. How do I reset the treadmill console?", label="Enter Question")
-    submit = gr.Button("Get Answer")
-    output = gr.Textbox(label="Answer")
-    submit.click(fn=run_query, inputs=question, outputs=output)
+    gr.Markdown("## 🧠 SmartManuals-AI — Ask Your PDF and Word Docs")
+    with gr.Row():
+        qbox = gr.Textbox(label="Ask a Question", placeholder="e.g. How do I calibrate SE3 console?")
+        model_pick = gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), label="Choose a Model", value="Mistral 7B")
+    answer = gr.Textbox(label="Answer", lines=8)
+    ask_btn = gr.Button("Ask")
+    ask_btn.click(fn=ask, inputs=[qbox, model_pick], outputs=[answer])
 
 demo.launch()