Upload folder using huggingface_hub

modal_train.py

@@ -1,21 +1,21 @@
 """
 clankerDiffusion — Modal L4 training.
 
-Bakes the project code into the image, keeps data + checkpoints on a persistent
-Modal Volume (/vol), injects HF_TOKEN from .env, trains a big hybrid model on an
-L4 (24 GB), and pushes every checkpoint to the HF repo:
+Builds a CUDA image, pulls the latest code from the HF code repo at runtime,
+keeps data + checkpoints on a persistent Modal Volume (/vol), injects HF_TOKEN
+from .env, trains a big hybrid model on an L4 (24 GB), and pushes every
+checkpoint to the HF repo:
 
     coderofpears/clankerDiffusion-checkpoints
 
-Launch (after `modal token new` on this machine):
-    modal run modal_train.py::train_on_l4 --hours 20 --size large
+Data is generated in-container by prep.py (fast HF egress), scaled by --scale,
+so no large file transfer is needed.
 
-Data: stage train.bin + tokenizer.json onto the volume once, e.g.
-    modal volume put clanker-vol C:\Users\User\CalcGPU\clankerDiffusion\data /data
-(or it is downloaded from the HF data repo if present there).
+Launch (after `modal token new` on this machine):
+    modal run modal_train.py::train_on_l4 --hours 20 --size large --scale 3
 """
 import os
-from modal import App, Image, Volume, Secret, function
+from modal import App, Image, Volume, Secret
 
 HERE = os.path.dirname(os.path.abspath(__file__))
 CODE_REPO = "coderofpears/clankerDiffusion-base"
@@ -25,17 +25,10 @@ DOTENV = os.path.join(HERE, ".env")
 
 image = (
     Image.debian_slim()
-    .pip_install(
-        "torch==2.11.0", index_url="https://download.pytorch.org/whl/cu128",
-        "numpy", "tokenizers", "datasets", "safetensors",
-        "huggingface_hub", "accelerate", "hf-transfer",
-    )
+    .pip_install("torch==2.11.0", index_url="https://download.pytorch.org/whl/cu128")
+    .pip_install("numpy", "tokenizers", "datasets", "safetensors",
+                 "huggingface_hub", "hf-transfer")
     .env({"HF_HUB_ENABLE_HF_TRANSFER": "1", "TOKENIZERS_PARALLELISM": "false"})
-    .add_local_dir(
-        HERE, "/root/clanker",
-        ignore=["data", "checkpoints", ".venv", "__pycache__", "*.bin",
-                "*.pt", ".git", "*.log", "temp_up"],
-    )
 )
 
 app = App("clanker-diffusion", image=image)
@@ -54,30 +47,27 @@ SIZES = {
     timeout=60 * 60 * 25,
     volumes={"/vol": volume},
     secrets=[Secret.from_dotenv(DOTENV)],
-    _allow_background_volume_commits=True,
 )
 def train_on_l4(hours: float = 20.0, ckpt_every: int = 250,
-                size: str = "large", batch: int = None):
-    import os
+                size: str = "large", batch: int = None, scale: float = 3.0):
     import subprocess
     import sys
 
+    # ---- pull latest code from HF ----
+    os.makedirs("/root/clanker", exist_ok=True)
+    subprocess.run(["hf", "download", CODE_REPO, "/root/clanker",
+                    "--repo-type", "model"], check=True)
     os.chdir("/root/clanker")
+
     data_dir, ckpt_dir = "/vol/data", "/vol/checkpoints"
     os.makedirs(data_dir, exist_ok=True)
     os.makedirs(ckpt_dir, exist_ok=True)
 
-    # ---- ensure data ----
+    # ---- ensure data (generate in-container; fast HF egress) ----
     if not os.path.exists(os.path.join(data_dir, "train.bin")):
-        print("[modal] train.bin missing; trying HF download ...")
-        try:
-            subprocess.run(["hf", "download", DATA_REPO, "train.bin", "-d", data_dir],
-                           check=True)
-            subprocess.run(["hf", "download", DATA_REPO, "tokenizer.json", "-d", data_dir],
-                           check=True)
-        except Exception as e:
-            print(f"[modal] HF data download failed: {e}; bailing.")
-            raise
+        print(f"[modal] train.bin missing; generating data locally (scale={scale}) ...")
+        subprocess.run([sys.executable, "prep.py", "--out-dir", data_dir,
+                        "--scale", str(scale)], check=True)
 
     spec = dict(SIZES.get(size, SIZES["large"]))
     if batch:
