Upload folder using huggingface_hub

agent.py

@@ -0,0 +1,99 @@
+"""
+clankerDiffusion — tool-using agent (ReAct loop) with mid-response RAG.
+
+The model emits <tool name="...">args</tool> blocks. We parse + execute
+them locally (tools.py) and feed the <result> back, looping until the
+model stops calling tools. Each turn alternates AR <-> diffusion, so the
+agent is literally "sometimes normal, sometimes diffusion".
+
+Knowledge injection happens in TWO ways, both mid-response:
+  * model-driven : the model calls <tool name="retrieve">q</tool>; we run the
+    retriever and feed <result> back, then keep generating.
+  * controller-driven : the RAG controller watches the partial response, builds
+    a query from the user question + what was just said, and if it finds
+    relevant passages it injects <context>...</context> straight into the
+    stream and the model continues -- no need for the model to ask.
+"""
+import argparse
+import infer
+from tools import execute_tool, parse_tool_calls
+from rag import default_kb
+
+
+SYSTEM = ("You are clanker, a helpful assistant that can THINK, USE TOOLS, and "
+          "READ CONTEXT. When you need to compute, inspect files, or look up "
+          "knowledge, emit a tool call like "
+          "<tool name=\"calc\">17 * 23</tool>, "
+          "<tool name=\"python\">print(2**10)</tool>, "
+          "<tool name=\"read_file\">notes.txt</tool>, "
+          "<tool name=\"list_dir\">.</tool>, or "
+          "<tool name=\"retrieve\">the capital of France</tool>. "
+          "Available tools: calc(expr), python(code), read_file(path), "
+          "list_dir(path), retrieve(query). If <context>...</context> is "
+          "provided, use it to answer. After a <result>...</result> appears, "
+          "continue and give the final answer. You may reason in <think>...</think>.")
+
+
+def run_agent(prompt, max_turns=8, mode_cycle=("diff", "ar"),
+              max_new=200, gen_len=160, steps=24,
+              rag=True, rag_k=3, rag_budget=2):
+    model, tok = infer._MODEL, infer._TOK
+    kb = default_kb()
+    injected = set()
+
+    ctx = [tok.bos_id] + tok.encode(
+        f"<system>{SYSTEM}</system><user>{prompt}</user><assistant>")
+
+    print(f"\n=== USER ===\n{prompt}\n")
+    for turn in range(max_turns):
+        mode = mode_cycle[turn % len(mode_cycle)]
+        prev_len = len(ctx)
+        print(f"--- turn {turn} [{mode}] ---", flush=True)
+        if mode == "ar":
+            ids = infer.generate_ar(model, tok, ctx, max_new=max_new)
+        else:
+            ids = infer.generate_diff(model, tok, ctx, gen_len=gen_len, steps=steps)
+        new_ids = ids[prev_len:]              # ONLY the newly generated tokens
+        ctx = ids                             # full updated context (no dup)
+        text_new = tok.decode(new_ids)
+        print(text_new, flush=True)
+
+        # 1) model-driven tool calls (incl. retrieve) -> mid-response injection
+        calls = parse_tool_calls(text_new)
+        if calls:
+            for name, arg in calls:
+                print(f"  [tool] {name}({arg[:120]})", flush=True)
+                res = execute_tool(name, arg)
+                print(f"  [result] {res[:300]}", flush=True)
+                ctx = ctx + tok.encode(f"<result>{res}</result>")
+            continue                          # keep going within the same answer
+
+        # 2) controller-driven RAG: inject context mid-response, then continue
+        if rag and len(injected) < rag_budget * rag_k:
+            q = prompt + " " + text_new[-300:]
+            fresh = [p for p in kb.retrieve(q, k=rag_k) if p not in injected]
+            if fresh:
+                injected.update(fresh)
+                block = "<context>\n" + "\n---\n".join(fresh) + "\n</context>"
+                ctx = ctx + tok.encode(block)
+                print(f"  [rag] injected {len(fresh)} passage(s) mid-response",
+                      flush=True)
+                continue
+
+        break                                 # nothing to do -> answer is final
+
+    final = tok.decode(ctx)
+    i = final.rfind("<assistant>")
+    print("\n=== clanker (final) ===\n" + (final[i + 10:] if i >= 0 else final))
+    return final
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("prompt", nargs="?", default="What is 17 * 23, and what files are in the current folder?")
+    ap.add_argument("--ckpt", default=None)
+    ap.add_argument("--turns", type=int, default=8)
+    ap.add_argument("--no-rag", action="store_true")
+    a = ap.parse_args()
+    infer.init(ckpt_path=a.ckpt)
+    run_agent(a.prompt, max_turns=a.turns, rag=not a.no_rag)

build_rag.py

@@ -0,0 +1,114 @@
+"""
+clankerDiffusion — build the RAG knowledge base + RAG training data.
+
+Run AFTER the tokenizer exists (data/tokenizer.json):
+    python build_rag.py
+
+  * writes a small default knowledge base into data/kb/ and saves its index
+  * writes data/rag_corpus.txt : many formatted examples teaching the model
+      - to call <tool name="retrieve">q</tool> and use the <result>, and
+      - to read <context>...</context> injected mid-response and answer from it.
+    This corpus is consumed by rag_finetune.py (continuation training).
+"""
+import os
+import json
+import random
+
+from rag import KnowledgeBase, DEFAULT_INDEX, KB_DIR
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+DATADIR = os.path.join(HERE, "data")
+CORPUS = os.path.join(DATADIR, "rag_corpus.txt")
+
+# A small factual knowledge base. (In practice, point build_rag.py at your own
+# docs via kb.ingest_path(...); this seed makes `retrieve` work out of the box.)
+FACTS = [
+    ("geography", "France is a country in Western Europe. Its capital and largest city is Paris. The currency is the euro.",
+     "What is the capital of France?", "The capital of France is Paris."),
+    ("geography", "Germany's capital is Berlin. The official language is German and the currency is the euro.",
+     "What is the capital of Germany?", "The capital of Germany is Berlin."),
+    ("geography", "Japan is an island nation in East Asia. Its capital is Tokyo and its highest mountain is Mount Fuji.",
+     "What is the capital of Japan?", "The capital of Japan is Tokyo."),
+    ("science", "Water is a chemical compound with the formula H2O. It boils at 100 degrees Celsius at standard pressure and freezes at 0 degrees Celsius.",
+     "At what temperature does water boil?", "Water boils at 100 degrees Celsius at standard pressure."),
+    ("science", "The speed of light in a vacuum is approximately 299,792 kilometers per second.",
+     "What is the speed of light?", "The speed of light in a vacuum is about 299,792 km/s."),
+    ("science", "The chemical symbol for gold is Au, for silver is Ag, and for iron is Fe.",
+     "What is the chemical symbol for gold?", "The chemical symbol for gold is Au."),
+    ("history", "World War II lasted from 1939 to 1945. It involved most of the world's nations and ended with the defeat of the Axis powers.",
+     "When did World War II end?", "World War II ended in 1945."),
+    ("history", "The Declaration of Independence of the United States was adopted on July 4, 1776.",
+     "When was the US Declaration of Independence adopted?", "It was adopted on July 4, 1776."),
+    ("tech", "Python is a high-level, interpreted programming language. It uses indentation to define blocks and is widely used for data science and AI.",
+     "What kind of language is Python?", "Python is a high-level, interpreted programming language."),
+    ("tech", "The Transformer is a neural network architecture introduced in 2017 that relies on self-attention instead of recurrence.",
+     "What is a Transformer in machine learning?", "A Transformer is a neural architecture from 2017 based on self-attention."),
+    ("math", "The value of pi is approximately 3.14159. It is the ratio of a circle's circumference to its diameter.",
+     "What is the value of pi?", "Pi is approximately 3.14159."),
+    ("math", "Euler's identity is e^(i*pi) + 1 = 0, linking the numbers e, i, pi, 1, and 0.",
+     "What is Euler's identity?", "Euler's identity is e^(i*pi) + 1 = 0."),
+    ("space", "The Sun is the star at the center of the Solar System. Earth orbits it at an average distance of about 149.6 million kilometers.",
+     "What is the Sun?", "The Sun is the star at the center of the Solar System."),
+    ("space", "Mars is the fourth planet from the Sun and is often called the Red Planet because of its iron-oxide surface.",
+     "Why is Mars called the Red Planet?", "Mars is called the Red Planet due to its iron-oxide (rusty) surface."),
+    ("biology", "DNA stands for deoxyribonucleic acid. It carries the genetic instructions used in the growth and functioning of all known living organisms.",
+     "What does DNA stand for?", "DNA stands for deoxyribonucleic acid."),
+    ("biology", "Photosynthesis is the process by which plants convert light energy, water, and carbon dioxide into glucose and oxygen.",
+     "What is photosynthesis?", "Photosynthesis is how plants turn light, water, and CO2 into glucose and oxygen."),
+    ("economics", "Inflation is the rate at which the general level of prices for goods and services rises, eroding purchasing power.",
+     "What is inflation?", "Inflation is the rise in the general price level, reducing purchasing power."),
+    ("economics", "Gross Domestic Product (GDP) is the total monetary value of all finished goods and services produced within a country in a period.",
+     "What is GDP?", "GDP is the total value of finished goods and services produced in a country."),
+    ("language", "The word 'clanker' in this project is the name of the assistant. It is a from-scratch hybrid diffusion language model.",
+     "What is clanker?", "clanker is the name of this assistant, a from-scratch hybrid diffusion language model."),
+    ("project", "clankerDiffusion alternates between AR mode (normal next-token) and DIFF mode (masked diffusion) using a learned mode embedding.",
+     "What are the two modes of clankerDiffusion?", "AR mode (autoregressive) and DIFF mode (masked diffusion)."),
+]
+
+SYSTEM = ("You are clanker, a helpful assistant that can use tools and read context. "
+          "When you need knowledge, call <tool name=\"retrieve\">query</tool>. "
+          "If <context>...</context> is provided, answer using it.")
+
+
+def build_kb():
+    os.makedirs(KB_DIR, exist_ok=True)
+    # a readable KB file so list_dir/retrieve demo works
+    kb_text = "\n\n".join(f"[{topic}]\n{doc}" for topic, doc, _, _ in FACTS)
+    with open(os.path.join(KB_DIR, "knowledge.txt"), "w", encoding="utf-8") as f:
+        f.write(kb_text)
+    kb = KnowledgeBase()
+    kb.ingest_path(KB_DIR)
+    kb.save(DEFAULT_INDEX)
+    return kb
+
+
+def build_corpus(n_each=400, seed=0):
+    random.seed(seed)
+    lines = []
+    facts = FACTS
+    for _ in range(n_each):
+        topic, doc, q, a = random.choice(facts)
+        # format A: model must RETRIEVE then answer
+        lines.append(
+            f"<system>{SYSTEM}</system><user>{q}</user>"
+            f"<assistant><tool name=\"retrieve\">{q}</tool>"
+            f"<result>{doc}</result>{a}</assistant>")
+        # format B: context already injected mid-response, model answers from it
+        lines.append(
+            f"<system>{SYSTEM}</system><user>{q}</user>"
+            f"<assistant><context>{doc}</context>{a}</assistant>")
+        # format C: a short chain-of-thought variant
+        lines.append(
+            f"<system>{SYSTEM}</system><user>{q}</user>"
+            f"<assistant><think>The relevant knowledge is: {doc}</think>"
+            f"{a}</assistant>")
+    random.shuffle(lines)
+    with open(CORPUS, "w", encoding="utf-8") as f:
+        for ln in lines:
+            f.write(ln + "\n")
+    print(f"[rag] wrote {len(lines)} training examples -> {CORPUS}")
+
+
+if __name__ == "__main__":
+    build_kb()
+    build_corpus()

