Upload folder using huggingface_hub

modal_train.py

@@ -52,6 +52,7 @@ def train_on_l4(hours: float = 20.0, ckpt_every: int = 250,
                 size: str = "large", batch: int = None, scale: float = 3.0):
     import subprocess
     import sys
+    import shutil
 
     # ---- pull latest code from HF (use huggingface_hub, not the `hf` CLI) ----
     from huggingface_hub import snapshot_download
@@ -65,6 +66,15 @@ def train_on_l4(hours: float = 20.0, ckpt_every: int = 250,
     os.makedirs(data_dir, exist_ok=True)
     os.makedirs(ckpt_dir, exist_ok=True)
 
+    # reuse the canonical tokenizer / RAG corpus already fetched with the code
+    for name in ("tokenizer.json", "meta.json", "rag_corpus.txt"):
+        src = os.path.join("/root/clanker/data", name)
+        if os.path.exists(src):
+            shutil.copy(src, os.path.join(data_dir, name))
+
+    os.makedirs(data_dir, exist_ok=True)
+    os.makedirs(ckpt_dir, exist_ok=True)
+
     # ---- ensure data (generate in-container; fast HF egress) ----
     if not os.path.exists(os.path.join(data_dir, "train.bin")):
         print(f"[modal] train.bin missing; generating data locally (scale={scale}) ...")

prep.py

@@ -41,7 +41,7 @@ TOK_BUDGET_RAG = 150_000_000        # scaled by --scale
 # --------------------------------------------------------------------------
 # 1) Tokenizer training
 # --------------------------------------------------------------------------
-def train_tokenizer():
+def train_tokenizer(out_dir=DATADIR):
     print("[prep] streaming FineWeb-edu to collect tokenizer training docs ...")
     ds = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT",
                       streaming=True, split="train")
@@ -53,7 +53,7 @@ def train_tokenizer():
     print(f"[prep] collected {len(texts)} docs for tokenizer")
     tok = YKTokenizer().train(
         iter(texts), vocab_size=32768,
-        save_path=os.path.join(DATADIR, "tokenizer.json"))
+        save_path=os.path.join(out_dir, "tokenizer.json"))
     print(f"[prep] tokenizer trained: vocab={tok.vocab_size}")
     return tok
 
@@ -248,18 +248,17 @@ def main():
         # try to reuse the canonical tokenizer from HF (keeps all runs compatible)
         try:
             print("[prep] no local tokenizer; downloading canonical one from HF ...")
-            import subprocess, shutil, tempfile
-            tmp = tempfile.mkdtemp()
-            subprocess.run(["hf", "download", "coderofpears/clankerDiffusion-base",
-                            "data/tokenizer.json", "-d", tmp], check=True)
-            for root, _, files in os.walk(tmp):
-                if "tokenizer.json" in files:
-                    shutil.copy(os.path.join(root, "tokenizer.json"), tok_path)
-                    break
+            from huggingface_hub import hf_hub_download
+            tok_path = hf_hub_download(
+                repo_id="coderofpears/clankerDiffusion-base",
+                filename="data/tokenizer.json",
+                repo_type="model",
+                local_dir=out_dir,
+                token=os.environ.get("HF_TOKEN"))
             tok = YKTokenizer.load(tok_path)
         except Exception as e:
             print(f"[prep] HF tokenizer download failed ({e}); training a new one.")
-            tok = train_tokenizer()
+            tok = train_tokenizer(out_dir)
 
     bin_path = os.path.join(out_dir, "train.bin")
     if os.path.exists(bin_path):

train.py

@@ -44,12 +44,24 @@ def build_cfg(args):
 
 
 def _push_hf(path, repo):
-    """Upload a single checkpoint file to HF via the `hf` CLI (background)."""
+    """Upload a single checkpoint file to HF (background thread)."""
     if not repo:
         return
     try:
-        subprocess.run(["hf", "upload", repo, path, "--repo-type", "model"],
-                       check=True, capture_output=True, timeout=600)
+        from huggingface_hub import HfApi
+        token = os.environ.get("HF_TOKEN")
+        if not token:
+            for p in (os.path.join(HERE, ".env"),
+                      os.path.join(os.path.dirname(HERE), ".env"),
+                      os.path.join(os.path.expanduser("~"), ".env")):
+                if os.path.exists(p):
+                    for line in open(p, encoding="utf-8"):
+                        if line.strip().startswith("HF_TOKEN"):
+                            token = line.split("=", 1)[1].strip().strip('"').strip("'")
+        api = HfApi(token=token)
+        api.upload_file(path_or_fileobj=path,
+                        path_in_repo=os.path.basename(path),
+                        repo_id=repo, repo_type="model")
         print(f"[hf] pushed {os.path.basename(path)} -> {repo}", flush=True)
     except Exception as e:
         print(f"[hf] push failed for {path}: {e}", flush=True)