Update app.py

app.py

@@ -8,19 +8,20 @@ import torch
 import nltk
 import traceback
 import docx2txt
-import logging
 from PIL import Image
 from io import BytesIO
 from tqdm import tqdm
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer, util
-from nltk.tokenize.punkt import PunktSentenceTokenizer, PunktTrainer
+from nltk.tokenize import sent_tokenize
 
-# ---------------- Logger Setup ----------------
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
-logger = logging.getLogger("SmartManuals")
+# Ensure punkt is downloaded
+try:
+    nltk.data.find("tokenizers/punkt")
+except LookupError:
+    nltk.download("punkt")
 
-# ---------------- Config ----------------
+# Configuration
 HF_TOKEN = os.getenv("HF_TOKEN")
 MANUALS_DIR = "Manuals"
 CHROMA_PATH = "chroma_store"
@@ -30,25 +31,18 @@ CHUNK_OVERLAP = 100
 MAX_CONTEXT_CHUNKS = 3
 MODEL_ID = "ibm-granite/granite-vision-3.2-2b"
 
+# Device selection
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ---------------- Sentence Tokenizer (Persistent) ----------------
-try:
-    nltk.data.find("tokenizers/punkt")
-except LookupError:
-    nltk.download("punkt")
-
-tokenizer_punkt = PunktSentenceTokenizer()
-
 # ---------------- Text Helpers ----------------
 def clean(text):
     return "\n".join([line.strip() for line in text.splitlines() if line.strip()])
 
 def split_sentences(text):
     try:
-        return tokenizer_punkt.tokenize(text)
-    except Exception as e:
-        logger.warning("Tokenizer fallback: simple split. Reason: %s", e)
+        return sent_tokenize(text)
+    except:
+        print("\u26a0\ufe0f Tokenizer fallback: simple split.")
         return text.split(". ")
 
 def split_chunks(sentences, max_tokens=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
@@ -80,71 +74,86 @@ def extract_pdf_text(path):
                 text = pytesseract.image_to_string(img)
             chunks.append((path, i + 1, clean(text)))
     except Exception as e:
-        logger.error("PDF read error [%s]: %s", path, e)
+        print("\u274c PDF read error:", path, e)
     return chunks
 
 def extract_docx_text(path):
     try:
         return [(path, 1, clean(docx2txt.process(path)))]
     except Exception as e:
-        logger.error("DOCX read error [%s]: %s", path, e)
+        print("\u274c DOCX read error:", path, e)
         return []
 
 # ---------------- Embedding ----------------
 def embed_all():
-    embedder = SentenceTransformer("all-MiniLM-L6-v2")
-    embedder.eval()
-    client = chromadb.PersistentClient(path=CHROMA_PATH)
+    try:
+        embedder = SentenceTransformer("all-MiniLM-L6-v2")
+        embedder.eval()
+    except Exception as e:
+        print("\u274c Failed to load SentenceTransformer:", e)
+        return None, None
 
     try:
+        client = chromadb.PersistentClient(path=CHROMA_PATH)
         client.delete_collection(COLLECTION_NAME)
-    except:
-        pass
-    collection = client.get_or_create_collection(COLLECTION_NAME)
+        collection = client.get_or_create_collection(COLLECTION_NAME)
+    except Exception as e:
+        print("\u274c Failed to initialize ChromaDB:", e)
+        return None, None
 
     docs, ids, metas = [], [], []
-    logger.info("📄 Processing manuals...")
-
-    for fname in os.listdir(MANUALS_DIR):
-        fpath = os.path.join(MANUALS_DIR, fname)
-        if fname.lower().endswith(".pdf"):
-            pages = extract_pdf_text(fpath)
-        elif fname.lower().endswith(".docx"):
-            pages = extract_docx_text(fpath)
-        else:
-            continue
-
-        for path, page, text in pages:
-            for i, chunk in enumerate(split_chunks(split_sentences(text))):
-                chunk_id = f"{fname}::{page}::{i}"
-                docs.append(chunk)
-                ids.append(chunk_id)
-                metas.append({"source": fname, "page": page})
-
-                if len(docs) >= 16:
-                    embs = embedder.encode(docs).tolist()
-                    collection.add(documents=docs, ids=ids, metadatas=metas, embeddings=embs)
-                    docs, ids, metas = [], [], []
-
-    if docs:
-        embs = embedder.encode(docs).tolist()
-        collection.add(documents=docs, ids=ids, metadatas=metas, embeddings=embs)
-
-    logger.info("✅ Embedded %d chunks.", len(ids))
-    return collection, embedder
+    print("\ud83d\udcc4 Processing manuals...")
+
+    try:
+        for fname in os.listdir(MANUALS_DIR):
+            fpath = os.path.join(MANUALS_DIR, fname)
+            if fname.lower().endswith(".pdf"):
+                pages = extract_pdf_text(fpath)
+            elif fname.lower().endswith(".docx"):
+                pages = extract_docx_text(fpath)
+            else:
+                continue
+
+            for path, page, text in pages:
+                for i, chunk in enumerate(split_chunks(split_sentences(text))):
+                    chunk_id = f"{fname}::{page}::{i}"
+                    docs.append(chunk)
+                    ids.append(chunk_id)
+                    metas.append({"source": fname, "page": page})
+
+                    if len(docs) >= 16:
+                        embs = embedder.encode(docs).tolist()
+                        collection.add(documents=docs, ids=ids, metadatas=metas, embeddings=embs)
+                        docs, ids, metas = [], [], []
+
+        if docs:
+            embs = embedder.encode(docs).tolist()
+            collection.add(documents=docs, ids=ids, metadatas=metas, embeddings=embs)
+
+        print(f"\u2705 Embedded {len(ids)} chunks.")
+        return collection, embedder
+
+    except Exception as e:
+        print("\u274c Error during embedding:", e)
+        return None, None
 
 # ---------------- Model Setup ----------------
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        token=HF_TOKEN,
-        device_map="auto" if torch.cuda.is_available() else None,
-        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
-    ).to(device)
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
-    return pipe, tokenizer
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            token=HF_TOKEN,
+            device_map="auto" if torch.cuda.is_available() else None,
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
+        ).to(device)
+        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
+        return pipe, tokenizer
+    except Exception as e:
+        print("\u274c Failed to load model:", e)
+        return None, None
 