colab_train.py

@@ -0,0 +1,215 @@
+"""
+============================================================================
+ clankerDiffusion — Colab A100-80GB training script
+============================================================================
+
+HOW TO RUN IN COLAB (A100 80GB):
+
+  [Cell 1 — setup, run once]
+      !pip install -q torch==2.11.0+cu124 -f https://download.pytorch.org/whl/cu124
+      !pip install -q transformers tokenizers datasets accelerate safetensors huggingface_hub numpy
+      import os
+      os.environ["HF_TOKEN"]   = "hf_xxx"          # your token (also set in Secrets)
+      os.environ["CODE_REPO"]  = "clankerDiffusion/base"        # from upload_artifacts.py
+      os.environ["CKPT_REPO"] = "clankerDiffusion/checkpoints"
+
+  [Cell 2 — launch in BACKGROUND, then disconnect safely]
+      !nohup python colab_train.py > colab_train.log 2>&1 &
+      # check later with:  !tail -n 30 colab_train.log
+      # it checkpoints + uploads to HF every 250 steps until the runtime dies
+
+The script trains the SAME from-scratch hybrid model (AR + masked
+diffusion) on FineWeb-edu, resuming if a checkpoint exists, and pushes
+a checkpoint to HuggingFace Hub every 250 steps in a background thread.
+============================================================================
+"""
+import os, sys, json, time, threading, argparse
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import snapshot_download, HfApi
+
+CODE_REPO = os.environ.get("CODE_REPO", "clankerDiffusion/base")
+CKPT_REPO = os.environ.get("CKPT_REPO", "clankerDiffusion/checkpoints")
+HF_TOKEN = os.environ.get("HF_TOKEN")
+api = HfApi(token=HF_TOKEN)
+
+# ---- pull our model code + tokenizer from HF -------------------------------
+print(f"[colab] downloading code from {CODE_REPO} ...")
+local = snapshot_download(CODE_REPO, repo_type="model")
+sys.path.insert(0, local)
+from model import YKDiff
+from tokenizer import YKTokenizer
+
+# ---- big architecture for A100 80GB -------------------------------------
+CFG = dict(
+    d_model=2048, n_layers=24, n_heads=16, d_ff=5504,
+    max_len=2048, vocab_size=32768,
+)
+
+tok = YKTokenizer.load(os.path.join(local, "tokenizer.json"))
+CFG["vocab_size"] = tok.vocab_size
+V = CFG["vocab_size"]
+mask_id, pad_id = tok.mask_id, tok.pad_id
+print(f"[colab] vocab={V}")
+
+# ---- streaming fineweb into a rolling token buffer -------------------------
+from datasets import load_dataset
+ds = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT",
+                  streaming=True, split="train")
+BUF_CAP = 60_000_000
+buf = []
+_buf_lock = threading.Lock()
+
+
+def _refill():
+    for ex in ds:
+        ids = tok.encode(ex["text"])
+        with _buf_lock:
+            buf.extend(ids)
+            if len(buf) > BUF_CAP:
+                del buf[: len(buf) - BUF_CAP]
+
+
+threading.Thread(target=_refill, daemon=True).start()
+
+
+def sample_batch(batch, seq_len):
+    with _buf_lock:
+        if len(buf) < seq_len + 1:
+            return None
+        N = len(buf)
+        starts = np.random.randint(0, N - seq_len, size=batch)
+        return torch.tensor(
+            [buf[s:s + seq_len] for s in starts], dtype=torch.long)
+
+
+# ---- model -----------------------------------------------------------------
+model = YKDiff(CFG).cuda()
+n_params = sum(p.numel() for p in model.parameters())
+print(f"[colab] params = {n_params/1e9:.2f}B")
+optim = torch.optim.AdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.95),
+                          weight_decay=0.1)
+
+# ---- resume ----------------------------------------------------------------
+CKPT_LOCAL = "/content/clanker_ckpts"
+os.makedirs(CKPT_LOCAL, exist_ok=True)
+step0 = 0
+existing = sorted(f for f in os.listdir(CKPT_LOCAL) if f.endswith(".pt"))
+if existing:
+    sd = torch.load(os.path.join(CKPT_LOCAL, existing[-1]), map_location="cuda")
+    model.load_state_dict(sd["model"]); optim.load_state_dict(sd["optim"])
+    step0 = sd["step"]
+    print(f"[colab] resumed step={step0}")
+else:
+    # try pulling latest from HF
+    try:
+        api.create_repo(CKPT_REPO, repo_type="model", exist_ok=True)
+    except Exception:
+        pass
+
+
+def _upload(path):
+    def _u():
+        try:
+            api.upload_file(repo_id=CKPT_REPO,
+                           path_in_repo=os.path.basename(path),
+                           path_or_fileobj=path, repo_type="model")
+            print(f"[colab] uploaded {os.path.basename(path)} -> {CKPT_REPO}",
+                  flush=True)
+        except Exception as e:
+            print(f"[colab] upload failed: {e}", flush=True)
+    threading.Thread(target=_u, daemon=True).start()
+
+
+# ---- training loop (hybrid) ----------------------------------------------
+@torch.no_grad()
+def _cosine_lr(step, warmup, total, base, minlr):
+    if step < warmup:
+        return base * step / warmup
+    p = (step - warmup) / max(total - warmup, 1)
+    return minlr + 0.5 * (base - minlr) * (1 + np.cos(np.pi * min(p, 1.0)))
+
+
+BATCH, SEQ, GRAD_ACCUM = 16, 2048, 4
+WARMUP, TOTAL_STEPS = 500, 200_000
+BASE_LR, MIN_LR = 1e-4, 1e-5
+CKPT_EVERY = 250
+amp = torch.cuda.amp.autocast(dtype=torch.bfloat16)
+
+step = step0
+model.train()
+t0 = time.time()
+print("[colab] training started.", flush=True)
+
+while True:                                   # run as long as possible
+    try:
+        optim.zero_grad(set_to_none=True)
+        for micro in range(GRAD_ACCUM):
+            idx = None
+            while idx is None:
+                idx = sample_batch(BATCH, SEQ)
+                time.sleep(0.02)
+            idx = idx.cuda()
+            mode_ar = (torch.rand(1).item() < 0.5)
+            with amp:
+                if mode_ar:
+                    m = torch.zeros(BATCH, dtype=torch.long, device="cuda")
+                    logits = model(idx, m, t=None)
+                    loss = F.cross_entropy(
+                        logits[:, :-1].reshape(-1, V),
+                        idx[:, 1:].reshape(-1), ignore_index=pad_id)
+                else:
+                    m = torch.ones(BATCH, dtype=torch.long, device="cuda")
+                    r = torch.rand(BATCH, device="cuda")
+                    is_mask = torch.rand(BATCH, SEQ, device="cuda") < r[:, None]
+                    not_pad = idx != pad_id
+                    masked = idx.clone(); masked[is_mask] = mask_id
+                    logits = model(masked, m, t=r)
+                    ce = F.cross_entropy(logits.reshape(-1, V),
+                                        idx.reshape(-1), reduction="none",
+                                        ignore_index=-100)
+                    ce = ce * is_mask.reshape(-1) * not_pad.reshape(-1)
+                    denom = (is_mask & not_pad).reshape(-1).sum().clamp(min=1)
+                    loss = ce.sum() / denom
+            (loss / GRAD_ACCUM).backward()
+
+        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        lr = _cosine_lr(step, WARMUP, TOTAL_STEPS, BASE_LR, MIN_LR)
+        for g in optim.param_groups:
+            g["lr"] = lr
+        optim.step()
+        step += 1
+
+        if step % 25 == 0:
+            print(f"[colab] step {step} loss~{loss.item():.3f} "
+                  f"lr={lr:.2e} t={(time.time()-t0)/60:.1f}m", flush=True)
+
+        if step % CKPT_EVERY == 0:
+            # full local ckpt (for resume)
+            full = os.path.join(CKPT_LOCAL, f"clanker_{step:07d}.pt")
+            torch.save({"model": model.state_dict(),
+                        "optim": optim.state_dict(),
+                        "step": step, "cfg": CFG, "vocab": V}, full)
+            # light bf16 model-only for HF upload
+            lite = os.path.join(CKPT_LOCAL, f"clanker_{step:07d}_lite.pt")
+            torch.save({"model": {k: v.to(torch.bfloat16)
+                                   for k, v in model.state_dict().items()},
+                        "cfg": CFG, "vocab": V, "step": step}, lite)
+            print(f"[colab] checkpoint {step}", flush=True)
+            _upload(lite)
+            # keep only last 2 local full ckpts to save disk
+            for old in sorted(f for f in os.listdir(CKPT_LOCAL)
+                           if f.endswith(".pt") and "lite" not in f)[:-2]:
+                os.remove(os.path.join(CKPT_LOCAL, old))
+
+    except torch.cuda.OutOfMemoryError:
+        print("[colab] OOM — skipping step", flush=True)
+        optim.zero_grad(set_to_none=True)
+        torch.cuda.empty_cache()
+    except Exception as e:
+        print(f"[colab] step error (continuing): {e}", flush=True)
+        torch.cuda.empty_cache()
+
+print("[colab] loop ended.")

