Create builder.py

builder.py

@@ -0,0 +1,159 @@
+import os, uuid, shutil
+from pathlib import Path
+from huggingface_hub import snapshot_download, upload_folder, HfApi
+
+# ---------- Config via secrets ----------
+HF_TOKEN  = os.getenv("HF_TOKEN")  # MUST be set
+CORPUS_DS = os.getenv("CORPUS_DS", "Azizahalq/materialmind-corpus")
+INDEX_DS  = os.getenv("INDEX_DS",  "Azizahalq/materialmind-index")
+
+ROOT       = Path(__file__).parent.resolve()
+MM_ROOT    = ROOT / "MaterialMind"
+SRC_DIR    = MM_ROOT / "sources"
+INDEX_BASE = MM_ROOT / "index" / "chroma_v3"  # we’ll create a <uuid> subdir here
+
+EMB_MODEL  = "BAAI/bge-small-en-v1.5"
+
+def log(*a): print(*a, flush=True)
+
+def ensure_dirs():
+    SRC_DIR.mkdir(parents=True, exist_ok=True)
+    INDEX_BASE.mkdir(parents=True, exist_ok=True)
+
+def download_corpus():
+    log("[Step] Downloading corpus dataset:", CORPUS_DS)
+    snapshot_download(repo_id=CORPUS_DS, repo_type="dataset",
+                      local_dir=str(SRC_DIR), local_dir_use_symlinks=False)
+    log("[OK] Corpus ready at", SRC_DIR)
+
+def build_index():
+    # Lazy embedder (FastEmbed -> ST)
+    try:
+        from fastembed import TextEmbedding
+        embedder = TextEmbedding(model_name=EMB_MODEL)
+        def embed(texts): return [v for v in embedder.embed(texts)]
+        log("[EMB] FastEmbed:", EMB_MODEL)
+    except Exception as e:
+        from sentence_transformers import SentenceTransformer
+        model = SentenceTransformer(EMB_MODEL)
+        def embed(texts): return model.encode(texts, normalize_embeddings=True).tolist()
+        log("[EMB] ST fallback:", EMB_MODEL, e)
+
+    # Readers
+    import re
+    def norm(s):
+        s = s.replace("\r","\n")
+        s = re.sub(r"[ \t]+"," ",s)
+        s = re.sub(r"\n{3,}","\n\n",s)
+        return s.strip()
+
+    def from_pdf(path:Path):
+        any_text=False
+        try:
+            import fitz
+            doc=fitz.open(str(path))
+            for i,p in enumerate(doc):
+                t=p.get_text("text").strip()
+                if t:
+                    any_text=True
+                    yield norm(t), i+1
+            doc.close()
+        except Exception:
+            pass
+        if not any_text:
+            try:
+                from pypdf import PdfReader
+                r=PdfReader(str(path))
+                for i,p in enumerate(r.pages):
+                    try: raw=p.extract_text() or ""
+                    except: raw=""
+                    t=norm(raw)
+                    if t:
+                        any_text=True
+                        yield t, i+1
+            except Exception as e:
+                log("[WARN] pdf read fail:", path.name, e)
+        if not any_text:
+            log("[HINT] no extractable text:", path.name)
+
+    def chunk(text, max_chars=1200, overlap=150):
+        n=len(text); 
+        if n<=max_chars:
+            if n>0: yield text
+            return
+        i=0
+        while i<n:
+            j=min(i+max_chars,n)
+            yield text[i:j]
+            i = j-overlap if j<n else j
+
+    # Build Chroma catalog under a fresh UUID directory
+    cat_dir = INDEX_BASE / str(uuid.uuid4())
+    cat_dir.mkdir(parents=True, exist_ok=True)
+    log("[Step] Building Chroma catalog at:", cat_dir)
+
+    import chromadb
+    client = chromadb.PersistentClient(path=str(cat_dir))
+    col = client.get_or_create_collection(name="materialmind")
+
+    # iterate files
+    batch_ids, batch_docs, batch_meta = [], [], []
+    def flush():
+        if not batch_ids: return
+        embs = embed(batch_docs)
+        col.add(ids=batch_ids, documents=batch_docs, metadatas=batch_meta, embeddings=embs)
+        batch_ids.clear(); batch_docs.clear(); batch_meta.clear()
+
+    added = 0
+    for f in SRC_DIR.rglob("*"):
+        if not f.is_file(): 
+            continue
+        if f.suffix.lower() != ".pdf":
+            continue
+        rel = f.relative_to(MM_ROOT).as_posix()
+        for page_text, page in from_pdf(f):
+            for c in chunk(page_text):
+                batch_ids.append(str(uuid.uuid4()))
+                batch_docs.append(c)
+                batch_meta.append({"source": rel, "page": page})
+                if len(batch_ids) >= 256:
+                    flush()
+            added += 1
+            if added % 200 == 0:
+                log(f"  +{added} chunks...")
+
+    flush()
+    log("[OK] Built. Total chunks ~", col.count())
+    return cat_dir  # MaterialMind/index/chroma_v3/<uuid>
+
+def upload_catalog(cat_dir:Path):
+    # Upload to dataset INDEX_DS under path: index/chroma_v3/<uuid>
+    # (the app will snapshot_download INDEX_DS later)
+    target_path_in_repo = f"index/chroma_v3/{cat_dir.name}"
+    log("[Step] Uploading catalog to dataset:", INDEX_DS, "at", target_path_in_repo)
+    api = HfApi(token=HF_TOKEN)
+    upload_folder(
+        repo_id=INDEX_DS,
+        repo_type="dataset",
+        path_in_repo=target_path_in_repo,
+        folder_path=str(cat_dir),
+        token=HF_TOKEN,
+        allow_patterns=None,
+        ignore_patterns=["**/__pycache__/**"],
+    )
+    log("[OK] Uploaded.")
+    log("NOTE: set Space secret INDEX_DS =", INDEX_DS)
+    log("      optional INDEX_DIR = MaterialMind/index/chroma_v3/" + cat_dir.name)
+
+def run():
+    print("==== MaterialMind Index Builder ====")
+    if not HF_TOKEN:
+        raise RuntimeError("HF_TOKEN secret is required.")
+    ensure_dirs()
+    download_corpus()
+    cat_dir = build_index()
+    upload_catalog(cat_dir)
+    print("==== Done. You can stop this Space. ====")
+
+if __name__ == "__main__":
+    run()