@@ -102,5 +92,7 @@ def train_on_l4(hours: float = 20.0, ckpt_every: int = 250,
 
 
 @app.local_entrypoint()
-def main(hours: float = 20.0, ckpt_every: int = 250, size: str = "large", batch: int = None):
-    train_on_l4.remote(hours=hours, ckpt_every=ckpt_every, size=size, batch=batch)
+def main(hours: float = 20.0, ckpt_every: int = 250, size: str = "large",
+         batch: int = None, scale: float = 3.0):
+    train_on_l4.remote(hours=hours, ckpt_every=ckpt_every, size=size,
+                        batch=batch, scale=scale)

prep.py

@@ -1,24 +1,27 @@
 """
-Data preparation for yk_diffusion.
-
-One streaming pass over FineWeb-edu:
-  1. first N docs  -> train the byte-level BPE tokenizer (from scratch)
-  2. remainder     -> tokenize and pack into a flat uint16 .bin until token budget
-Then append:
-  - synthetic tool-use conversations (guarantees the agent's exact tag format)
-  - Glaive function-calling data (best-effort) if it loads
-
-Outputs (under ./data):
-  tokenizer.json + .meta.json
-  train.bin   (flat uint16 tokens)
-  meta.json   ({n_tokens, seq_len, vocab_size})
+Data preparation for clankerDiffusion.
+
+One streaming pass over FineWeb-edu (+ a Wikipedia slice for world knowledge,
++ synthetic tool-use and RAG/retrieval examples that teach the special tags):
+  1. first N docs -> train the byte-level BPE tokenizer (from scratch)
+  2. remainder    -> tokenize and pack into a flat uint16 .bin until token budget
+
+Outputs (under --out-dir, default ./data):
+  tokenizer.json  meta.json
+  train.bin       (flat uint16 tokens)
+  meta.json       ({n_tokens, seq_len, vocab_size})
+
+Designed so EVERY training environment (local / Modal / Kaggle / TPU) can
+regenerate the corpus itself with fast HF egress -- no 2 GB file transfer needed.
+Use --scale to grow the corpus (e.g. --scale 4 for a multi-billion-token run).
 """
-import os, sys, json, random, argparse
+import os, json, random, argparse
 import numpy as np
 from datasets import load_dataset
 
 import tokenizer as tokmod
 from tokenizer import YKTokenizer, SPECIAL
+from build_rag import FACTS, SYSTEM as RAG_SYSTEM
 
 random.seed(1234)
 np.random.seed(1234)
@@ -28,9 +31,11 @@ DATADIR = os.path.join(OUT, "data")
 os.makedirs(DATADIR, exist_ok=True)
 
 SEQ_LEN = 1024
-TOK_BUDGET = 900_000_000      # fineweb portion target
-TOK_BUDGET_TOOL = 200_000_000   # synthetic + glaive portion cap
-TOK_TRAIN_DOCS = 60_000          # docs used to train the tokenizer
+TOK_TRAIN_DOCS = 80_000
+TOK_BUDGET_FINEWEB = 900_000_000     # scaled by --scale
+TOK_BUDGET_WIKI = 300_000_000       # scaled by --scale
+TOK_BUDGET_TOOL = 250_000_000       # scaled by --scale
+TOK_BUDGET_RAG = 150_000_000        # scaled by --scale
 
 
 # --------------------------------------------------------------------------
@@ -58,7 +63,8 @@ def train_tokenizer():
 # --------------------------------------------------------------------------
 CALC_TEMPLATES = [
     "What is {a} {op} {b}?", "Compute {a} {op} {b} for me.",
-    "Calculate the result of {a} {op} {b}.", "If I have {a} and add/subtract/multiply/divide by {b} what do I get?",
+    "Calculate the result of {a} {op} {b}.",
+    "If I start at {a} and apply {op} {b}, what do I get?",
 ]
 OPS = {"+": "plus", "-": "minus", "*": "times", "/": "divided by"}
 PY_SNIPPETS = [
@@ -66,27 +72,29 @@ PY_SNIPPETS = [
     "import math\nprint(round(math.sqrt({n}), 4))",
     "print(sorted([{a}, {b}, {c}]))",
     "print({n} ** 2 + {n})",
+    "s='clanker'; print(s[::-1])",
 ]
 FILE_Q = [
     "Read the file {path} and tell me what is on the first line.",
-    "What is inside {path}?",
-    "List the files in {dir}.",
+    "What is inside {path}?", "List the files in {dir}.",
 ]
 SYSTEM = ("You are clanker, a helpful assistant that can THINK and USE TOOLS. "
           "When you need to compute or inspect something, wrap a tool call in "
           "<tool name=\"...\">arguments</tool>. Available tools: calc(expr), "
-          "python(code), read_file(path), list_dir(path). After a tool result "
-          "appears in <result>...</result>, continue and give the final answer. "
-          "You may use <think>...</think> to reason first.")
+          "python(code), read_file(path), list_dir(path), retrieve(query). After "
+          "a tool result appears in <result>...</result>, continue and give the "
+          "final answer. If <context>...</context> is provided, use it. You may "
+          "use <think>...</think> to reason first.")
 
 
-def gen_synthetic(n=25000):
+def gen_synthetic(n=60000):
     out = []
     for _ in range(n):
         kind = random.random()
-        if kind < 0.5:
+        if kind < 0.45:
             a = random.randint(2, 999); b = random.randint(2, 999)
-            op = random.choice(["+", "-", "*", "/"]); b = max(2, b if op != "/" else random.randint(2, 50))
+            op = random.choice(["+", "-", "*", "/"])
+            b = max(2, b if op != "/" else random.randint(2, 50))
             if op == "/":
                 a = a * b
             ans = eval(f"{a}{op}{b}")
@@ -101,7 +109,7 @@ def gen_synthetic(n=25000):
             tool = f'<tool name="python">{code}</tool>'
             try:
                 import io, contextlib
-                buf = io.StringIO(); 
+                buf = io.StringIO()
                 with contextlib.redirect_stdout(buf):
                     exec(code, {"__builtins__": __builtins__}, {})
                 result = buf.getvalue().strip()
@@ -128,7 +136,29 @@ def gen_synthetic(n=25000):
 
 
 # --------------------------------------------------------------------------
-# 3) Glaive function-calling (best-effort)
+# 3) RAG / retrieval examples (teach <tool name="retrieve"> and <context>)
+# --------------------------------------------------------------------------
+def gen_rag(n=40000):
+    out = []
+    for _ in range(n):
+        topic, doc, q, a = random.choice(FACTS)
+        mode = random.random()
+        if mode < 0.5:
+            conv = (f"<bos><system>{RAG_SYSTEM}</system><user>{q}</user>"
+                    f"<assistant><tool name=\"retrieve\">{q}</tool>"
+                    f"<result>{doc}</result>{a}</assistant><eos>")
+        elif mode < 0.85:
+            conv = (f"<bos><system>{RAG_SYSTEM}</system><user>{q}</user>"
+                    f"<assistant><context>{doc}</context>{a}</assistant><eos>")
+        else:
+            conv = (f"<bos><system>{RAG_SYSTEM}</system><user>{q}</user>"
+                    f"<assistant><think>{doc}</think>{a}</assistant><eos>")
+        out.append(conv)
+    return out
+
+
+# --------------------------------------------------------------------------
+# 4) Glaive function-calling (best-effort)
 # --------------------------------------------------------------------------
 def gen_glaive(max_examples=20000):
     out = []
@@ -152,7 +182,6 @@ def gen_glaive(max_examples=20000):
                 elif role in ("human", "user"):
                     parts.append(f"<user>{val}</user>")
                 elif role in ("gpt", "assistant", "function"):
-                    # function_call style -> our <tool> tag
                     val = val.replace("{\"name\":", "<tool name=\"").replace("\"function_call\"", "")
                     parts.append(f"<assistant>{val}</assistant>")
                 elif role == "tool":
@@ -166,12 +195,9 @@ def gen_glaive(max_examples=20000):
 
 
 # --------------------------------------------------------------------------
-# 4) Packing
+# 5) Packing
 # --------------------------------------------------------------------------
 def pack(tok, texts, bin_path, budget, seq_len):
-    """Append documents (packed into seq_len chunks) to bin_path until `budget`
-    tokens are written. Each call tracks its OWN counter so budgets are
-    independent across calls (fineweb vs synthetic vs glaive)."""
     n = 0
     buf = []
     with open(bin_path, "ab") as f:
@@ -198,52 +224,101 @@ def pack(tok, texts, bin_path, budget, seq_len):
 
 def main():
     ap = argparse.ArgumentParser()
-    ap.add_argument("--seed-tok-only", action="store_true")
+    ap.add_argument("--out-dir", default=DATADIR)
+    ap.add_argument("--scale", type=float, default=1.0,
+                    help="multiply token budgets (e.g. 4 -> ~4x more data)")
+    ap.add_argument("--no-wiki", action="store_true")
+    ap.add_argument("--no-rag", action="store_true")
+    ap.add_argument("--no-glaive", action="store_true")
     args = ap.parse_args()
 
-    tok_path = os.path.join(DATADIR, "tokenizer.json")
+    out_dir = args.out_dir
+    os.makedirs(out_dir, exist_ok=True)
+    scale = args.scale
+    bw_fw = int(TOK_BUDGET_FINEWEB * scale)
+    bw_wiki = int(TOK_BUDGET_WIKI * scale)
+    bw_tool = int(TOK_BUDGET_TOOL * scale)
+    bw_rag = int(TOK_BUDGET_RAG * scale)
+
+    tok_path = os.path.join(out_dir, "tokenizer.json")
     if os.path.exists(tok_path):
         print("[prep] loading existing tokenizer")
         tok = YKTokenizer.load(tok_path)
     else:
-        tok = train_tokenizer()
+        # 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
+            tok = YKTokenizer.load(tok_path)
+        except Exception as e:
+            print(f"[prep] HF tokenizer download failed ({e}); training a new one.")
+            tok = train_tokenizer()
 
-    bin_path = os.path.join(DATADIR, "train.bin")
+    bin_path = os.path.join(out_dir, "train.bin")
     if os.path.exists(bin_path):
         os.remove(bin_path)
 
-    # --- fineweb-edu remainder (same stream, continue) ---
-    print("[prep] tokenizing FineWeb-edu (this is the bandwidth-heavy step) ...")
+    # --- fineweb-edu ---
+    print(f"[prep] FineWeb-edu (budget {bw_fw:,}) ...")
     ds = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT",
                       streaming=True, split="train")
     gen = iter(ds)