infer.py

@@ -0,0 +1,135 @@
+"""
+clankerDiffusion — inference.
+
+Three generation modes, all from the SAME weights:
+  ar      normal autoregressive next-token decoding (causal)
+  diff    masked discrete diffusion: start the answer fully masked,
+          iteratively unmask the most-confident tokens (LLaDA-style remasking)
+  hybrid  a few AR steps to "think", then diffusion for the answer
+  random  flips a coin each call -> "sometimes normal, sometimes diffusion"
+"""
+import os, json, argparse, random
+import torch
+import torch.nn.functional as F
+from model import YKDiff
+from tokenizer import YKTokenizer
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+DATADIR = os.path.join(HERE, "data")
+CKPTDIR = os.path.join(HERE, "checkpoints")
+
+
+def load(tok_path=None, ckpt_path=None):
+    tok_path = tok_path or os.path.join(DATADIR, "tokenizer.json")
+    tok = YKTokenizer.load(tok_path)
+    if ckpt_path is None:
+        ck = sorted([f for f in os.listdir(CKPTDIR) if f.endswith(".pt")])
+        if not ck:
+            raise SystemExit("no checkpoint found in ./checkpoints")
+        ckpt_path = os.path.join(CKPTDIR, ck[-1])
+    sd = torch.load(ckpt_path, map_location="cuda")
+    cfg = sd["cfg"]
+    model = YKDiff(cfg).cuda().eval()
+    model.load_state_dict(sd["model"])
+    print(f"[infer] loaded {ckpt_path}  params={sum(p.numel() for p in model.parameters())/1e6:.1f}M")
+    return model, tok
+
+
+@torch.no_grad()
+def _topp(logits, temp, top_p):
+    logits = logits / max(temp, 1e-6)
+    if top_p >= 1.0:
+        return torch.multinomial(F.softmax(logits, -1), 1).item()
+    s, order = torch.sort(logits, descending=True)
+    p = F.softmax(s, -1)
+    c = torch.cumsum(p, -1)
+    keep = order[c <= top_p]
+    if len(keep) == 0:
+        keep = order[:1]
+    return keep[torch.multinomial(F.softmax(logits[keep], -1), 1).item()].item()
+
+
+@torch.no_grad()
+def generate_ar(model, tok, prompt_ids, max_new=256, temp=0.9, top_p=0.95):
+    ids = list(prompt_ids)
+    max_len = model.max_len
+    for _ in range(max_new):
+        ctx = torch.tensor([ids[-max_len:]], device="cuda")
+        logits = model(ctx, torch.zeros(1, dtype=torch.long, device="cuda"))[:, -1]
+        nxt = _topp(logits[0], temp, top_p)
+        if nxt == tok.eos_id:
+            break
+        ids.append(nxt)
+    return ids
+
+
+@torch.no_grad()
+def generate_diff(model, tok, prompt_ids, gen_len=128, steps=24, temp=1.0, sample=True):
+    L = len(prompt_ids) + gen_len
+    seq = list(prompt_ids) + [tok.mask_id] * gen_len
+    p0 = len(prompt_ids)
+    for step in range(steps):
+        x = torch.tensor([seq], device="cuda")
+        t = torch.full((1,), (steps - step - 1) / steps, device="cuda")
+        logits = model(x, torch.ones(1, dtype=torch.long, device="cuda"), t=t)
+        gl = logits[0, p0:]                         # generated-region logits
+        probs = F.softmax(gl / max(temp, 1e-6), -1)
+        if sample:
+            preds = torch.multinomial(probs, 1).squeeze(-1).tolist()
+        else:
+            preds = probs.argmax(-1).tolist()
+        conf = probs.max(-1).values            # [gen_len]
+        n_mask = int(round((steps - step - 1) / steps * gen_len))
+        order = conf.argsort().tolist()        # ascending confidence
+        mask_set = set(order[:n_mask])
+        for j in range(gen_len):
+            seq[p0 + j] = tok.mask_id if j in mask_set else preds[j]
+    return seq
+
+
+def build_prompt(tok, system, user):
+    return [tok.bos_id] + tok.encode(f"<system>{system}</system><user>{user}</user><assistant>")
+
+
+def generate(prompt_ids, mode="random", **kw):
+    model, tok = _MODEL, _TOK
+    if mode == "random":
+        mode = "ar" if random.random() < 0.5 else "diff"
+    if mode == "ar":
+        ids = generate_ar(model, tok, prompt_ids, **kw)
+    elif mode == "diff":
+        ids = generate_diff(model, tok, prompt_ids, **kw)
+    elif mode == "hybrid":
+        # think a little in AR, then diffuse the answer
+        think = generate_ar(model, tok, prompt_ids, max_new=64, **kw)
+        # generate_diff returns the FULL sequence (think + answer)
+        ids = generate_diff(model, tok, think, **kw)
+    else:
+        raise ValueError(mode)
+    return tok.decode(ids)
+
+
+# module-level cache so agent.py can call generate() directly
+_MODEL, _TOK = None, None
+
+
+def init(tok_path=None, ckpt_path=None):
+    global _MODEL, _TOK
+    _MODEL, _TOK = load(tok_path, ckpt_path)
+    return _MODEL, _TOK
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("prompt", nargs="?", default="What is 17 * 23?")
+    ap.add_argument("--mode", default="random")
+    ap.add_argument("--max-new", type=int, default=200)
+    ap.add_argument("--gen-len", type=int, default=160)
+    ap.add_argument("--steps", type=int, default=24)
+    ap.add_argument("--ckpt", default=None)
+    a = ap.parse_args()
+    m, t = init(ckpt_path=a.ckpt)
+    ids = build_prompt(t, "You are clanker, a helpful assistant.", a.prompt)
+    out = generate(ids, mode=a.mode, max_new=a.max_new,
+                   gen_len=a.gen_len, steps=a.steps)
+    print("\n=== clanker ===\n" + out)

