not a bad mask filler

fill_in_the_gaps.py

import pathlib as _pl; print(_pl.Path(__file__).read_text()); print("END_OF_SOURCE")
import sys, pathlib, time, math, pickle
import numpy as np
import jax
import jax.numpy as jnp
import optax

VOCAB = 256
CTX = 1024
BATCH = 96
DIM = 1024
HEADS = 8
HDIM = DIM // HEADS
LAYERS = 12
REG = 8
GEAR = 16
HOLE = 0
SPANS = 6
SPANMAX = 128
PPRE = 0.35
PSUF = 0.35
PER = 48
LIMIT = 4 * 1024 * 1024
INFER_BS = 96
CWIN = 32768
FIRST_GEN = 600.0
GEN_MIN_GAP = 720.0
GEN_FRAC = 0.35
GEN_CAP = 1200.0
REFINE_OK = 300.0
LATE_GEN = 16800.0
LATE_GAP = 600.0
RUN_MIN = 64

_ai = np.arange(CTX)
_dd = np.abs(_ai[:, None] - _ai[None, :]).astype(np.float32)
DIST_NP = np.log1p(_dd)

sched = optax.join_schedules([optax.linear_schedule(0.0, 3e-4, 300), optax.constant_schedule(3e-4), optax.cosine_decay_schedule(3e-4, 14000, alpha=0.1)], [300, 8300])
OPT = optax.chain(optax.clip_by_global_norm(1.0), optax.adamw(sched, b1=0.9, b2=0.95, eps=1e-8, weight_decay=0.01))

def main():
    if len(sys.argv) == 4 and sys.argv[1] == "train":
        train(sys.argv[2], sys.argv[3])
        return
    if len(sys.argv) == 5 and sys.argv[1] == "infer":
        p = load_params(sys.argv[3])
        gen(p, sys.argv[2], sys.argv[4])
        return
    raise SystemExit

def loopfiles(root):
    root = pathlib.Path(root)
    while True:
        for p in sorted(root.rglob("*")):
            try:
                if p.is_file() and p.stat().st_size > 0:
                    yield p
            except Exception:
                continue

def maskgen(key):
    key, k1, k2, k3, k4, k5 = jax.random.split(key, 6)
    u = jax.random.uniform(k1, (BATCH, SPANS), minval=1e-6, maxval=1.0)
    lens = jnp.clip(jnp.floor(jnp.exp(u * math.log(SPANMAX))), 1, SPANMAX).astype(jnp.int32)
    starts = jax.random.randint(k2, (BATCH, SPANS), 0, CTX)
    pos = jnp.arange(CTX)[None, None, :]
    sm = ((pos >= starts[..., None]) & (pos < (starts + lens)[..., None])).any(axis=1)
    up = jax.random.uniform(k3, (BATCH,), minval=1e-6, maxval=1.0)
    plen = jnp.clip(jnp.floor(jnp.exp(up * math.log(CTX))), 1, CTX).astype(jnp.int32)
    us = jax.random.uniform(k4, (BATCH,), minval=1e-6, maxval=1.0)
    slen = jnp.clip(jnp.floor(jnp.exp(us * math.log(CTX))), 1, CTX).astype(jnp.int32)
    fl = jax.random.uniform(k5, (BATCH, 2))
    pos1 = jnp.arange(CTX)[None, :]
    pm = (pos1 < plen[:, None]) & (fl[:, 0:1] < PPRE)
    qm = (pos1 >= (CTX - slen)[:, None]) & (fl[:, 1:2] < PSUF)
    m = sm | pm | qm
    m = m.at[:, 0].set(True)
    return key, m

def save_params(p, path):
    with open(path, "wb") as f:
        pickle.dump(jax.tree_util.tree_map(lambda x: np.asarray(x, dtype=np.float32), p), f, protocol=4)