-    for _ in range(TOK_TRAIN_DOCS):       # advance past tokenizer docs
+    for _ in range(TOK_TRAIN_DOCS):
         next(gen)
     def fineweb_iter():
         for ex in gen:
             yield ex["text"]
-    n = pack(tok, fineweb_iter(), bin_path, TOK_BUDGET, SEQ_LEN)
+    n = pack(tok, fineweb_iter(), bin_path, bw_fw, SEQ_LEN)
     print(f"[prep] fineweb packed: {n:,} tokens")
 
+    # --- wikipedia (world knowledge) ---
+    if not args.no_wiki:
+        print(f"[prep] Wikipedia (budget {bw_wiki:,}) ...")
+        try:
+            wds = load_dataset("wikipedia", "20220301.en",
+                               streaming=True, split="train")
+            def wiki_iter():
+                for ex in wds:
+                    yield ex["text"]
+            nw = pack(tok, wiki_iter(), bin_path, bw_wiki, SEQ_LEN)
+            print(f"[prep] wikipedia packed: {nw:,} tokens")
+            n += nw
+        except Exception as e:
+            print(f"[prep] wikipedia skipped: {e}")
+
     # --- synthetic tool data ---
-    print("[prep] generating synthetic tool-use conversations ...")
-    synth = gen_synthetic(60000)
-    n2 = pack(tok, synth, bin_path, TOK_BUDGET_TOOL, SEQ_LEN)
-    print(f"[prep] synthetic packed: {n2:,} tokens")
+    synth = gen_synthetic(int(60_000 * scale) + 60000)
+    n2 = pack(tok, synth, bin_path, bw_tool, SEQ_LEN)
+    print(f"[prep] synthetic tool packed: {n2:,} tokens")
+
+    # --- RAG / retrieval data ---
+    if not args.no_rag:
+        rag = gen_rag(int(40_000 * scale) + 40000)
+        nr = pack(tok, rag, bin_path, bw_rag, SEQ_LEN)
+        print(f"[prep] RAG packed: {nr:,} tokens")
+        n2 += nr
 
     # --- glaive ---