model.py

@@ -0,0 +1,181 @@
+"""
+yk_diffusion: a from-scratch hybrid language model.
+
+A single Transformer runs in two modes, selected by a learned mode embedding:
+  - AR mode      (mode=0): causal attention -> standard next-token autoregressive LM ("normal")
+  - DIFF mode    (mode=1): bidirectional attention -> masked discrete diffusion denoising
+
+A time embedding conditions the diffusion mask ratio (MDLM / LLaDA-style absorbing-state training).
+The same weights serve both behaviours; a <MODE> signal picks which one at inference time.
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+
+# ----------------------------------------------------------------------------
+# Building blocks
+# ----------------------------------------------------------------------------
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d))
+
+    def forward(self, x):
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return x * self.weight
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, head_dim, max_len=8192, base=10000.0):
+        super().__init__()
+        inv = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+        self.register_buffer("inv_freq", inv)
+
+    def forward(self, seq_len, device):
+        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)            # [L, head_dim/2]
+        emb = torch.cat([freqs, freqs], dim=-1)          # [L, head_dim]
+        return emb.cos(), emb.sin()
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rope(x, cos, sin):
+    # x: [B, h, L, hd];  cos/sin: [1, 1, L, hd]
+    return x * cos + rotate_half(x) * sin
+
+
+class Attention(nn.Module):
+    def __init__(self, d, n_heads):
+        super().__init__()
+        assert d % n_heads == 0
+        self.d = d
+        self.n_heads = n_heads
+        self.hd = d // n_heads
+        self.scale = self.hd ** -0.5
+        self.qkv = nn.Linear(d, 3 * d, bias=False)
+        self.proj = nn.Linear(d, d, bias=False)
+
+    def forward(self, x, cos, sin, attn_mask=None):
+        B, L, _ = x.shape
+        qkv = self.qkv(x).reshape(B, L, 3, self.n_heads, self.hd)
+        qkv = qkv.permute(2, 0, 3, 1, 4)             # [3, B, h, L, hd]
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+        scores = (q @ k.transpose(-2, -1)) * self.scale
+        if attn_mask is not None:
+            scores = scores + attn_mask
+        attn = scores.softmax(dim=-1)
+        out = (attn @ v).transpose(1, 2).reshape(B, L, self.d)
+        return self.proj(out)
+
+
+class MLP(nn.Module):
+    def __init__(self, d, d_ff):
+        super().__init__()
+        self.w1 = nn.Linear(d, d_ff, bias=False)
+        self.w3 = nn.Linear(d, d_ff, bias=False)
+        self.w2 = nn.Linear(d_ff, d, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class Block(nn.Module):
+    def __init__(self, d, n_heads, d_ff):
+        super().__init__()
+        self.ln1 = RMSNorm(d)
+        self.attn = Attention(d, n_heads)
+        self.ln2 = RMSNorm(d)
+        self.mlp = MLP(d, d_ff)
+
+    def forward(self, x, cos, sin, attn_mask):
+        x = x + self.attn(self.ln1(x), cos, sin, attn_mask)
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+# ----------------------------------------------------------------------------
+# The model
+# ----------------------------------------------------------------------------
+
+class YKDiff(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        d = cfg["d_model"]
+        self.vocab = cfg["vocab_size"]
+        self.max_len = cfg["max_len"]
+
+        self.tok_emb = nn.Embedding(self.vocab, d)
+        self.mode_emb = nn.Embedding(2, d)          # 0 = AR, 1 = DIFF
+        self.time_emb = nn.Linear(1, d, bias=False)  # diffusion mask ratio conditioning
+        self.rope = RotaryEmbedding(d // cfg["n_heads"], max_len=self.max_len)
+
+        self.blocks = nn.ModuleList([
+            Block(d, cfg["n_heads"], cfg["d_ff"]) for _ in range(cfg["n_layers"])
+        ])
+        self.norm = RMSNorm(d)
+        self.lm_head = nn.Linear(d, self.vocab, bias=False)
+        with torch.no_grad():
+            self.lm_head.weight.copy_(self.tok_emb.weight)   # weight tying
+
+        self._causal = None
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def _causal_mask(self, L, device):
+        if self._causal is None or self._causal.shape[-1] < L:
+            m = torch.full((L, L), float("-inf"), device=device)
+            m = torch.triu(m, diagonal=1)
+            self._causal = m
+        return self._causal[:L, :L]
+
+    def forward(self, idx, mode, t=None, attn_mask=None):
+        """
+        idx:        [B, L] long   token ids
+        mode:       [B] long      (0 AR, 1 DIFF)
+        t:          [B] float|None  diffusion mask ratio (None -> 0)
+        attn_mask:  [L, L]|None    explicit mask; if None, AR uses causal, DIFF uses none
+        """
+        B, L = idx.shape
+        x = self.tok_emb(idx)
+        x = x + self.mode_emb(mode).unsqueeze(1)
+        if t is None:
+            t = torch.zeros(B, device=idx.device)
+        x = x + self.time_emb(t.unsqueeze(-1)).unsqueeze(1)
+
+        cos, sin = self.rope(L, idx.device)
+        cos = cos.unsqueeze(0).unsqueeze(0)        # [1,1,L,hd]
+        sin = sin.unsqueeze(0).unsqueeze(0)
+
+        if attn_mask is None:
+            # AR default causal; DIFF default bidirectional (None)
+            attn_mask = self._causal_mask(L, idx.device) if mode[0].item() == 0 else None
+        else:
+            attn_mask = attn_mask.to(idx.device)
+
+        # gradient checkpointing: trade ~20% compute for ~4x less
+        # activation memory, so a big batch fits on 16 GB.
+        training = self.training and torch.is_grad_enabled()
+        for blk in self.blocks:
+            if training:
+                x = checkpoint(blk, x, cos, sin, attn_mask,
+                               use_reentrant=False)
+            else:
+                x = blk(x, cos, sin, attn_mask)
+        x = self.norm(x)
+        return self.lm_head(x)                     # [B, L, vocab]

prep.py

@@ -0,0 +1,250 @@
+"""
+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})
+"""
+import os, sys, json, random, argparse
+import numpy as np
+from datasets import load_dataset
+
+import tokenizer as tokmod
+from tokenizer import YKTokenizer, SPECIAL
+
+random.seed(1234)
+np.random.seed(1234)
+
+OUT = os.path.dirname(os.path.abspath(__file__))
+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
+
+
+# --------------------------------------------------------------------------
+# 1) Tokenizer training
+# --------------------------------------------------------------------------
+def train_tokenizer():
+    print("[prep] streaming FineWeb-edu to collect tokenizer training docs ...")
+    ds = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT",
+                      streaming=True, split="train")
+    texts = []
+    for i, ex in enumerate(ds):
+        texts.append(ex["text"])
+        if i + 1 >= TOK_TRAIN_DOCS:
+            break
+    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"))
+    print(f"[prep] tokenizer trained: vocab={tok.vocab_size}")
+    return tok
+
+
+# --------------------------------------------------------------------------
+# 2) Synthetic tool-use conversations (tag format == agent.py)
+# --------------------------------------------------------------------------
+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?",
+]
+OPS = {"+": "plus", "-": "minus", "*": "times", "/": "divided by"}
+PY_SNIPPETS = [
+    "print(sum(range(1, {n}+1)))",
+    "import math\nprint(round(math.sqrt({n}), 4))",
+    "print(sorted([{a}, {b}, {c}]))",
+    "print({n} ** 2 + {n})",
+]
+FILE_Q = [
+    "Read the file {path} and tell me what is on the first line.",
+    "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.")
+
+
+def gen_synthetic(n=25000):
+    out = []
+    for _ in range(n):
+        kind = random.random()
+        if kind < 0.5:
+            a = random.randint(2, 999); b = random.randint(2, 999)
+            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}")
+            q = random.choice(CALC_TEMPLATES).format(a=a, b=b, op=OPS[op])
+            tool = f'<tool name="calc">{a} {op} {b}</tool>'
+            result = str(ans)
+            think = f"<think>The user wants {a} {OPS[op]} {b}. I'll use the calculator.</think>"
+        elif kind < 0.8:
+            n_ = random.randint(3, 200); a = random.randint(1, 50); b = random.randint(1, 50); c = random.randint(1, 50)
+            code = random.choice(PY_SNIPPETS).format(n=n_, a=a, b=b, c=c)
+            q = f"Run this tiny Python snippet and report the output:\n{code}"
+            tool = f'<tool name="python">{code}</tool>'
+            try:
+                import io, contextlib
+                buf = io.StringIO(); 
+                with contextlib.redirect_stdout(buf):
+                    exec(code, {"__builtins__": __builtins__}, {})
+                result = buf.getvalue().strip()
+            except Exception as e:
+                result = f"error: {e}"
+            think = "<think>I can execute this with the python tool.</think>"
+        else:
+            path = random.choice(["notes.txt", "data/log.csv", "README.md", "config.json"])
+            q = random.choice(FILE_Q).format(path=path, dir=random.choice(["src", "data", "."]))
+            if "List" in q or "list" in q:
+                tool = f'<tool name="list_dir">{path}</tool>'
+                result = f"{path}/\n  file_a.txt\n  file_b.csv"
+            else:
+                tool = f'<tool name="read_file">{path}</tool>'
+                result = f"line 1: hello from {path}"
+            think = "<think>I should read the file with the read_file tool.</think>"
+
+        conv = (f"<bos><system>{SYSTEM}</system>"
+                f"<user>{q}</user>"
+                f"<assistant>{think}{tool}<result>{result}</result>"
+                f"Based on the tool result, the answer is {result}.</assistant><eos>")
+        out.append(conv)
+    return out
+
+
+# --------------------------------------------------------------------------
+# 3) Glaive function-calling (best-effort)
+# --------------------------------------------------------------------------
+def gen_glaive(max_examples=20000):
+    out = []
+    try:
+        ds = load_dataset("glaiveai/glaive-function-calling-v2",
+                          streaming=True, split="train")
+    except Exception as e:
+        print(f"[prep] Glaive unavailable ({e}); skipping.")
+        return out
+    for i, ex in enumerate(ds):
+        if i >= max_examples:
+            break
+        try:
+            conv = ex["conversations"]
+            parts = ["<bos>"]
+            for m in conv:
+                role = m.get("role") or m.get("from")
+                val = m.get("value") or m.get("content") or ""
+                if role in ("system", "system_prompt"):
+                    parts.append(f"<system>{val}</system>")
+                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":
+                    parts.append(f"<result>{val}</result>")
+            parts.append("<eos>")
+            out.append("".join(parts))
+        except Exception:
+            continue
+    print(f"[prep] Glaive converted: {len(out)} examples")
+    return out
+
+
+# --------------------------------------------------------------------------
+# 4) 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:
+        for text in texts:
+            ids = tok.encode(text)
+            if not ids:
+                continue
+            buf.extend(ids)
+            while len(buf) >= seq_len:
+                chunk = np.array(buf[:seq_len], dtype=np.uint16)
+                f.write(chunk.tobytes())
+                buf = buf[seq_len:]
+                n += seq_len
+                if n >= budget:
+                    return n
+    if buf:
+        chunk = np.array(buf[:seq_len], dtype=np.uint16)
+        if len(chunk) == seq_len:
+            with open(bin_path, "ab") as f:
+                f.write(chunk.tobytes())
+            n += seq_len
+    return n
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--seed-tok-only", action="store_true")
+    args = ap.parse_args()
+
+    tok_path = os.path.join(DATADIR, "tokenizer.json")
+    if os.path.exists(tok_path):
+        print("[prep] loading existing tokenizer")
+        tok = YKTokenizer.load(tok_path)
+    else:
+        tok = train_tokenizer()
+
+    bin_path = os.path.join(DATADIR, "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) ...")
+    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
+        next(gen)
+    def fineweb_iter():
+        for ex in gen:
+            yield ex["text"]
+    n = pack(tok, fineweb_iter(), bin_path, TOK_BUDGET, SEQ_LEN)
+    print(f"[prep] fineweb packed: {n:,} tokens")
+
+    # --- 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")
+
+    # --- glaive ---
+    gl = gen_glaive(20000)
+    n3 = 0
+    if gl:
+        n3 = pack(tok, gl, bin_path, TOK_BUDGET_TOOL, SEQ_LEN)
+
+    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:
+        json.dump(meta, f)
+    print(f"[prep] DONE. total tokens={total:,}  vocab={tok.vocab_size}")
+    print(f"[prep] files: {bin_path}")
+
+
+if __name__ == "__main__":
+    main()

push_hf.py

@@ -0,0 +1,83 @@
+"""
+clankerDiffusion — push artifacts to HuggingFace (via the `hf` CLI, which
+correctly creates the first commit on a brand-new repo).
+
+  python push_hf.py --what all        # code + tokenizer + rag corpus + train.bin
+  python push_hf.py --what code       # .py sources + tokenizer + rag corpus
+  python push_hf.py --what data       # train.bin + tokenizer + rag corpus
+  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)
+"""
+import os
+import subprocess
+import argparse
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+DATADIR = os.path.join(HERE, "data")
+CKPTDIR = os.path.join(HERE, "checkpoints")
+
+CODE_REPO = "coderofpears/clankerDiffusion-base"
+DATA_REPO = "coderofpears/clankerDiffusion-data"
+CKPT_REPO = "coderofpears/clankerDiffusion-checkpoints"
+
+PY_SOURCES = ["model.py", "tokenizer.py", "prep.py", "train.py", "infer.py",
+              "tools.py", "agent.py", "rag.py", "build_rag.py", "rag_finetune.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 _upload(repo, local, patterns):
+    # upload only the matching files via a temp staging dir
+    import tempfile, shutil, fnmatch
+    stage = tempfile.mkdtemp()
+    try:
+        for root, _, files in os.walk(local):
+            for f in files:
+                rel = os.path.relpath(os.path.join(root, f), local)
+                if any(fnmatch.fnmatch(rel, p) or fnmatch.fnmatch(f, p) for p in patterns):
+                    dst = os.path.join(stage, rel)
+                    os.makedirs(os.path.dirname(dst), exist_ok=True)
+                    shutil.copy(os.path.join(root, f), dst)
+        if not os.listdir(stage):
+            print(f"[hf] nothing matched for {repo}")
+            return
+        _run(["hf", "upload", repo, stage, "--repo-type", "model"])
+        print(f"[hf] -> {repo}")
+    finally:
+        shutil.rmtree(stage, ignore_errors=True)
+
+
+def push_code():
+    _upload(CODE_REPO, HERE, PY_SOURCES + ["data/tokenizer.json", "data/rag_corpus.txt", "README.md"])
+
+
+def push_data():
+    _upload(DATA_REPO, DATADIR, ["train.bin", "tokenizer.json", "rag_corpus.txt"])
+
+
+def push_ckpt():
+    _upload(CKPT_REPO, CKPTDIR, ["*.pt"])
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--what", default="all", choices=["all", "code", "data", "ckpt"])
+    a = ap.parse_args()
+    if a.what in ("all", "code"):
+        push_code()
+    if a.what in ("all", "data"):
+        push_data()
+    if a.what in ("all", "ckpt"):
+        push_ckpt()
+
+
+if __name__ == "__main__":
+    main()

rag.py

@@ -0,0 +1,163 @@
+"""
+clankerDiffusion — local RAG / knowledge-injection subsystem.
+
+A small, dependency-free BM25 retriever so the agent can pull relevant
+passages from a local knowledge base and inject them as <context> blocks.
+
+Two ways knowledge gets injected:
+  1. Model-driven : the model emits <tool name="retrieve">query</tool> and the
+     agent runs it, appending the <result> and continuing (ReAct).
+  2. Controller-driven : the agent's RAG controller watches the response as it
+     is generated and, mid-turn, retrieves passages for the current question and
+     injects <context>...</context> right into the stream -- "knowledge injection
+     even in the middle of a response" -- without the model having to ask.
+
+BM25 is used by default (no model downloads, runs anywhere). If
+sentence-transformers is available it is used as an optional re-ranker.
+"""
+import os
+import re
+import json
+import math
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+DATADIR = os.path.join(HERE, "data")
+KB_DIR = os.path.join(DATADIR, "kb")
+DEFAULT_INDEX = os.path.join(KB_DIR, "index.json")
+
+_TOKEN_RE = re.compile(r"[a-z0-9]+")
+
+
+def _tok(s):
+    return _TOKEN_RE.findall(s.lower())
+
+
+def _chunk(text, size=180, stride=90):
+    """Split text into token-window chunks with overlap."""
+    toks = _tok(text)
+    if not toks:
+        return []
+    out = []
+    i = 0
+    while i < len(toks):
+        out.append(" ".join(toks[i:i + size]))
+        if i + size >= len(toks):
+            break
+        i += stride
+    return out
+
+
+class KnowledgeBase:
+    def __init__(self):
+        self.docs = []          # list[str]  chunk texts
+        self._df = {}           # term -> doc freq
+        self._N = 0
+        self._avgdl = 1.0
+        self._built = False
+
+    # -- ingest -------------------------------------------------------------
+    def add_text(self, text, source=""):
+        for ch in _chunk(text):
+            if ch:
+                self.docs.append(ch)
+        self._built = False
+
+    def ingest_path(self, path):
+        """Ingest a file or a directory of files into the KB."""
+        if os.path.isdir(path):
+            files = []
+            for root, _, fs in os.walk(path):
+                for f in fs:
+                    if f.lower().endswith((".txt", ".md", ".py", ".csv", ".json", ".log")):
+                        files.append(os.path.join(root, f))
+        else:
+            files = [path]
+        for fp in files:
+            try:
+                with open(fp, "r", errors="replace") as fh:
+                    self.add_text(fh.read(), source=fp)
+            except Exception as e:
+                print(f"[rag] skip {fp}: {e}")
+        print(f"[rag] ingested {len(files)} file(s) -> {len(self.docs)} chunks")
+
+    # -- index --------------------------------------------------------------
+    def _build(self):
+        self._df = {}
+        self._N = len(self.docs)
+        lengths = []
+        for d in self.docs:
+            seen = set()
+            for t in _tok(d):
+                seen.add(t)
+            lengths.append(len(_tok(d)))
+            for t in seen:
+                self._df[t] = self._df.get(t, 0) + 1
+        self._avgdl = (sum(lengths) / self._N) if self._N else 1.0
+        self._built = True
+
+    def retrieve(self, query, k=4):
+        if not self._built:
+            self._build()
+        if self._N == 0:
+            return []
+        q_toks = _tok(query)
+        if not q_toks:
+            return []
+        k1, b = 1.5, 0.75
+        scores = []
+        for d in self.docs:
+            dtoks = _tok(d)
+            dl = len(dtoks)
+            tf = {}
+            for t in dtoks:
+                tf[t] = tf.get(t, 0) + 1
+            s = 0.0
+            for t in q_toks:
+                if t not in self._df:
+                    continue
+                idf = math.log((self._N - self._df[t] + 0.5) / (self._df[t] + 0.5) + 1.0)
+                f = tf.get(t, 0)
+                s += idf * (f * (k1 + 1)) / (f + k1 * (1 - b + b * dl / self._avgdl))
+            scores.append(s)
+        order = sorted(range(self._N), key=lambda i: scores[i], reverse=True)
+        return [self.docs[i] for i in order[:k] if scores[i] > 0]
+
+    # -- persist ------------------------------------------------------------
+    def save(self, path=DEFAULT_INDEX):
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        json.dump({"docs": self.docs}, open(path, "w"), ensure_ascii=False)
+        print(f"[rag] saved index -> {path} ({len(self.docs)} chunks)")
+
+    @classmethod
+    def load(cls, path=DEFAULT_INDEX):
+        kb = cls()
+        if os.path.exists(path):
+            data = json.load(open(path, "r", encoding="utf-8"))
+            kb.docs = data.get("docs", [])
+            kb._built = False
+            print(f"[rag] loaded index {path} ({len(kb.docs)} chunks)")
+        return kb
+
+
+# A process-wide default KB, lazily built from data/kb/.
+_DEFAULT_KB = None
+
+
+def default_kb():
+    global _DEFAULT_KB
+    if _DEFAULT_KB is None:
+        if os.path.exists(DEFAULT_INDEX):
+            _DEFAULT_KB = KnowledgeBase.load(DEFAULT_INDEX)
+        elif os.path.isdir(KB_DIR):
+            kb = KnowledgeBase()
+            kb.ingest_path(KB_DIR)
+            _DEFAULT_KB = kb
+        else:
+            _DEFAULT_KB = KnowledgeBase()
+    return _DEFAULT_KB
+
+
+def retrieve(query, k=4, kb=None):
+    """Top-level retrieval used by the `retrieve` tool + RAG controller."""
+    kb = kb or default_kb()
+    return kb.retrieve(query, k=k)

rag_finetune.py

@@ -0,0 +1,129 @@
+"""
+clankerDiffusion — RAG continuation fine-tune.
+
+Loads the latest base checkpoint and continues training on the RAG corpus
+(data/rag_corpus.txt) so the model actually learns to (a) emit
+<tool name="retrieve">q</tool> and use the <result>, and (b) read
+<context>...</context> injected mid-response. Run this AFTER the base
+24h run (or anytime you have a checkpoint).
+
+    python rag_finetune.py --hours 3 --batch 24 --ckpt-every 500
+
+It reuses the hybrid AR/DIFF loss from train.py's logic (self-contained copy
+below so it doesn't import the training script's __main__).
+"""
+import os, json, argparse, time, random
+import torch
+import torch.nn.functional as F
+from torch.utils.data import IterableDataset
+
+from model import YKDiff
+from tokenizer import YKTokenizer
+from rag import DEFAULT_INDEX, KnowledgeBase
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+DATADIR = os.path.join(HERE, "data")
+CKPTDIR = os.path.join(HERE, "checkpoints")
+CORPUS = os.path.join(DATADIR, "rag_corpus.txt")
+SEQ = 1024
+
+
+# ---- hybrid loss (mirrors train.py) ---------------------------------------
+def ar_loss(logits, ids, mask):
+    sl = ids[:, 1:]; lm = logits[:, :-1]
+    return F.cross_entropy(lm.reshape(-1, lm.size(-1)), sl.reshape(-1), reduction="none").mean()
+
+
+def diff_loss(logits, ids, mask, t):
+    B, L, V = logits.shape
+    ce = F.cross_entropy(logits.reshape(-1, V), ids.reshape(-1), reduction="none").view(B, L)
+    w = (mask * (1.0 - t)[:, None]).reshape(B, L)
+    return (ce * w).sum() / w.sum().clamp_min(1.0)
+
+
+# ---- corpus -> token stream ----------------------------------------------
+class RagTokens(IterableDataset):
+    def __init__(self, path, tok, seq):
+        self.lines = [l.rstrip("\n") for l in open(path, encoding="utf-8") if l.strip()]
+        self.tok = tok
+        self.seq = seq
+        self.buf = []
+
+    def _fill(self):
+        while len(self.buf) < self.seq + 1:
+            line = random.choice(self.lines)
+            self.buf.extend(self.tok.encode(line))
+
+    def __iter__(self):
+        while True:
+            self._fill()
+            chunk = self.buf[: self.seq + 1]
+            self.buf = self.buf[self.seq:]
+            yield torch.tensor(chunk, dtype=torch.long)
+
+
+def pack_batch(ds, n):
+    ids = torch.stack([next(iter(ds)) for _ in range(n)])         # [B, L+1]
+    mask = (ids != ds.tok.pad_id).long()
+    return ids, mask
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--hours", type=float, default=3.0)
+    ap.add_argument("--batch", type=int, default=24)
+    ap.add_argument("--ckpt-every", type=int, default=500)
+    ap.add_argument("--log-every", type=int, default=25)
+    ap.add_argument("--ckpt", default=None)
+    a = ap.parse_args()
+
+    tok = YKTokenizer.load(os.path.join(DATADIR, "tokenizer.json"))
+    if a.ckpt is None:
+        ck = sorted([f for f in os.listdir(CKPTDIR) if f.endswith(".pt")])
+        if not ck:
+            raise SystemExit("no checkpoint")
+        a.ckpt = os.path.join(CKPTDIR, ck[-1])
+    sd = torch.load(a.ckpt, map_location="cuda")
+    model = YKDiff(sd["cfg"]).cuda().train()
+    model.load_state_dict(sd["model"])
+    step0 = sd.get("step", 0)
+    print(f"[rag-ft] resume {a.ckpt} step={step0} params={sum(p.numel() for p in model.parameters())/1e6:.1f}M")
+
+    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01, betas=(0.9, 0.95))
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(1, int(a.hours * 60)), eta_min=1e-5)
+    ds = RagTokens(CORPUS, tok, SEQ)
+
+    t0 = time.time()
+    limit = a.hours * 3600
+    step = step0
+    while time.time() - t0 < limit:
+        step += 1
+        ids, mask = pack_batch(ds, a.batch)
+        x = ids[:, :-1].cuda(); y = ids[:, 1:].cuda(); m = mask[:, 1:].cuda()
+        if random.random() < 0.5:
+            mode = torch.zeros(a.batch, dtype=torch.long, device="cuda")
+            t = None
+            loss = ar_loss(model(x, mode), y, m)
+        else:
+            mode = torch.ones(a.batch, dtype=torch.long, device="cuda")
+            t = torch.rand(a.batch, device="cuda")
+            r = (0.1 + 0.9 * torch.rand(a.batch, device="cuda"))
+            inps = torch.where(torch.rand_like(y.float()) < r[:, None], tok.mask_id, y)
+            loss = diff_loss(model(inps, mode, t=t), y, m, t)
+        opt.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step(); sched.step()
+        if step % a.log_every == 0:
+            print(f"[rag-ft] step {step} loss={loss.item():.3f} t={ (time.time()-t0)/60:.1f}m", flush=True)
+        if step % a.ckpt_every == 0:
+            torch.save({"model": model.state_dict(), "cfg": sd["cfg"], "step": step,
+                        "rag_ft": True}, os.path.join(CKPTDIR, f"ragft_{step:06d}.pt"))
+            print(f"[rag-ft] saved ragft_{step:06d}.pt", flush=True)
+    torch.save({"model": model.state_dict(), "cfg": sd["cfg"], "step": step, "rag_ft": True},
+               os.path.join(CKPTDIR, "ragft_final.pt"))
+    print("[rag-ft] done", flush=True)
+
+
+if __name__ == "__main__":
+    main()

tokenizer.py

@@ -0,0 +1,105 @@
+"""
+From-scratch byte-level BPE tokenizer, trained on the corpus (no pretrained weights).
+Special tokens are reserved for the hybrid model:
+  <pad>  <bos>  <eos>  <mask>  <think>  </think>  <tool>  </tool>  <result>
+"""
+import json
+import os
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer
+from tokenizers.pre_tokenizers import ByteLevel
+from tokenizers.decoders import ByteLevel as ByteLevelDecoder
+from tokenizers import AddedToken
+
+SPECIAL = ["<pad>", "<bos>", "<eos>", "<mask>",
+            "<think>", "</think>", "<tool>", "</tool>", "<result>"]
+
+
+class YKTokenizer:
+    def __init__(self, path=None):
+        self.path = path
+        self.tok = None
+        self._ids = {}
+
+    # ---- training --------------------------------------------------------
+    def train(self, iterator, vocab_size=32768, save_path=None):
+        assert vocab_size > len(SPECIAL)
+        model = BPE(unk_token="<unk>")
+        self.tok = Tokenizer(model)
+        self.tok.pre_tokenizer = ByteLevel()
+        self.tok.decoder = ByteLevelDecoder()
+        trainer = BpeTrainer(
+            vocab_size=vocab_size,
+            special_tokens=SPECIAL + ["<unk>"],
+            show_progress=True,
+        )
+        self.tok.train_from_iterator(iterator, trainer)
+        self._refresh()
+        if save_path:
+            self.save(save_path)
+        return self
+
+    # ---- load / save ----------------------------------------------------
+    def save(self, path):
+        self.path = path
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        self.tok.save(path)
+        with open(path + ".meta.json", "w") as f:
+            json.dump({"vocab_size": self.vocab_size,
+                       "special_ids": self._ids}, f)
+
+    @classmethod
+    def load(cls, path):
+        obj = cls(path)
+        obj.tok = Tokenizer.from_file(path)
+        meta = path + ".meta.json"
+        if os.path.exists(meta):
+            with open(meta) as f:
+                m = json.load(f)
+            obj._ids = m["special_ids"]
+        obj._refresh()
+        return obj
+
+    def _refresh(self):
+        self._ids = {s: self.tok.token_to_id(s) for s in SPECIAL}
+        self._ids["<unk>"] = self.tok.token_to_id("<unk>")
+        self.vocab_size = self.tok.get_vocab_size()
+
+    # ---- ids -------------------------------------------------------------
+    @property
+    def pad_id(self):   return self._ids["<pad>"]
+    @property
+    def bos_id(self):   return self._ids["<bos>"]
+    @property
+    def eos_id(self):   return self._ids["<eos>"]
+    @property
+    def mask_id(self):  return self._ids["<mask>"]
+    @property
+    def think_id(self):  return self._ids["<think>"]
+    @property
+    def endthink_id(self): return self._ids["</think>"]
+    @property
+    def tool_id(self):  return self._ids["<tool>"]
+    @property
+    def endtool_id(self): return self._ids["</tool>"]
+    @property
+    def result_id(self): return self._ids["<result>"]
+
+    # ---- encode / decode ------------------------------------------------
+    def encode(self, text, add_special=False):
+        ids = self.tok.encode(text).ids
+        if add_special:
+            ids = [self.bos_id] + ids + [self.eos_id]
+        return ids
+
+    def encode_batch(self, texts):
+        return [t.ids for t in self.tok.encode_batch(texts)]
+
+    def decode(self, ids, skip_special=True):
+        if skip_special:
+            ids = [i for i in ids if i not in self._ids.values() or i == self._ids.get("<unk>", -1)]
+        return self.tok.decode(ids, skip_special_tokens=skip_special)
+
+    def __len__(self):
+        return self.vocab_size

tools.py

@@ -0,0 +1,127 @@
+"""
+clankerDiffusion — local tool executor for the agent loop.
+
+Tools the model can call via  <tool name="...">args</tool> :
+  calc(expr)      safe arithmetic / math expression
+  python(code)    run Python, stdout captured (threaded timeout)
+  read_file(p)    read a text file
+  list_dir(p)     list a directory
+  retrieve(q)     RAG search over the local knowledge base (returns passages)
+
+NOTE: python() executes locally on YOUR machine. It is sandboxed with a
+restricted builtins set and a wall-clock timeout, but treat it as you would
+any local code. Don't point clanker at untrusted input.
+"""
+import ast, math, io, contextlib, os, time, threading, builtins
+from rag import retrieve as _rag_retrieve
+
+SAFE_NAMES = {k: getattr(math, k) for k in dir(math) if not k.startswith("_")}
+SAFE_NAMES.update({"abs": abs, "min": min, "max": max, "round": round,
+                   "pow": pow, "len": len, "sum": sum, "sorted": sorted,
+                   "int": int, "float": float, "str": str, "list": list,
+                   "range": range, "True": True, "False": False, "None": None})
+
+
+def _safe_eval(expr):
+    node = ast.parse(expr, mode="eval").body
+
+    def ev(n):
+        if isinstance(n, ast.BinOp):
+            a, b = ev(n.left), ev(n.right)
+            if isinstance(n.op, ast.Add): return a + b
+            if isinstance(n.op, ast.Sub): return a - b
+            if isinstance(n.op, ast.Mult): return a * b
+            if isinstance(n.op, ast.Div): return a / b
+            if isinstance(n.op, ast.FloorDiv): return a // b
+            if isinstance(n.op, ast.Mod): return a % b
+            if isinstance(n.op, ast.Pow): return a ** b
+            raise ValueError("op")
+        if isinstance(n, ast.UnaryOp):
+            v = ev(n.operand)
+            return -v if isinstance(n.op, ast.USub) else +v
+        if isinstance(n, ast.Constant) and isinstance(n.value, (int, float)):
+            return n.value
+        if isinstance(n, ast.Call) and isinstance(n.func, ast.Name):
+            fn = SAFE_NAMES.get(n.func.id)
+            if fn is None: raise ValueError(f"no {n.func.id}")
+            return fn(*[ev(a) for a in n.args])
+        if isinstance(n, ast.Name):
+            if n.id in SAFE_NAMES: return SAFE_NAMES[n.id]
+            raise ValueError(f"name {n.id}")
+        raise ValueError("bad expr")
+    return ev(node)
+
+
+def _run_python(code, timeout=10.0):
+    ns = {"__builtins__": {k: getattr(builtins, k)
+                            for k in ("print", "len", "range", "list", "dict",
+                                      "tuple", "set", "str", "int", "float",
+                                      "bool", "min", "max", "sum", "sorted",
+                                      "abs", "enumerate", "zip", "map", "filter",
+                                      "open", "round", "pow") if hasattr(builtins, k)}}
+    ns["math"] = math
+    buf = io.StringIO()
+    result = {}
+
+    def target():
+        try:
+            with contextlib.redirect_stdout(buf):
+                exec(code, ns)
+            result["out"] = buf.getvalue()
+        except Exception as e:
+            result["out"] = buf.getvalue() + f"\nERROR: {type(e).__name__}: {e}"
+
+    t = threading.Thread(target=target, daemon=True)
+    t.start(); t.join(timeout)
+    if t.is_alive():
+        return f"ERROR: execution timed out after {timeout}s"
+    return result.get("out", "").strip() or "(no output)"
+
+
+def execute_tool(name, arg):
+    try:
+        if name == "calc":
+            val = _safe_eval(arg.strip())
+            return f"{arg.strip()} = {val}"
+        if name == "python":
+            return _run_python(arg)
+        if name == "read_file":
+            p = arg.strip().strip('"\'')
+            if not os.path.exists(p): return f"ERROR: no such file: {p}"
+            with open(p, "r", errors="replace") as f:
+                return f.read()[:4000]
+        if name == "list_dir":
+            p = arg.strip().strip('"\'') or "."
+            if not os.path.isdir(p): return f"ERROR: no such dir: {p}"
+            return "\n".join(sorted(os.listdir(p))[:200])
+        if name == "retrieve":
+            try:
+                passages = _rag_retrieve(arg.strip(), k=4)
+                if not passages:
+                    return "(no relevant context found in the knowledge base)"
+                return "\n---\n".join(passages)
+            except Exception as e:
+                return f"ERROR: retrieve failed: {e}"
+        return f"ERROR: unknown tool '{name}'"
+    except Exception as e:
+        return f"ERROR: {type(e).__name__}: {e}"
+
+
+def parse_tool_calls(text):
+    """Return list of (name, argstring) from <tool name=\"...\">..</tool> blocks."""
+    calls = []
+    i = 0
+    while True:
+        s = text.find("<tool", i)
+        if s < 0: break
+        e = text.find("</tool>", s)
+        if e < 0: break
+        head = text[s:text.find(">", s) + 1]
+        # head like: <tool name="calc">
+        import re
+        m = re.search(r'name="([^"]+)"', head)
+        name = m.group(1) if m else "?"
+        arg = text[s + len(head):e].strip()
+        calls.append((name, arg))
+        i = e + len("</tool>")
+    return calls

