Update app.py

app.py

@@ -1,156 +1,177 @@
-# app.py
-# SmartManuals-AI: Hugging Face Space version
-
-import os, json, fitz, nltk, chromadb, io
-import torch
+import os
+import json
+import fitz  # PyMuPDF
+import re
 from tqdm import tqdm
-from PIL import Image
 from docx import Document
-from sentence_transformers import SentenceTransformer, util
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-from nltk.tokenize import sent_tokenize
+from PIL import Image
 import pytesseract
+import io
+import torch
+import chromadb
+from sentence_transformers import SentenceTransformer, util
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import gradio as gr
 
-# ----------------------
-# Configuration
-# ----------------------
+# ---------------------------
+# 📁 Configuration
+# ---------------------------
 MANUALS_FOLDER = "./Manuals"
-CHUNKS_JSONL = "chunks.jsonl"
 CHROMA_PATH = "./chroma_store"
 COLLECTION_NAME = "manual_chunks"
 CHUNK_SIZE = 750
 CHUNK_OVERLAP = 100
-MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-# ----------------------
-# Ensure punkt is downloaded
-# ----------------------
-nltk.download("punkt")
-
-# ----------------------
-# Utilities
-# ----------------------
-def extract_text_from_pdf(path):
-    doc = fitz.open(path)
-    text = ""
-    for page in doc:
-        t = page.get_text()
-        if not t.strip():
-            pix = page.get_pixmap(dpi=300)
-            img = Image.open(io.BytesIO(pix.tobytes("png")))
-            t = pytesseract.image_to_string(img)
-        text += t + "\n"
-    return text
-
-def extract_text_from_docx(path):
-    doc = Document(path)
-    return "\n".join(p.text for p in doc.paragraphs if p.text.strip())
+MAX_CONTEXT_CHUNKS = 3
+HF_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
+# ---------------------------
+# 🧹 Helpers
+# ---------------------------
 def clean(text):
-    return "\n".join([line.strip() for line in text.splitlines() if line.strip()])
+    lines = text.splitlines()
+    return "\n".join(line.strip() for line in lines if line.strip())
 
 def split_sentences(text):
-    return sent_tokenize(text)
+    return re.split(r'(?<=[.!?])\s+', text.strip())
 
-def chunk_sentences(sentences, max_tokens=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
-    chunks, chunk, count = [], [], 0
-    for s in sentences:
-        words = s.split()
-        if count + len(words) > max_tokens:
+def chunk_sentences(sentences, max_len=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 > 0 else []
-            count = sum(len(x.split()) for x in chunk)
-        chunk.append(s)
-        count += len(words)
+            chunk = chunk[-overlap:] if overlap else []
+            length = sum(len(s.split()) for s in chunk)
+        chunk.append(sent)
+        length += tokens
     if chunk:
         chunks.append(" ".join(chunk))
     return chunks
 
-def get_metadata(filename):
+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()
-    return {
-        "source_file": filename,
-        "doc_type": "service manual" if "sm" in name else "owner manual" if "om" in name else "unknown",
-        "model": "se3hd" if "se3hd" in name else "unknown"
-    }
-
-# ----------------------
-# Embedding
-# ----------------------
+    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
+# ---------------------------
 def embed_all():
-    embedder = SentenceTransformer("all-MiniLM-L6-v2")
     client = chromadb.PersistentClient(path=CHROMA_PATH)
-    try:
+    if COLLECTION_NAME in [c.name for c in client.list_collections()]:
         client.delete_collection(COLLECTION_NAME)
-    except:
-        pass
     collection = client.create_collection(COLLECTION_NAME)
-    chunks, metadatas, ids = [], [], []
-    files = os.listdir(MANUALS_FOLDER)
-    idx = 0
-    for file in tqdm(files):
-        path = os.path.join(MANUALS_FOLDER, file)
-        text = extract_text_from_pdf(path) if file.endswith(".pdf") else extract_text_from_docx(path)
-        meta = get_metadata(file)
-        sents = split_sentences(clean(text))
-        for i, chunk in enumerate(chunk_sentences(sents)):
-            chunks.append(chunk)
-            ids.append(f"{file}::chunk_{i}")
-            metadatas.append(meta)
-            if len(chunks) >= 16:
-                emb = embedder.encode(chunks).tolist()
-                collection.add(documents=chunks, ids=ids, metadatas=metadatas, embeddings=emb)
-                chunks, ids, metadatas = [], [], []
-    if chunks:
-        emb = embedder.encode(chunks).tolist()
-        collection.add(documents=chunks, ids=ids, metadatas=metadatas, embeddings=emb)
-    return collection, embedder
 
-# ----------------------
-# Model setup
-# ----------------------
-def load_model():
-    device = 0 if torch.cuda.is_available() else -1
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, token=HF_TOKEN)
-    return pipeline("text-generation", model=model, tokenizer=tokenizer, device=device, max_new_tokens=512)
-
-# ----------------------
-# RAG Pipeline
-# ----------------------
-def answer_query(question):
-    results = db.query(query_texts=[question], n_results=5)
-    context = "\n\n".join(results["documents"][0])
-    prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
-You are a helpful assistant. Use the provided context to answer questions. If you don't know, say 'I don't know.'
-<context>
-{context}
-</context>
-<|start_header_id|>user<|end_header_id|>
-{question}<|start_header_id|>assistant<|end_header_id|>"""
-    return llm(prompt)[0]["generated_text"].split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
-
-# ----------------------
-# UI
-# ----------------------
-with gr.Blocks() as demo:
-    status = gr.Textbox(label="Status", value="Embedding manuals... Please wait.", interactive=False)
-    question = gr.Textbox(label="Ask a Question")
-    submit = gr.Button("🔍 Ask")
-    answer = gr.Textbox(label="Answer", lines=8)
+    embedder = SentenceTransformer("all-MiniLM-L6-v2")
+    records = []
 
-    def handle_query(q):
-        return answer_query(q)
+    for fname in os.listdir(MANUALS_FOLDER):
+        path = os.path.join(MANUALS_FOLDER, fname)
+        if not fname.lower().endswith((".pdf", ".docx")):
+            continue
+        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)
+        for i, chunk in enumerate(chunks):
+            records.append({
+                "id": f"{fname}::chunk_{i+1}",
+                "text": chunk,
+                "metadata": {"source_file": fname, "model": model, "doc_type": doc_type}
+            })
+
+    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)
 
-    submit.click(fn=handle_query, inputs=question, outputs=answer)
+    return collection, embedder
 
-# ----------------------
-# Startup
-# ----------------------
-status_text = "Embedding manuals and loading model..."
+# ---------------------------
+# 💬 Load HF Model
+# ---------------------------
+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)
+
+# ---------------------------
+# 🔎 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
+# ---------------------------
 db, embedder = embed_all()
-llm = load_model()
-status_text = "Ready!"
+
+# ---------------------------
+# 🎛️ 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)
+
 demo.launch()