def load_params(path):
    with open(path, "rb") as f:
        q = pickle.load(f)
    return jax.tree_util.tree_map(lambda x: jnp.asarray(x, dtype=jnp.float32), q)

def init(key):
    def r(shape, s=0.02):
        nonlocal key
        key, sub = jax.random.split(key)
        return jax.random.normal(sub, shape, dtype=jnp.float32) * s
    p = {
        "emb": r((VOCAB, DIM)),
        "posemb": r((CTX, DIM)),
        "holemb": r((2, DIM)),
        "gear": r((VOCAB, GEAR), 0.1),
        "gearproj": r((GEAR, DIM)),
        "depthw": jnp.zeros((LAYERS + 1,), jnp.float32),
        "lnfw": jnp.ones((DIM,), jnp.float32),
        "lnfb": jnp.zeros((DIM,), jnp.float32),
        "bytew": r((DIM, VOCAB)),
        "byteb": jnp.zeros((VOCAB,), jnp.float32),
        "bytew2": r((DIM, VOCAB)),
        "byteb2": jnp.zeros((VOCAB,), jnp.float32),
        "layers": {
            "l1w": jnp.ones((LAYERS, DIM), jnp.float32),
            "l1b": jnp.zeros((LAYERS, DIM), jnp.float32),
            "wq": r((LAYERS, DIM, DIM)),
            "wk": r((LAYERS, DIM, DIM)),
            "wv": r((LAYERS, DIM, DIM)),
            "wo": r((LAYERS, DIM, DIM)),
            "db": jnp.full((LAYERS, HEADS), -2.0, jnp.float32),
            "regb": jnp.zeros((LAYERS, HEADS, REG), jnp.float32),
            "kreg": r((LAYERS, REG, HEADS, HDIM)),
            "vreg": r((LAYERS, REG, HEADS, HDIM)),
            "l2w": jnp.ones((LAYERS, DIM), jnp.float32),
            "l2b": jnp.zeros((LAYERS, DIM), jnp.float32),
            "f1w": r((LAYERS, DIM, DIM * 4)),
            "f1b": jnp.zeros((LAYERS, DIM * 4), jnp.float32),
            "f2w": r((LAYERS, DIM * 4, DIM)),
            "f2b": jnp.zeros((LAYERS, DIM), jnp.float32),
        },
    }
    return p

def ln(x, w, b):
    xf = x.astype(jnp.float32)
    m = jnp.mean(xf, axis=-1, keepdims=True)
    v = jnp.mean((xf - m) ** 2, axis=-1, keepdims=True)
    return ((xf - m) * jax.lax.rsqrt(v + 1e-5) * w.astype(jnp.float32) + b.astype(jnp.float32)).astype(jnp.bfloat16)

def attn(x, q):
    B, T, _ = x.shape
    bf = jnp.bfloat16
    qq = (x @ q["wq"].astype(bf)).reshape(B, T, HEADS, HDIM)
    kk = (x @ q["wk"].astype(bf)).reshape(B, T, HEADS, HDIM)
    vv = (x @ q["wv"].astype(bf)).reshape(B, T, HEADS, HDIM)
    R = q["kreg"].shape[0]
    kreg = jnp.broadcast_to(q["kreg"].astype(bf)[None], (B, R, HEADS, HDIM))
    vreg = jnp.broadcast_to(q["vreg"].astype(bf)[None], (B, R, HEADS, HDIM))
    kk = jnp.concatenate([kreg, kk], axis=1)
    vv = jnp.concatenate([vreg, vv], axis=1)
    logits = jnp.einsum("bthd,bshd->bhts", qq, kk) * (1.0 / math.sqrt(HDIM))
    dist = jnp.asarray(DIST_NP, dtype=bf)
    sp = jax.nn.softplus(q["db"].astype(jnp.float32)).astype(bf)
    bias_tt = -sp[:, None, None] * dist[None]
    bias_reg = jnp.broadcast_to(q["regb"].astype(bf)[:, None, :], (HEADS, T, R))
    logits = logits + jnp.concatenate([bias_tt, bias_reg], axis=-1)[None]
    att = jax.nn.softmax(logits.astype(jnp.float32), axis=-1).astype(bf)
    out = jnp.einsum("bhts,bshd->bthd", att, vv).reshape(B, T, DIM)
    return out @ q["wo"].astype(bf)