train.py

@@ -0,0 +1,145 @@
+"""
+clankerDiffusion — local training loop (RTX 5060 Ti, ~5h budget).
+
+Hybrid objective (per step, mode chosen at random, p(AR)=0.5):
+  AR   (mode 0): causal LM cross-entropy over the whole window.
+  DIFF (mode 1): MDLM absorbing-state masked diffusion — mask each token
+                 independently with ratio r~U(0,1); reconstruct masked tokens
+                 with bidirectional attention, conditioned on r via the time embed.
+
+Runs in bf16 (Blackwell), AdamW + cosine LR, grad-clip, checkpoints locally.
+Stops after --hours (default 5) of wall-clock time.
+"""
+import os, json, time, argparse
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from model import YKDiff
+from tokenizer import YKTokenizer
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+DATADIR = os.path.join(HERE, "data")
+CKPTDIR = os.path.join(HERE, "checkpoints")
+os.makedirs(CKPTDIR, exist_ok=True)
+
+# ---- architecture for 16 GB -------------------------------------------------
+CFG = dict(
+    d_model=768, n_layers=12, n_heads=12, d_ff=2048,
+    max_len=1024, vocab_size=32768,
+)
+
+
+def load_data():
+    meta = json.load(open(os.path.join(DATADIR, "meta.json")))
+    arr = np.memmap(os.path.join(DATADIR, "train.bin"), dtype=np.uint16, mode="r")
+    CFG["vocab_size"] = meta["vocab_size"]
+    CFG["max_len"] = meta["seq_len"]
+    print(f"[train] data n_tokens={meta['n_tokens']:,} seq_len={meta['seq_len']} "
+          f"vocab={meta['vocab_size']}")
+    return arr, meta["seq_len"]
+
+
+def sample_batch(arr, seq_len, batch):
+    N = len(arr)
+    starts = np.random.randint(0, N - seq_len, size=batch)
+    out = np.stack([arr[s:s + seq_len].astype(np.int64) for s in starts])
+    return torch.from_numpy(out).long()
+
+
+def train(args):
+    tok = YKTokenizer.load(os.path.join(DATADIR, "tokenizer.json"))
+    arr, seq_len = load_data()
+
+    model = YKDiff(CFG).cuda()
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"[train] model params = {n_params/1e6:.1f}M")
+
+    optim = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95),
+                              weight_decay=0.1)
+    V = CFG["vocab_size"]
+    pad_id = tok.pad_id
+    mask_id = tok.mask_id
+
+    # resume
+    step0 = 0
+    ckpts = sorted([f for f in os.listdir(CKPTDIR) if f.endswith(".pt")])
+    if ckpts and not args.fresh:
+        path = os.path.join(CKPTDIR, ckpts[-1])
+        sd = torch.load(path, map_location="cuda")
+        model.load_state_dict(sd["model"]); optim.load_state_dict(sd["optim"])
+        step0 = sd["step"]
+        print(f"[train] resumed from {path} step={step0}")
+
+    model.train()
+    amp = torch.cuda.amp.autocast(dtype=torch.bfloat16)
+    t0 = time.time()
+    limit = args.hours * 3600.0
+    step = step0
+    running = 0.0
+
+    while True:
+        if time.time() - t0 > limit:
+            print(f"[train] wall-clock limit {args.hours}h reached at step {step}")
+            break
+
+        optim.zero_grad(set_to_none=True)
+        mode_ar = (torch.rand(1).item() < 0.5)
+        idx = sample_batch(arr, seq_len, args.batch).cuda()
+
+        with amp:
+            if mode_ar:
+                m = torch.zeros(args.batch, dtype=torch.long, device="cuda")
+                logits = model(idx, m, t=None)                 # causal
+                loss = F.cross_entropy(
+                    logits[:, :-1].reshape(-1, V),
+                    idx[:, 1:].reshape(-1), ignore_index=pad_id)
+                mname = "AR"
+            else:
+                m = torch.ones(args.batch, dtype=torch.long, device="cuda")
+                r = torch.rand(args.batch, device="cuda")       # per-sample ratio
+                is_mask = torch.rand(args.batch, seq_len, device="cuda") < r[:, None]
+                not_pad = idx != pad_id
+                masked = idx.clone(); masked[is_mask] = mask_id
+                logits = model(masked, m, t=r)
+                ce = F.cross_entropy(logits.reshape(-1, V), idx.reshape(-1),
+                                    reduction="none", ignore_index=-100)
+                ce = ce * is_mask.reshape(-1) * not_pad.reshape(-1)
+                denom = (is_mask & not_pad).reshape(-1).sum().clamp(min=1)
+                loss = ce.sum() / denom
+                mname = "DIFF"
+
+        loss.backward()
+        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optim.step()
+
+        running = running * 0.9 + float(loss.item()) * 0.1
+        step += 1
+        if step % args.log_every == 0:
+            print(f"[train] step {step} [{mname}] loss={running:.3f} "
+                  f"t={(time.time()-t0)/60:.1f}m", flush=True)
+
+        if step % args.ckpt_every == 0:
+            path = os.path.join(CKPTDIR, f"clanker_{step:07d}.pt")
+            torch.save({"model": model.state_dict(), "optim": optim.state_dict(),
+                        "step": step, "cfg": CFG, "vocab": V}, path)
+            print(f"[train] checkpoint -> {path}", flush=True)
+
+    # final save
+    path = os.path.join(CKPTDIR, f"clanker_{step:07d}_final.pt")
+    torch.save({"model": model.state_dict(), "optim": optim.state_dict(),
+                "step": step, "cfg": CFG, "vocab": V}, path)
+    json.dump(CFG, open(os.path.join(CKPTDIR, "config.json"), "w"))
+    print(f"[train] DONE final={path} steps={step}")
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--hours", type=float, default=5.0)
+    ap.add_argument("--batch", type=int, default=32)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--log-every", type=int, default=25)
+    ap.add_argument("--ckpt-every", type=int, default=500)
+    ap.add_argument("--fresh", action="store_true")
+    train(ap.parse_args())

