Update app.py

app.py

@@ -31,7 +31,6 @@ 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"
 
 # ---------------- Text Helpers ----------------
@@ -86,74 +85,59 @@ def extract_docx_text(path):
 
 # ---------------- Embedding ----------------
 def embed_all():
-    try:
-        embedder = SentenceTransformer("all-MiniLM-L6-v2")
-        embedder.eval()
-    except Exception as e:
-        print("\u274c Failed to load SentenceTransformer:", e)
-        return None, None
+    embedder = SentenceTransformer("all-MiniLM-L6-v2")
+    embedder.eval()
+    client = chromadb.PersistentClient(path=CHROMA_PATH)
 
     try:
-        client = chromadb.PersistentClient(path=CHROMA_PATH)
         client.delete_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
+    except:
+        pass
+    collection = client.get_or_create_collection(COLLECTION_NAME)
 
     docs, ids, metas = [], [], []
     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
+    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
 
 # ---------------- Model Setup ----------------
 def load_model():
-    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
+    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
 
-# ---------------- 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.\"
 
@@ -166,10 +150,12 @@ 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."
+        return "⚠️ The system is still initializing or failed to load. Please try again later."
     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)
@@ -184,31 +170,20 @@ with gr.Blocks() as demo:
         question = gr.Textbox(label="Ask your question")
         ask = gr.Button("Ask")
     answer = gr.Textbox(label="Answer", lines=8)
-    status = gr.Markdown(visible=False)
+    ask.click(fn=get_answer, inputs=question, outputs=answer)
 
-    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."
+# ---------------- Startup ----------------
+embedder = db = model_pipe = model_tokenizer = None
 
-# Embed + Load Model at Startup
 try:
     db, embedder = embed_all()
 except Exception as 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__":
     demo.launch()