def block(h, q):
    h = h + attn(ln(h, q["l1w"], q["l1b"]), q)
    g = ln(h, q["l2w"], q["l2b"])
    g = jax.nn.gelu(g @ q["f1w"].astype(jnp.bfloat16) + q["f1b"].astype(jnp.bfloat16))
    g = g @ q["f2w"].astype(jnp.bfloat16) + q["f2b"].astype(jnp.bfloat16)
    return h + g

def fwd(p, x, aux):
    bf = jnp.bfloat16
    z = p["emb"].astype(bf)[x]
    z = z + p["holemb"].astype(bf)[(x == HOLE).astype(jnp.int32)]
    z = z + p["posemb"].astype(bf)[None, :, :]
    z = z + p["gear"].astype(bf)[x] @ p["gearproj"].astype(bf)
    w = jax.nn.softmax(p["depthw"].astype(jnp.float32)).astype(bf)
    acc = z * w[0]
    def f(carry, xs):
        h, a = carry
        pl, wl = xs
        h2 = block(h, pl)
        return (h2, a + h2 * wl), None
    (z, acc), _ = jax.lax.scan(jax.remat(f), (z, acc), (p["layers"], w[1:]))
    z = ln(acc, p["lnfw"], p["lnfb"])
    logits = z @ p["bytew"].astype(bf) + p["byteb"].astype(bf)
    if aux:
        logits2 = z @ p["bytew2"].astype(bf) + p["byteb2"].astype(bf)
        return logits, logits2
    return logits

def loss(p, x, y, yn, w, m):
    lm, la = fwd(p, x, True)
    ce = optax.softmax_cross_entropy_with_integer_labels(lm.astype(jnp.float32), y)
    l1 = (ce * w).sum() / (w.sum() + 1e-6)
    ce2 = optax.softmax_cross_entropy_with_integer_labels(la.astype(jnp.float32), yn)
    l2 = ce2.mean()
    acc = ((jnp.argmax(lm, axis=-1) == y).astype(jnp.float32) * m).sum() / (m.astype(jnp.float32).sum() + 1e-6)
    return l1 + 0.1 * l2, (l1, l2, acc)

def train_step(p, st, ema, key, y, yn, beta):
    key, m = maskgen(key)
    x = jnp.where(m, jnp.int32(0), y)
    w = m.astype(jnp.float32) * 0.97 + 0.03
    (l, aux), g = jax.value_and_grad(loss, has_aux=True)(p, x, y, yn, w, m)
    up, st = OPT.update(g, st, p)
    p = optax.apply_updates(p, up)
    ema = jax.tree_util.tree_map(lambda e, qq: e * beta + qq * (1.0 - beta), ema, p)
    return p, st, ema, key, l, aux

train_step = jax.jit(train_step, donate_argnums=(0, 1, 2))

def infer_step(p, x):
    logits = fwd(p, x.astype(jnp.int32), False)
    logits = logits.at[:, :, HOLE].set(jnp.bfloat16(-1e9))
    return jnp.argmax(logits, axis=-1).astype(jnp.uint8)

infer_step = jax.jit(infer_step)

