Update app.py

app.py

@@ -4,77 +4,95 @@ import os
 import pandas as pd
 from sentence_transformers import SentenceTransformer
 
-# Load model from HF
+from pypdf import PdfReader
+import docx
+
 model = SentenceTransformer("csAhmad/zoraiz-model")
 
-# Ensure temp directory exists
 EXTRACT_PATH = "temp/extracted"
-os.makedirs(EXTRACT_PATH, exist_ok=True)
 
 
+# -------------------------
+# TEXT EXTRACTORS
+# -------------------------
+def extract_text(file_path):
+    ext = file_path.lower()
+
+    # TXT
+    if ext.endswith(".txt"):
+        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+            return f.read()
+
+    # PDF
+    elif ext.endswith(".pdf"):
+        reader = PdfReader(file_path)
+        return " ".join([page.extract_text() or "" for page in reader.pages])
+
+    # DOCX
+    elif ext.endswith(".docx"):
+        doc = docx.Document(file_path)
+        return "\n".join([para.text for para in doc.paragraphs])
+
+    return None
+
+
+# -------------------------
+# MAIN FUNCTION
+# -------------------------
 def process_zip(zip_file):
-    # ---------------------------
-    # 1. Validate input
-    # ---------------------------
+
     if zip_file is None:
         raise gr.Error("Please upload a ZIP file.")
 
-    # reset extraction folder
+    # clean folder
     if os.path.exists(EXTRACT_PATH):
-        for f in os.listdir(EXTRACT_PATH):
-            try:
-                os.remove(os.path.join(EXTRACT_PATH, f))
-            except:
-                pass
+        for root, dirs, files in os.walk(EXTRACT_PATH):
+            for f in files:
+                try:
+                    os.remove(os.path.join(root, f))
+                except:
+                    pass
 
     os.makedirs(EXTRACT_PATH, exist_ok=True)
 
     zip_path = zip_file.name
 
-    # ---------------------------
-    # 2. Extract ZIP safely
-    # ---------------------------
+    # extract ZIP
     try:
         with zipfile.ZipFile(zip_path, "r") as zip_ref:
             zip_ref.extractall(EXTRACT_PATH)
     except zipfile.BadZipFile:
-        raise gr.Error("Invalid ZIP file uploaded.")
+        raise gr.Error("Invalid ZIP file.")
 
-    # ---------------------------
-    # 3. Process files
-    # ---------------------------
     results = []
 
+    # recursive scan
     for root, dirs, files in os.walk(EXTRACT_PATH):
         for file in files:
             file_path = os.path.join(root, file)
 
-            # only process text files
-            if file.lower().endswith(".txt"):
-                with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
-                    text = f.read()
+            text = extract_text(file_path)
 
-                if not text.strip():
-                    continue
+            if not text or not text.strip():
+                continue
 
+            try:
                 emb = model.encode(text)
 
                 results.append({
-                    "file": file,
+                    "file": os.path.relpath(file_path, EXTRACT_PATH),
+                    "type": file.split(".")[-1],
                     "text_length": len(text),
                     "embedding_dim": len(emb),
                     "preview_embedding": str(emb[:10])
                 })
 
-    # ---------------------------
-    # 4. Handle empty results
-    # ---------------------------
+            except Exception as e:
+                print(f"Embedding error: {file_path} -> {e}")
+
     if len(results) == 0:
-        raise gr.Error("No valid .txt files found inside ZIP.")
+        raise gr.Error("No readable TXT, PDF, or DOCX files found in ZIP.")
 
-    # ---------------------------
-    # 5. Create Excel file
-    # ---------------------------
     df = pd.DataFrame(results)
 
     output_file = "output.xlsx"
@@ -83,15 +101,15 @@ def process_zip(zip_file):
     return output_file
 
 
-# ---------------------------
-# 6. Gradio UI
-# ---------------------------
+# -------------------------
+# GRADIO UI
+# -------------------------
 demo = gr.Interface(
     fn=process_zip,
     inputs=gr.File(file_types=[".zip"]),
     outputs=gr.File(label="Download Excel"),
-    title="ZIP → Embeddings → Excel Generator",
-    description="Upload a ZIP containing .txt files. It will generate embeddings using SentenceTransformer and export to Excel."
+    title="ZIP → Multiformat Embeddings",
+    description="Upload ZIP containing TXT, PDF, DOCX (even in folders). Generates embeddings and exports Excel."
 )
 
 demo.launch()