-    gl = gen_glaive(20000)
-    n3 = 0
-    if gl:
-        n3 = pack(tok, gl, bin_path, TOK_BUDGET_TOOL, SEQ_LEN)
+    if not args.no_glaive:
+        gl = gen_glaive(20000)
+        n3 = pack(tok, gl, bin_path, bw_tool, SEQ_LEN) if gl else 0
+    else:
+        n3 = 0
 
     total = n + n2 + n3
     meta = {"n_tokens": int(total), "seq_len": SEQ_LEN,
-             "vocab_size": tok.vocab_size, "path": "data/train.bin"}
-    with open(os.path.join(DATADIR, "meta.json"), "w") as f:
+            "vocab_size": tok.vocab_size, "path": "train.bin", "scale": scale}
+    with open(os.path.join(out_dir, "meta.json"), "w") as f:
         json.dump(meta, f)
     print(f"[prep] DONE. total tokens={total:,}  vocab={tok.vocab_size}")
-    print(f"[prep] files: {bin_path}")
+    print(f"[prep] files in {out_dir}")
 
 
 if __name__ == "__main__":

push_hf.py

@@ -1,6 +1,5 @@
 """
-clankerDiffusion — push artifacts to HuggingFace (via the `hf` CLI, which
-correctly creates the first commit on a brand-new repo).
+clankerDiffusion — push artifacts to HuggingFace (via huggingface_hub).
 
   python push_hf.py --what all        # code + tokenizer + rag corpus + train.bin
   python push_hf.py --what code       # .py sources + tokenizer + rag corpus
@@ -8,13 +7,15 @@ correctly creates the first commit on a brand-new repo).
   python push_hf.py --what ckpt       # latest checkpoints
 
 Repos (under your HF user):
-  clankerDiffusion-base        code + tokenizer + rag corpus
-  clankerDiffusion-data        train.bin  (packed training corpus)
-  clankerDiffusion-checkpoints model checkpoints (.pt)
+  coderofpears/clankerDiffusion-base        code + tokenizer + rag corpus
+  coderofpears/clankerDiffusion-data        train.bin  (packed training corpus)
+  coderofpears/clankerDiffusion-checkpoints model checkpoints (.pt)
 """
 import os
