Update app.py

app.py

@@ -1,153 +1,144 @@
-# ✅ SmartManuals-AI: Hugging Face Space App (RAM Safe, Multi-model, No Preview)
-
-import os, json, fitz, torch, chromadb, docx
+import os
+import json
+import fitz  # PyMuPDF
+import docx
+import chromadb
+import torch
+import nltk
 import gradio as gr
+from tqdm import tqdm
+from typing import List
 from PIL import Image
 from nltk.tokenize import sent_tokenize
 from sentence_transformers import SentenceTransformer, util
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from tqdm import tqdm
 
-# ---------------------------
-# ⚙️ Constants
-# ---------------------------
-MANUALS_DIR = "Manuals"
+# --- Ensure punkt tokenizer is available ---
+try:
+    nltk.data.find("tokenizers/punkt")
+except LookupError:
+    nltk.download("punkt")
+
+# --- Configuration ---
+MANUALS_FOLDER = "./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
-TOP_K = 3
-
 MODEL_OPTIONS = {
-    "LLaMA 3.1 (8B)": "meta-llama/Llama-3.1-8B-Instruct",
+    "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"
+    "Gemma 7B": "google/gemma-1.1-7b-it"
 }
+HF_TOKEN = os.environ.get("HF_TOKEN")
+MAX_CONTEXT_CHUNKS = 3
 
-# ---------------------------
-# 📄 Extract Text from PDFs and DOCX
-# ---------------------------
+# --- Utility Functions ---
 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
+        return "\n".join([page.get_text().strip() for page in doc])
+    except:
+        return ""
 
 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 "\n".join([para.text.strip() for para in doc.paragraphs])
+    except:
         return ""
 
-# ---------------------------
-# 🧹 Clean + Chunk
-# ---------------------------
 def clean(text):
     return "\n".join([line.strip() for line in text.splitlines() if line.strip()])
 
 def split_sentences(text):
-    return sent_tokenize(text)
-
-def chunk_text(sentences, size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
-    chunks, chunk, length = [], [], 0
-    for sent in sentences:
-        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 += n
-    if chunk:
-        chunks.append(" ".join(chunk))
+    try:
+        return sent_tokenize(text)
+    except Exception as e:
+        print(f"[Tokenizer Error] {e}. Falling back to simple split.")
+        return text.split(". ")
+
+def chunk_sentences(sentences, max_tokens=500, overlap=50):
+    chunks = []
+    current = []
+    total = 0
+    for sentence in sentences:
+        count = len(sentence.split())
+        if total + count > max_tokens:
+            chunks.append(" ".join(current))
+            current = current[-overlap:]
+            total = sum(len(s.split()) for s in current)
+        current.append(sentence)
+        total += count
+    if current:
+        chunks.append(" ".join(current))
     return chunks
 
-# ---------------------------
-# 📦 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)
+    db = chromadb.PersistentClient(path=CHROMA_PATH)
+    if COLLECTION_NAME in [c.name for c in db.list_collections()]:
+        db.delete_collection(COLLECTION_NAME)
+    collection = db.create_collection(COLLECTION_NAME)
 
+    embedder = SentenceTransformer("all-MiniLM-L6-v2")
     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)
-        text = clean(text)
-        sents = split_sentences(text)
-        chunks = chunk_text(sents)
-        for i, chunk in enumerate(chunks):
-            all_chunks.append({
-                "id": f"{fname}::chunk_{i+1}",
-                "text": chunk,
-                "metadata": {"source": fname}
-            })
-
-    # Batch embed and store
+
+    for fname in os.listdir(MANUALS_FOLDER):
+        path = os.path.join(MANUALS_FOLDER, fname)
+        text = ""
+        if fname.lower().endswith(".pdf"):
+            text = extract_text_from_pdf(path)
+        elif fname.lower().endswith(".docx"):
+            text = extract_text_from_docx(path)
+        else:
+            continue
+
+        sents = split_sentences(clean(text))
+        chunks = chunk_sentences(sents)
+        for idx, chunk in enumerate(chunks):
+            chunk_id = f"{fname}::chunk_{idx}"
+            all_chunks.append({"id": chunk_id, "text": chunk, "metadata": {"source": fname}})
+
     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]
+        docs = [x["text"] for x in batch]
+        ids = [x["id"] for x in batch]
+        metas = [x["metadata"] for x 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
 
-# ---------------------------
-# 🔍 RAG Search & LLM Answer
-# ---------------------------
-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}"
+def answer_query(query, model_choice):
+    db, embedder = embed_all()
+    results = db.get_collection(COLLECTION_NAME).query(query_texts=[query], n_results=MAX_CONTEXT_CHUNKS)
 
-    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:"
+    context = "\n\n".join(results["documents"][0])
+    model_id = MODEL_OPTIONS.get(model_choice)
 
-    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}"
+    tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
+    model = AutoModelForCausalLM.from_pretrained(model_id, token=HF_TOKEN)
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+    prompt = f"""
+    Context:
+    {context}
 
-# ---------------------------
-# ▶️ UI
-# ---------------------------
-db, embedder = embed_all()
+    Question: {query}
+    Answer:"""
+
+    out = pipe(prompt, max_new_tokens=300, do_sample=False)
+    return out[0]["generated_text"].split("Answer:")[-1].strip()
+
+# --- UI ---
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 SmartManuals-AI — Ask Your PDF and Word Docs")
+    gr.Markdown("""# 📘 SmartManuals-AI
+    Ask technical questions from manuals (PDF & DOCX) with LLM + OCR + RAG.
+    """)
+
     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])
+        question = gr.Textbox(label="Your Question", placeholder="e.g., How do I reset the console?")
+        model_choice = gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), value="LLaMA 3.1 8B", label="Model")
+    answer = gr.Textbox(label="Answer")
+    submit = gr.Button("Ask")
+    submit.click(fn=answer_query, inputs=[question, model_choice], outputs=answer)
 
 demo.launch()