upload_artifacts.py

@@ -0,0 +1,58 @@
+"""
+upload_artifacts.py  —  run THIS locally before launching the Colab notebook.
+
+Publishes the model code + trained tokenizer + (optionally) your local
+checkpoint to HuggingFace Hub so the Colab script can pull them.
+
+  set HF_TOKEN=your_huggingface_token
+  python upload_artifacts.py            # code + tokenizer
+  python upload_artifacts.py --with-ckpt   # also upload latest local checkpoint
+"""
+import os, sys, argparse
+from huggingface_hub import HfApi
+
+ROOT = r"C:\Users\User\CalcGPU\clankerDiffusion"
+CODE_REPO = "clankerDiffusion/base"
+CKPT_REPO = "clankerDiffusion/checkpoints"
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--with-ckpt", action="store_true")
+    a = ap.parse_args()
+
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        sys.exit("set HF_TOKEN first:  set HF_TOKEN=hf_xxx")
+    api = HfApi(token=token)
+
+    api.create_repo(CODE_REPO, repo_type="model", exist_ok=True)
+    code_files = ["model.py", "tokenizer.py", "train.py", "prep.py", "infer.py",
+                  os.path.join("data", "tokenizer.json"),
+                  os.path.join("data", "tokenizer.json.meta.json")]
+    for rel in code_files:
+        p = os.path.join(ROOT, rel)
+        if os.path.exists(p):
+            api.upload_file(repo_id=CODE_REPO,
+                           path_in_repo=os.path.basename(p),
+                           path_or_fileobj=p)
+            print("uploaded", rel)
+        else:
+            print("skip (missing)", rel)
+
+    if a.with_ckpt:
+        api.create_repo(CKPT_REPO, repo_type="model", exist_ok=True)
+        ckpt_dir = os.path.join(ROOT, "checkpoints")
+        if os.path.isdir(ckpt_dir):
+            pts = sorted(f for f in os.listdir(ckpt_dir) if f.endswith(".pt"))
+            if pts:
+                p = os.path.join(ckpt_dir, pts[-1])
+                api.upload_file(repo_id=CKPT_REPO,
+                               path_in_repo=os.path.basename(p),
+                               path_or_fileobj=p)
+                print("uploaded checkpoint", pts[-1])
+    print("DONE. Colab CODE_REPO =", CODE_REPO)
+
+
+if __name__ == "__main__":
+    main()