-import subprocess
 import argparse
+import tempfile
+import shutil
+import fnmatch
 
 HERE = os.path.dirname(os.path.abspath(__file__))
 DATADIR = os.path.join(HERE, "data")
@@ -29,14 +30,28 @@ PY_SOURCES = ["model.py", "tokenizer.py", "prep.py", "train.py", "infer.py",
               "colab_train.py", "modal_train.py", "push_hf.py", "upload_artifacts.py"]
 
 
-def _run(cmd):
-    print("+", " ".join(cmd), flush=True)
-    subprocess.run(cmd, check=True)
+def _token():
+    # prefer explicit env, else read from .env
+    tok = os.environ.get("HF_TOKEN")
+    if tok:
+        return tok
+    for p in (os.path.join(HERE, ".env"), os.path.join(os.path.expanduser("~"), ".env")):
+        if os.path.exists(p):
+            for line in open(p, encoding="utf-8"):
+                line = line.strip()
+                if line.startswith("HF_TOKEN"):
+                    return line.split("=", 1)[1].strip().strip('"').strip("'")
+    raise SystemExit("HF_TOKEN not found in env or .env")
 
 
 def _upload(repo, local, patterns):
-    # upload only the matching files via a temp staging dir
-    import tempfile, shutil, fnmatch
+    from huggingface_hub import HfApi
+    api = HfApi(token=_token())
+    try:
+        api.create_repo(repo_id=repo, repo_type="model", exist_ok=True)
+    except Exception as e:
+        print(f"[hf] create_repo note: {e}")
+
     stage = tempfile.mkdtemp()
     try:
         for root, _, files in os.walk(local):
@@ -49,7 +64,7 @@ def _upload(repo, local, patterns):
         if not os.listdir(stage):
             print(f"[hf] nothing matched for {repo}")
             return
-        _run(["hf", "upload", repo, stage, "--repo-type", "model"])
+        api.upload_folder(folder_path=stage, repo_id=repo, repo_type="model")
         print(f"[hf] -> {repo}")
     finally:
         shutil.rmtree(stage, ignore_errors=True)