Update app.py

app.py

@@ -41,7 +41,7 @@ def split_sentences(text):
     try:
         return sent_tokenize(text)
     except:
-        print("\u26a0\ufe0f Tokenizer fallback: simple split.")
+        print("Tokenizer fallback: simple split.")
         return text.split(". ")
 
 def split_chunks(sentences, max_tokens=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
@@ -73,70 +73,80 @@ def extract_pdf_text(path):
                 text = pytesseract.image_to_string(img)
             chunks.append((path, i + 1, clean(text)))
     except Exception as e:
-        print("\u274c PDF read error:", path, e)
+        print("PDF read error:", path, e)
     return chunks
 
 def extract_docx_text(path):
     try:
         return [(path, 1, clean(docx2txt.process(path)))]
     except Exception as e:
-        print("\u274c DOCX read error:", path, e)
+        print("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:
-        client.delete_collection(COLLECTION_NAME)
-    except:
-        pass
-    collection = client.get_or_create_collection(COLLECTION_NAME)
-
-    docs, ids, metas = [], [], []
-    print("\ud83d\udcc4 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)
-
-    print(f"\u2705 Embedded {len(ids)} chunks.")
-    return collection, embedder
+        embedder = SentenceTransformer("all-MiniLM-L6-v2")
+        embedder.eval()
+        client = chromadb.PersistentClient(path=CHROMA_PATH)
+
+        try:
+            client.delete_collection(COLLECTION_NAME)
+        except:
+            pass
+
+        collection = client.get_or_create_collection(COLLECTION_NAME)
+
+        docs, ids, metas = [], [], []
+        print("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)
+
+        print(f"Embedded {len(ids)} chunks.")
+        return collection, embedder
+
+    except Exception as e:
+        print("Embedding startup failed:", 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("Model loading failed:", e)
+        return None, None
 
 def ask_model(question, context, pipe, tokenizer):
     prompt = f"""Use only the following context to answer. If uncertain, say \"I don't know.\"
@@ -152,38 +162,41 @@ A:"""
 
 # ---------------- Query ----------------
 def get_answer(question):
-    if not all([embedder, db, model_pipe, model_tokenizer]):
-        return "⚠️ The system is still initializing or failed to load. Please try again later."
+    if not embedder or not db or not model_pipe:
+        return "System not ready. Try again after initialization."
     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])
         return ask_model(question, context, model_pipe, model_tokenizer)
     except Exception as e:
-        print("\u274c Query error:", e)
+        print("Query error:", e)
         return f"Error: {e}"
 
 # ---------------- UI ----------------
 with gr.Blocks() as demo:
-    gr.Markdown("## \ud83e\udd16 SmartManuals-AI (Granite 3.2-2B)")
+    gr.Markdown("## 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=8)
     ask.click(fn=get_answer, inputs=question, outputs=answer)
 
-# ---------------- Startup ----------------
-embedder = db = model_pipe = model_tokenizer = None
+# Startup Initialization
+embedder = None
+model_pipe = None
+model_tokenizer = None
 
 try:
     db, embedder = embed_all()
 except Exception as e:
-    print("\u274c Embedding failed:", e)
+    print("❌ Embedding failed:", e)
 
 try:
     model_pipe, model_tokenizer = load_model()
 except Exception as e:
-    print("\u274c Model loading failed:", e)
+    print("❌ Model load failed:", e)
 
+# Launch
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)