Update app.py

app.py

@@ -1,30 +1,52 @@
+
 import os
+import fitz
 import json
-import fitz  # PyMuPDF
+import gradio as gr
 import pytesseract
-from PIL import Image
-import io
-import nltk
 import chromadb
+import torch
+import asyncio
+import docx2txt
+import nltk
+import traceback
+from PIL import Image
+from io import BytesIO
 from tqdm import tqdm
-from sentence_transformers import SentenceTransformer
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-import gradio as gr
-
-# ---------------------------
-# 📦 Paths and Constants
-# ---------------------------
-MANUALS_DIR = "./Manuals"
-CHROMA_PATH = "./chroma_store"
-COLLECTION_NAME = "manual_chunks"
-
-# Ensure NLTK punkt is available
-nltk.download("punkt")
+from transformers import (
+    pipeline,
+    AutoModelForCausalLM,
+    AutoTokenizer
+)
+from sentence_transformers import SentenceTransformer, util
 from nltk.tokenize import sent_tokenize
 
-# ---------------------------
-# 🧼 Text cleaning utilities
-# ---------------------------
+# Ensure punkt is available
+try:
+    nltk.data.find("tokenizers/punkt")
+except LookupError:
+    nltk.download("punkt")
+
+# ---------------- Config ----------------
+MANUALS_DIR = "Manuals"
+CHROMA_PATH = "chroma_store"
+COLLECTION_NAME = "manual_chunks"
+CHUNK_SIZE = 750
+CHUNK_OVERLAP = 100
+MAX_CONTEXT_CHUNKS = 3
+
+MODELS = {
+    "LLaMA 3 (8B)": "meta-llama/Llama-3.1-8B-Instruct",
+    "Mistral 7B": "mistralai/Mistral-7B-Instruct-v0.3",
+    "Gemma 2B": "google/gemma-1.1-2b-it",
+    "LLaMA 4 (Scout 17B)": "meta-llama/Llama-4-Scout-17B-16E",
+    "Qwen 30B": "Qwen/Qwen3-30B-A3B"
+}
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ---------------- Utils ----------------
 def clean(text):
     return "\n".join([line.strip() for line in text.splitlines() if line.strip()])
 
@@ -32,110 +54,150 @@ def split_sentences(text):
     try:
         return sent_tokenize(text)
     except Exception as e:
-        print("[Tokenizer Error]", e, "\nFalling back to simple split.")
+        print("[Tokenizer Error]", e)
         return text.split(". ")
 
-# ---------------------------
-# 📄 PDF and DOCX extraction
-# ---------------------------
+def split_into_chunks(sentences, max_tokens=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
+    chunks = []
+    current_chunk, current_len = [], 0
+
+    for sentence in sentences:
+        words = sentence.split()
+        if current_len + len(words) > max_tokens and current_chunk:
+            chunks.append(" ".join(current_chunk))
+            current_chunk = current_chunk[-overlap:]
+            current_len = sum(len(s.split()) for s in current_chunk)
+
+        current_chunk.append(sentence)
+        current_len += len(words)
+
+    if current_chunk:
+        chunks.append(" ".join(current_chunk))
+
+    return chunks
+
 def extract_pdf_text(pdf_path):
-    doc = fitz.open(pdf_path)
-    pages = []
-    for i, page in enumerate(doc):
-        text = page.get_text().strip()
-        if not text:
-            try:
+    text_chunks = []
+    try:
+        doc = fitz.open(pdf_path)
+        for i, page in enumerate(doc):
+            text = page.get_text().strip()
+            if not text:
                 pix = page.get_pixmap(dpi=300)
-                img = Image.open(io.BytesIO(pix.tobytes("png")))
+                img = Image.open(BytesIO(pix.tobytes("png")))
                 text = pytesseract.image_to_string(img)
-            except pytesseract.TesseractNotFoundError:
-                print("❌ Tesseract not found. Skipping OCR for page.")
-                text = ""
-        pages.append((i + 1, text))
-    return pages
-
-# ---------------------------
-# 🧠 Embed text using MiniLM
-# ---------------------------
+            text_chunks.append((pdf_path, i + 1, clean(text)))
+    except Exception as e:
+        print("❌ Error reading PDF:", pdf_path, e)
+    return text_chunks
+
+def extract_docx_text(docx_path):
+    try:
+        text = clean(docx2txt.process(docx_path))
+        return [(docx_path, 1, text)]
+    except Exception as e:
+        print("❌ Error reading DOCX:", docx_path, e)
+        return []
+
+# ---------------- Background Embed ----------------
 def embed_all():
+    embedder = SentenceTransformer("all-MiniLM-L6-v2")
+    embedder.eval()
     client = chromadb.PersistentClient(path=CHROMA_PATH)
-    if COLLECTION_NAME in [c.name for c in client.list_collections()]:
+
+    try:
         client.delete_collection(COLLECTION_NAME)
-    collection = client.create_collection(COLLECTION_NAME)
+    except:
+        pass
+    collection = client.get_or_create_collection(COLLECTION_NAME)
 
-    embedder = SentenceTransformer("all-MiniLM-L6-v2")
+    chunks, ids, metas = [], [], []
+    idx = 0
+    print("📄 Scanning Manuals folder...")
 
-    chunk_id = 0
     for fname in os.listdir(MANUALS_DIR):
         fpath = os.path.join(MANUALS_DIR, fname)
         if fname.lower().endswith(".pdf"):
             pages = extract_pdf_text(fpath)
-            for page_num, text in pages:
-                sents = split_sentences(clean(text))
-                for i in range(0, len(sents), 5):
-                    chunk = " ".join(sents[i:i + 5])
-                    if chunk.strip():
-                        collection.add(
-                            documents=[chunk],
-                            metadatas=[{"source": fname, "page": page_num}],
-                            ids=[f"{fname}-{page_num}-{i}-{chunk_id}"]
-                        )
-                        chunk_id += 1
-
-    print(f"✅ Embedded {chunk_id} chunks.")
+        elif fname.lower().endswith(".docx"):
+            pages = extract_docx_text(fpath)
+        else:
+            continue
+
+        for filepath, page, text in pages:
+            sentences = split_sentences(text)
+            subchunks = split_into_chunks(sentences)
+            for i, subchunk in enumerate(subchunks):
+                chunks.append(subchunk)
+                ids.append(f"{fname}::{page}::{i}")
+                metas.append({"source": fname, "page": page})
+
+                if len(chunks) >= 16:
+                    embs = embedder.encode(chunks).tolist()
+                    collection.add(documents=chunks, ids=ids, metadatas=metas, embeddings=embs)
+                    chunks, ids, metas = [], [], []
+
+    if chunks:
+        embs = embedder.encode(chunks).tolist()
+        collection.add(documents=chunks, ids=ids, metadatas=metas, embeddings=embs)
+
+    print(f"✅ Embedded {len(ids)} chunks.")
     return collection, embedder
 
-# ---------------------------
-# 🤖 Load model
-# ---------------------------
-def load_llm():
-    model_id = "meta-llama/Llama-3.1-8B-Instruct"
-    token = os.environ.get("HF_TOKEN")
-    tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id, token=token, torch_dtype=None, device_map="auto"
-    )
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
+# ---------------- Model Loader ----------------
+def load_model(model_id):
+    tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
+    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, token=HF_TOKEN)
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
     return pipe, tokenizer
 