def draft(srcroot, outroot):
    t0 = time.time%                                                                                         ()
    outroot = pathlib.Path(outroot)
    outroot.mkdir(parents=True, exist_ok=True)
    srcroot_p = pathlib.Path(srcroot)
    count = 0
    for src in sorted(srcroot_p.rglob("*")):
        if not src.is_file():
            continue
        rel = src.relative_to(srcroot_p)
        dst = outroot / rel
        dst.parent.mkdir(parents=True, exist_ok=True)
        try:
            raw = src.read_bytes()
        except Exception:
            continue
        if not raw:
            dst.write_bytes(b"")
            count += 1
            continue
        a = np.frombuffer(raw, dtype=np.uint8).copy()
        z = a == 0
        if z.any():
            nz = a[~z]
            if nz.size > (8 << 20):
                nz = nz[:8 << 20]
            mode = int(np.bincount(nz, minlength=256).argmax()) if nz.size else 32
            a[z] = mode
        dst.write_bytes(a.tobytes())
        count += 1
    print(f"draft done files={count} elapsed={time.time()-t0:.1f}", flush=True)

def runsplit(hm):
    rs = hm & ~np.concatenate(([False], hm[:-1]))
    re_ = hm & ~np.concatenate((hm[1:], [False]))
    si = np.flatnonzero(rs)
    ei = np.flatnonzero(re_)
    lens = ei - si + 1
    hp = np.flatnonzero(hm)
    rid = np.searchsorted(si, hp, side="right") - 1
    ln = lens[rid]
    stt = si[rid]
    iir = hp - stt
    longm = ln > RUN_MIN
    return hp[longm & ((iir & 1) == 0)], hp[longm & ((iir & 1) == 1)]