+# ---------------- QA Logic ----------------
 def ask_model(question, context, pipe, tokenizer):
     prompt = f"""Use only the following context to answer. If uncertain, say \"I don't know.\"
 
@@ -157,37 +166,48 @@ A:"""
     output = pipe(prompt, max_new_tokens=512)[0]["generated_text"]
     return output.split("A:")[-1].strip()
 
-# ---------------- Query ----------------
 def get_answer(question):
+    if not all([embedder, db, model_pipe, model_tokenizer]):
+        return "\u274c System not initialized. Check logs or try restarting the app."
     try:
-        query_emb = embedder.encode(question, convert_to_tensor=True)
         results = db.query(query_texts=[question], n_results=MAX_CONTEXT_CHUNKS)
         context = "\n\n".join(results["documents"][0])
-        source_info = "\n\n".join([
-            f"📄 Source: {m.get('source', 'N/A')} (Page {m.get('page', 'N/A')})" for m in results["metadatas"][0]
-        ])
-        answer = ask_model(question, context, model_pipe, model_tokenizer)
-        return f"{answer}\n\n---\n{source_info}"
+        return ask_model(question, context, model_pipe, model_tokenizer)
     except Exception as e:
-        logger.error("❌ Query error: %s", e)
+        print("\u274c Query error:", e)
         return f"Error: {e}"
 
 # ---------------- UI ----------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 🤖 SmartManuals-AI (Granite 3.2-2B)")
+    gr.Markdown("## \ud83e\udd16 SmartManuals-AI (Granite 3.2-2B)")
     with gr.Row():
         question = gr.Textbox(label="Ask your question")
         ask = gr.Button("Ask")
-    answer = gr.Textbox(label="Answer", lines=10)
-    ask.click(fn=get_answer, inputs=question, outputs=answer)
+    answer = gr.Textbox(label="Answer", lines=8)
+    status = gr.Markdown(visible=False)
+
+    def wrapped_get_answer(q):
+        ans = get_answer(q)
+        return ans, ""  # hide status after success
+
+    ask.click(fn=wrapped_get_answer, inputs=question, outputs=[answer, status])
+
+    # Show status on startup error
+    if not all([embedder, db, model_pipe, model_tokenizer]):
+        status.visible = True
+        status.value = "\u26a0\ufe0f Initialization failed. Check logs or your HF_TOKEN."
 
 # Embed + Load Model at Startup
 try:
     db, embedder = embed_all()
-    model_pipe, model_tokenizer = load_model()
 except Exception as e:
-    logger.exception("❌ Startup failure: %s", e)
+    print("\u274c Embedding failed:", e)
     db, embedder = None, None
+
+try:
+    model_pipe, model_tokenizer = load_model()
+except Exception as e:
+    print("\u274c Model loading failed:", e)
     model_pipe, model_tokenizer = None, None
 
 if __name__ == "__main__":