Update app.py

app.py

@@ -8,20 +8,19 @@ 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 import sent_tokenize
+from nltk.tokenize.punkt import PunktSentenceTokenizer, PunktTrainer
 
-# Ensure punkt is downloaded
-try:
-    nltk.data.find("tokenizers/punkt")
-except LookupError:
-    nltk.download("punkt")
+# ---------------- Logger Setup ----------------
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger("SmartManuals")
 
-# Configuration
+# ---------------- Config ----------------
 HF_TOKEN = os.getenv("HF_TOKEN")
 MANUALS_DIR = "Manuals"
 CHROMA_PATH = "chroma_store"
@@ -33,15 +32,23 @@ MODEL_ID = "ibm-granite/granite-vision-3.2-2b"
 
 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 sent_tokenize(text)
-    except:
-        print("⚠️ Tokenizer fallback: simple split.")
+        return tokenizer_punkt.tokenize(text)
+    except Exception as e:
+        logger.warning("Tokenizer fallback: simple split. Reason: %s", e)
         return text.split(". ")
 
 def split_chunks(sentences, max_tokens=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
@@ -73,14 +80,14 @@ def extract_pdf_text(path):
                 text = pytesseract.image_to_string(img)
             chunks.append((path, i + 1, clean(text)))
     except Exception as e:
-        print("❌ PDF read error:", path, e)
+        logger.error("PDF read error [%s]: %s", path, e)
     return chunks
 
 def extract_docx_text(path):
     try:
         return [(path, 1, clean(docx2txt.process(path)))]
     except Exception as e:
-        print("❌ DOCX read error:", path, e)
+        logger.error("DOCX read error [%s]: %s", path, e)
         return []
 
 # ---------------- Embedding ----------------
@@ -96,7 +103,7 @@ def embed_all():
     collection = client.get_or_create_collection(COLLECTION_NAME)
 
     docs, ids, metas = [], [], []
-    print("📄 Processing manuals...")
+    logger.info("📄 Processing manuals...")
 
     for fname in os.listdir(MANUALS_DIR):
         fpath = os.path.join(MANUALS_DIR, fname)
@@ -123,7 +130,7 @@ def embed_all():
         embs = embedder.encode(docs).tolist()
         collection.add(documents=docs, ids=ids, metadatas=metas, embeddings=embs)
 
-    print(f"✅ Embedded {len(ids)} chunks.")
+    logger.info("✅ Embedded %d chunks.", len(ids))
     return collection, embedder
 
 # ---------------- Model Setup ----------------
@@ -156,9 +163,13 @@ def get_answer(question):
         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])
-        return ask_model(question, context, model_pipe, model_tokenizer)
+        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}"
     except Exception as e:
-        print("❌ Query error:", e)
+        logger.error("❌ Query error: %s", e)
         return f"Error: {e}"
 
 # ---------------- UI ----------------
@@ -167,7 +178,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         question = gr.Textbox(label="Ask your question")
         ask = gr.Button("Ask")
-    answer = gr.Textbox(label="Answer", lines=8)
+    answer = gr.Textbox(label="Answer", lines=10)
     ask.click(fn=get_answer, inputs=question, outputs=answer)
 
 # Embed + Load Model at Startup
@@ -175,7 +186,7 @@ try:
     db, embedder = embed_all()
     model_pipe, model_tokenizer = load_model()
 except Exception as e:
-    print("❌ Startup failure:", e)
+    logger.exception("❌ Startup failure: %s", e)
     db, embedder = None, None
     model_pipe, model_tokenizer = None, None