def fill_windows(p, apad, n, deadline):
    ar = np.arange(CTX)
    starts = np.arange(0, n, CTX)
    tot = 0
    for c0 in range(0, starts.shape[0], CWIN):
        if time.time() > deadline:
            v = apad[:n]
            z = v == 0
            if z.any():
                v[z] = 32
            return apad, tot
        s = starts[c0:c0 + CWIN]
        X = apad[s[:, None] + ar[None, :]]
        idx = np.nonzero((X == 0).any(axis=1))[0]
        if idx.size == 0:
            continue
        Xs = X[idx]
        m = Xs.shape[0]
        N = ((m + INFER_BS - 1) // INFER_BS) * INFER_BS
        Xp = np.zeros((N, CTX), dtype=np.uint8)
        Xp[:m] = Xs
        outs = []
        for o in range(0, N, INFER_BS):
            pr = infer_step(p, jnp.asarray(Xp[o:o + INFER_BS]))
            outs.append(np.asarray(pr))
        P = np.concatenate(outs, axis=0)[:m]
        zm = Xs == 0
        Xs[zm] = P[zm]
        X[idx] = Xs
        apad[s[0]:s[0] + s.shape[0] * CTX] = X.reshape(-1)
        tot += m
    return apad, tot

def gen(p, srcroot, outroot):
    t0 = time.time()
    deadline = t0 + GEN_CAP
    p = jax.tree_util.tree_map(lambda a: a.astype(jnp.bfloat16) if a.dtype == jnp.float32 else a, p)
    outroot = pathlib.Path(outroot)
    outroot.mkdir(parents=True, exist_ok=True)
    srcroot_p = pathlib.Path(srcroot)
    files_done = 0
    totwin = 0
    print("gen start", flush=True)
    for src in sorted(srcroot_p.rglob("*")):
        if not src.is_file():
            continue
        rel = src.relative_to(srcroot_p)
        dst = outroot / rel
        dst.parent.mkdir(parents=True, exist_ok=True)
        try:
            raw = src.read_bytes()
        except Exception:
            continue
        n = len(raw)
        if n == 0:
            dst.write_bytes(b"")
            files_done += 1
            continue
        apad = np.empty(n + CTX, dtype=np.uint8)
        apad[:n] = np.frombuffer(raw, dtype=np.uint8)
        apad[n:] = HOLE
        hm = apad[:n] == 0
        if not hm.any():
            dst.write_bytes(raw)
            files_done += 1
            continue
        if n <= (1 << 28) and hm.mean() <= 0.5:
            subA, subB = runsplit(hm)
        else:
            subA = subB = np.empty(0, np.int64)
        apad, w1 = fill_windows(p, apad, n, deadline)
        totwin += w1
        if time.time() < t0 + REFINE_OK and (subA.size or subB.size):
            if subA.size:
                base = apad.copy()
                base[subA] = HOLE
                apad, w2 = fill_windows(p, base, n, deadline)
                totwin += w2
            if subB.size:
                base = apad.copy()
                base[subB] = HOLE
                apad, w3 = fill_windows(p, base, n, deadline)
                totwin += w3
        dst.write_bytes(apad[:n].tobytes())
        del apad
        files_done += 1
        if files_done % 1000 == 0:
            print(f"gen files={files_done} win={totwin} elapsed={time.time()-t0:.1f}", flush=True)
        if time.time() > deadline:
            print("gen cap hit, remaining files keep previous predictions", flush=True)
            break
    print(f"gen done files={files_done} win={totwin} elapsed={time.time()-t0:.1f}", flush=True)

def train(root, out):
    print(f"start device={jax.devices()[0]} ctx={CTX} batch={BATCH} dim={DIM} layers={LAYERS}", flush=True)
    out = pathlib.Path(out)
    out.mkdir(parents=True, exist_ok=True)
    (out / "newvalpred").mkdir(parents=True, exist_ok=True)
    print("drafting newvalpred immediately", flush=True)
    draft("/newvalin", out / "newvalpred")
    print("draft validation complete", flush=True)
    key = jax.random.PRNGKey(0)
    p = init(key)
    nparams = sum(int(x.size) for x in jax.tree_util.tree_leaves(p))
    print(f"params={nparams} approx_mb={nparams * 4 / 1e6:.0f}", flush=True)
    ema = jax.tree_util.tree_map(lambda a: a.copy(), p)
    st = OPT.init(p)
    rng = np.random.default_rng(7)
    ar = np.arange(CTX)
    step = 0
    fcount = 0
    t0 = time.time()
    next_gen = FIRST_GEN
    for fp in loopfiles(root):
        try:
            fsize = fp.stat().st_size
        except Exception:
            continue
        if fsize < CTX + 1:
            continue
        try:
            if fsize > LIMIT:
                off = int(rng.integers(0, fsize - LIMIT + 1))
                with open(fp, "rb") as fh:
                    fh.seek(off)
                    raw = fh.read(LIMIT)
            else:
                raw = fp.read_bytes()
        except Exception:
            continue
        n = len(raw)
        if n < CTX + 1:
            continue
        data = np.frombuffer(raw, dtype=np.uint8)
        fcount += 1
        if fcount <= 10 or fcount % 200 == 0:
            print(f"file={fcount} bytes={n} name={fp.name} elapsed={time.time()-t0:.1f}", flush=True)
        for _ in range(PER):
            s = rng.integers(0, n - CTX, (BATCH, 1))
            idx = s + ar[None, :]
            y = jnp.asarray(data[idx].astype(np.int32))
            yn = jnp.asarray(data[idx + 1].astype(np.int32))
            beta = min(0.999, (step + 1.0) / (step + 10.0))
            p, st, ema, key, l, aux = train_step(p, st, ema, key, y, yn, beta)
            step += 1
            if step <= 5 or step % 100 == 0:
                print(f"step={step} loss={float(l):.4f} hole={float(aux[0]):.4f} aux={float(aux[1]):.4f} acc={float(aux[2]):.4f} elapsed={time.time()-t0:.1f}", flush=True)
            if step % 500 == 0:
                save_params(ema, out / "model.pkl")
            el = time.time() - t0
            if el >= next_gen:
                save_params(ema, out / "model.pkl")
                gen(ema, "/newvalin", out / "newvalpred")
                el = time.time() - t0
                gap = LATE_GAP if el > LATE_GEN else max(GEN_MIN_GAP, GEN_FRAC * el)
                next_gen = el + gap
    save_params(ema, out / "model.pkl")

if __name__ == "__main__":