-# ---------------------------
-# ❓ Ask a question
-# ---------------------------
-def ask_question(question, db, embedder, pipe, tokenizer):
-    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 that answers questions from technical manuals using only the provided context.
-<context>
+# ---------------- Query ----------------
+def query_llm(context, question, pipe, tokenizer):
+    prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+You are a helpful assistant. Use only the following context to answer. If uncertain, say: 'I don't know.'
+
 {context}
-</context>
 <|start_header_id|>user<|end_header_id|>
-{question}<|start_header_id|>assistant<|end_header_id|>
+{question}
+<|start_header_id|>assistant<|end_header_id|>
 """
+    out = pipe(prompt, max_new_tokens=512)[0]["generated_text"]
+    return out.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
+
+def answer_question(question, model_choice):
+    try:
+        model_id = MODELS[model_choice]
+        pipe, tokenizer = load_model(model_id)
+        query_emb = embedder.encode(question, convert_to_tensor=True)
+
+        results = db.query(query_texts=[question], n_results=MAX_CONTEXT_CHUNKS)
+        context_chunks = results["documents"][0]
+        context = "\n\n".join(context_chunks)
 
-    out = pipe(prompt)[0]["generated_text"]
-    final = out.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
-    return final
+        answer = query_llm(context, question, pipe, tokenizer)
+        return answer
+    except Exception as e:
+        traceback.print_exc()
+        return f"❌ Error: {str(e)}"
 
-# ---------------------------
-# 🚀 Build interface
-# ---------------------------
+# ---------------- Run App ----------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 SmartManuals-AI (Hugging Face Space Edition)")
+    gr.Markdown("### 📘 Ask Questions About Your Manuals")
 
-    with gr.Row():
-        qbox = gr.Textbox(label="Ask a Question", placeholder="e.g. How do I access diagnostics on the SE3 console?")
-        submit = gr.Button("🔍 Ask")
+    model_choice = gr.Dropdown(label="Select Model", choices=list(MODELS.keys()), value="LLaMA 3 (8B)")
+    question = gr.Textbox(label="Your Question", placeholder="e.g. How do I reset the treadmill?")
+    submit = gr.Button("🔍 Get Answer")
+    answer = gr.Textbox(label="Answer", lines=10)
 
-    abox = gr.Textbox(label="Answer", lines=8)
+    submit.click(fn=answer_question, inputs=[question, model_choice], outputs=answer)
 
+# Run background embed on startup
+try:
     db, embedder = embed_all()
-    pipe, tokenizer = load_llm()
-
-    submit.click(fn=lambda q: ask_question(q, db, embedder, pipe, tokenizer), inputs=qbox, outputs=abox)
+except Exception as e:
+    print("❌ Failed to embed docs:", e)
+    db, embedder = None, None
 
-# For Hugging Face Spaces
+# Only launch if in HF Space
 if __name__ == "__main__":
     demo.launch()