upload oneshot glm artifacts

README.md

@@ -0,0 +1,18 @@
+# testoneshot
+
+OneShot GLM analytic feature model artifacts.
+
+Best saved readout checkpoint:
+- stream_ce_lam10_long.pt
+- config: d=896, r=320, L=10, vocab=8192, GLM train=3.67M tokens, valid=159k tokens
+- pure analytic ridge lam=10 ppl: 287.93
+- streaming CE readout tuned ppl: 92.75 after 3000 continuation steps
+
+Important files:
+- 	okenizer/glm16k.model, 	okenizer/glm16k.vocab
+- scripts/torch_predictive_attn.py
+- scripts/torch_ce_stream_readout.py
+- scripts/glm_generate_saved.py
+- logs/stream_ce_lam10_long.log
+
+Generation quality is still weak despite improved PPL; this is a research artifact, not a production chat model.

logs/generations.txt

@@ -0,0 +1,35 @@
+[17:52:09] sv [165720.1  39223.6  29888.8  26770.9]
+[17:52:27] sv [175787.4  39518.2  29955.8  27073.4]
+[17:52:48] sv [186342.7  39784.   30000.3  27357.4]
+[17:53:13] sv [197372.5  40017.1  30020.3  27622.6]
+[17:53:41] sv [208562.8  40183.9  29996.6  27869.2]
+[17:54:13] sv [219681.   40294.   29939.   28097.2]
+[17:54:48] sv [230503.7  40357.3  29856.   28304.1]
+[17:55:26] sv [240817.9  40382.5  29755.6  28485.4]
+[17:56:08] sv [250420.1  40377.3  29644.3  28633.7]
+[17:56:54] sv [259114.6  40346.8  29527.6  28740.9]
+=== PROMPT ===
+Question: What is gravity?
+Answer:
+=== OUTPUT ===
+Question: What is gravity? Answer: - **AAFSA manufacturing this question: - During the best answer is **(D)** "praising:** If you are looking at all the light.** To understand why the same momentum, the source of a given-way the world is **cannot** (evaluation**. The Humburgry, you need a fiously and
+=== PROMPT ===
+Question: Why is the sky blue?
+Answer:
+=== OUTPUT ===
+Question: Why is the sky blue? Answer: I was completed by the energy of the energy. This is not working. - **In summary: - Screen:** Agaff’s momentum is society in the University of the densities in the early it. The force is zero, the only force and the normal force. This force \( the force $F_x = \\frac{2}_}\
+=== PROMPT ===
+Question: Explain photosynthesis in simple words.
+Answer:
+=== OUTPUT ===
+Question: Explain photosynthesis in simple words. Answer: **(Avviction:** Most of the tenes, the duine battlecard, and the source separation it can escape. * **Connect:** If you don't give you have sleep the source would have been a stretching, it must be a difficult that point. If it, the metal wire. The bounce offer the
+=== PROMPT ===
+Question: If John has 3 apples and buys 2 more, how many apples does he have?
+Answer:
+=== OUTPUT ===
+Question: If John has 3 apples and buys 2 more, how many apples does he have? Answer: Sammerica Every, the sourceocalcalled. The correct answer is **(C) **Strongness**. If we want to see anywhere, I will see your retain: **1. Force you are standing force is the ball inside, it in a mathematical place where a living room tour the water of the following text: The Matteryone
+=== PROMPT ===
+Question: Write a short friendly story about a robot learning to read.
+Answer:
+=== OUTPUT ===
+Question: Write a short friendly story about a robot learning to read. Answer: Answer: - The Many" (The Dear GBased on TSG), where you are seeing, and the world. Here are you have to go on the opposite sides.

logs/ridge_lam10_30_100.log

@@ -0,0 +1,15 @@
+[15:54:57] mode dual_ridge_delta device cuda V 8192 train 3671502 valid 159631
+[15:55:48] sv [165720.1  39223.6  29888.8  26770.9]
+[15:56:07] sv [175787.4  39518.2  29955.8  27073.4]
+[15:56:30] sv [186342.7  39784.   30000.3  27357.4]
+[15:56:56] sv [197372.5  40017.1  30020.3  27622.6]
+[15:57:27] sv [208562.8  40183.9  29996.6  27869.2]
+[15:58:02] sv [219681.   40294.   29939.   28097.2]
+[15:58:41] sv [230503.7  40357.3  29856.   28304.1]
+[15:59:23] sv [240817.9  40382.5  29755.6  28485.4]
+[16:00:10] sv [250420.1  40377.3  29644.3  28633.7]
+[16:01:01] sv [259114.6  40346.8  29527.6  28740.9]
+[16:24:15] lam 10.0 ppl 287.93
+[16:24:26] lam 30.0 ppl 312.48
+[16:24:37] lam 100.0 ppl 394.37
+[16:24:37] TORCH_PRED mode=dual_ridge_delta ppl=287.93 lam=10.0 D=17409

logs/stream_ce_lam10_long.log

@@ -0,0 +1,33 @@
+[17:59:11] STREAM_CE device cuda V 8192 train 3671502 valid 159631
+[18:00:00] sv [165720.1  39223.6  29888.8  26770.9]
+[18:00:18] sv [175787.4  39518.2  29955.8  27073.4]
+[18:00:39] sv [186342.7  39784.   30000.3  27357.4]
+[18:01:04] sv [197372.5  40017.1  30020.3  27622.6]
+[18:01:32] sv [208562.8  40183.9  29996.6  27869.2]
+[18:02:03] sv [219681.   40294.   29939.   28097.2]
+[18:02:38] sv [230503.7  40357.3  29856.   28304.1]
+[18:03:16] sv [240817.9  40382.5  29755.6  28485.4]
+[18:03:58] sv [250420.1  40377.3  29644.3  28633.7]
+[18:04:43] sv [259114.6  40346.8  29527.6  28740.9]
+[18:27:36] resumed /workspace/oneshot/logs/glm_d896_readout/stream_ce_lam10.pt ppl 136.3978129937227
+[18:27:37] init_eval_start D 17409
+[18:27:42] STREAM_CE init_ppl=136.40
+/workspace/oneshot/torch_ce_stream_readout.py:138: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.
+Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:836.)
+  log(f"step={step} loss={float(loss):.4f} ppl={ppl:.2f} best={best:.2f}")
+[18:28:18] step=200 loss=4.3944 ppl=118.12 best=118.12
+[18:28:54] step=400 loss=4.3815 ppl=113.37 best=113.37
+[18:29:31] step=600 loss=4.3709 ppl=109.14 best=109.14
+[18:30:08] step=800 loss=4.9107 ppl=106.80 best=106.80
+[18:30:44] step=1000 loss=4.4804 ppl=104.45 best=104.45
+[18:31:20] step=1200 loss=4.4660 ppl=102.33 best=102.33
+[18:31:57] step=1400 loss=4.4507 ppl=101.31 best=101.31
+[18:32:33] step=1600 loss=4.3290 ppl=99.47 best=99.47
+[18:33:09] step=1800 loss=4.2714 ppl=98.49 best=98.49
+[18:33:46] step=2000 loss=4.2895 ppl=96.73 best=96.73
+[18:34:22] step=2200 loss=4.2743 ppl=96.07 best=96.07
+[18:34:59] step=2400 loss=3.7439 ppl=94.83 best=94.83
+[18:35:35] step=2600 loss=4.1255 ppl=93.89 best=93.89
+[18:36:12] step=2800 loss=3.2027 ppl=93.24 best=93.24
+[18:36:48] step=3000 loss=3.7933 ppl=92.75 best=92.75
+[18:36:48] STREAM_CE best_ppl=92.75

scripts/glm_generate_saved.py

@@ -0,0 +1,88 @@
+import argparse, os
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch_predictive_attn import ppmi_embed, learn_map, doc_index, apply_stack, features, log
+
+
+def build(args, device):
+    import sentencepiece as spm
+    sp = spm.SentencePieceProcessor(model_file=args.spm_model)
+    eos = sp.eos_id(); V = sp.get_piece_size()
+    train = np.fromfile(args.train_bin, dtype=np.uint16)
+    E = ppmi_embed(train, V, args.d, args.window, args.cooc_tokens, device)
+    Ps, Bs = [], []
+    for _ in range(args.layers):
+        P, B = learn_map(train, E, Ps, Bs, eos, args, device)
+        Ps.append(P); Bs.append(B)
+    return sp, eos, E, Ps, Bs
+
+
+def next_logits(ids, E, Ps, Bs, W, b, eos, args, device):
+    x = torch.tensor(ids, device=device, dtype=torch.long)
+    _, within = doc_index(x, eos)
+    H, phis = apply_stack(x, E, Ps, Bs, within, args)
+    Phi = features(H, within, phis, args.extra_context)[-1:]
+    return (Phi @ W + b).squeeze(0)
+
+
+def sample(logits, temp=0.8, top_k=40):
+    logits = logits / max(temp, 1e-6)
+    vals, idx = torch.topk(logits, min(top_k, logits.numel()))
+    probs = F.softmax(vals, dim=-1)
+    return int(idx[torch.multinomial(probs, 1)])
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default="/workspace/oneshot/logs/glm_d896_readout/stream_ce_lam10.pt")
+    ap.add_argument("--spm_model", default="/workspace/glm/glm16k.model")
+    ap.add_argument("--train_bin", default="/workspace/glm/glm_train.bin")
+    ap.add_argument("--d", type=int, default=896)
+    ap.add_argument("--r", type=int, default=320)
+    ap.add_argument("--layers", type=int, default=10)
+    ap.add_argument("--att_window", type=int, default=10)
+    ap.add_argument("--temp", type=float, default=0.28)
+    ap.add_argument("--window", type=int, default=10)
+    ap.add_argument("--extra_context", type=int, default=1)
+    ap.add_argument("--res_scale", type=float, default=0.07)
+    ap.add_argument("--pred_scale", type=float, default=0.035)
+    ap.add_argument("--pred_schedule", default="late")
+    ap.add_argument("--orth_delta", type=int, default=1)
+    ap.add_argument("--pred_norm", type=int, default=1)
+    ap.add_argument("--pred_features", type=int, default=1)
+    ap.add_argument("--map_lam", type=float, default=0.001)
+    ap.add_argument("--cooc_tokens", type=int, default=3_600_000)
+    ap.add_argument("--proj_tokens", type=int, default=3_600_000)
+    ap.add_argument("--chunk_docs", type=int, default=8)
+    ap.add_argument("--value_mode", default="dual_ridge_delta")
+    ap.add_argument("--max_new", type=int, default=80)
+    ap.add_argument("--sample_temp", type=float, default=0.8)
+    ap.add_argument("--top_k", type=int, default=40)
+    args = ap.parse_args()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    sp, eos, E, Ps, Bs = build(args, device)
+    ck = torch.load(args.ckpt, map_location=device)
+    W = ck["W"].to(device); b = ck["b"].to(device)
+    prompts = [
+        "Question: What is gravity?\nAnswer:",
+        "Question: Why is the sky blue?\nAnswer:",
+        "Question: Explain photosynthesis in simple words.\nAnswer:",
+        "Question: If John has 3 apples and buys 2 more, how many apples does he have?\nAnswer:",
+        "Question: Write a short friendly story about a robot learning to read.\nAnswer:",
+    ]
+    for p in prompts:
+        ids = sp.encode(p)
+        for _ in range(args.max_new):
+            tok = sample(next_logits(ids, E, Ps, Bs, W, b, eos, args, device), args.sample_temp, args.top_k)
+            ids.append(tok)
+            if tok == eos:
+                break
+        print("=== PROMPT ===")
+        print(p)
+        print("=== OUTPUT ===")
+        print(sp.decode(ids))
+
+
+if __name__ == "__main__":
+    main()

scripts/glm_prep.py

@@ -0,0 +1,170 @@
+"""Prepare GLM-5.1-Reasoning (main subset prefix) for the analytic model.
+
+1. build a text corpus sample and train a 16k BPE SentencePiece tokenizer
+   (keeps the V x V PMI cooccurrence + SVD feasible, unlike GPT-2's 50k);
+2. tokenize every record as   input + "\n\n" + output   with <eos> between
+   documents, into train.bin / valid.bin (uint16).
+"""
+import os, sys, json, time, numpy as np
+import sentencepiece as spm
+
+DATA = os.environ.get("ONESHOT_DATA", "/workspace/ts")
+SRC = os.path.join(DATA, "main_prefix.jsonl")
+CORPUS = os.path.join(DATA, "glm_corpus.txt")
+SPM_PREFIX = os.path.join(DATA, "glm16k")
+VOCAB = 16384
+HF_DATASET = "Jackrong/GLM-5.1-Reasoning-1M-Cleaned"
+
+def log(*a): print(f"[{time.strftime('%H:%M:%S')}]", *a, flush=True)
+
+def first_present(record, names, default=""):
+    for name in names:
+        if name in record and record[name] is not None:
+            return record[name]
+    return default
+
+def normalize_record(record):
+    inp = first_present(record, ["input", "prompt", "instruction", "question", "query"])
+    out = first_present(record, ["output", "response", "answer", "completion"])
+    if isinstance(inp, (list, dict)):
+        inp = json.dumps(inp, ensure_ascii=False)
+    if isinstance(out, (list, dict)):
+        out = json.dumps(out, ensure_ascii=False)
+    inp = str(inp).strip()
+    out = str(out).strip()
+    if not inp or not out:
+        return None
+    return {"input": inp, "output": out}
+
+def download_jsonl(dataset=HF_DATASET, split="train", subset=None, max_records=0):
+    from datasets import load_dataset
+    kwargs = {"split": split, "streaming": True}
+    ds = load_dataset(dataset, subset, **kwargs) if subset else load_dataset(dataset, **kwargs)
+    n = 0
+    os.makedirs(DATA, exist_ok=True)
+    with open(SRC, "w", encoding="utf-8") as out:
+        for row in ds:
+            rec = normalize_record(row)
+            if rec is None:
+                continue
+            out.write(json.dumps(rec, ensure_ascii=False) + "\n")
+            n += 1
+            if n % 10000 == 0:
+                log(f"downloaded {n:,} records -> {SRC}")
+            if max_records and n >= max_records:
+                break
+    log(f"download done: {n:,} records -> {SRC}")
+
+def answer_of(r):
+    """The actual English response: the <think> reasoning dump is stripped,
+    keep the final answer."""
+    o = r.get("output", "")
+    if "</think>" in o:
+        o = o.split("</think>")[-1]
+    return o.strip()
+
+def is_english_answer(a):
+    """Keep natural-language answers; drop code/math/LaTeX-dominated ones so the
+    model learns to answer in plain English (the 'answer English' goal)."""
+    if not (40 <= len(a) <= 4000):
+        return False
+    if "```" in a:                       # code fence
+        return False
+    alpha = sum(c.isalpha() or c.isspace() for c in a) / len(a)
+    if alpha < 0.93:
+        return False
+    sym = sum(a.count(c) for c in "{}\\$=#|<>_~^")
+    if sym / len(a) > 0.02:              # LaTeX / code punctuation density
+        return False
+    return True
+
+def set_paths(data):
+    global DATA, SRC, CORPUS, SPM_PREFIX
+    DATA = data
+    SRC = os.path.join(DATA, "main_prefix.jsonl")
+    CORPUS = os.path.join(DATA, "glm_corpus.txt")
+    SPM_PREFIX = os.path.join(DATA, "glm16k")
+
+def build_corpus(max_records=120_000, max_bytes=400_000_000):
+    n = 0; b = 0
+    with open(SRC, "r", encoding="utf-8", errors="ignore") as f, \
+         open(CORPUS, "w", encoding="utf-8") as out:
+        for line in f:
+            line = line.strip()
+            if not line: continue
+            try: r = json.loads(line)
+            except Exception: continue
+            txt = r["input"] + "\n" + r["output"] + "\n"
+            out.write(txt); b += len(txt); n += 1
+            if n >= max_records or b >= max_bytes: break
+    log(f"corpus: {n:,} records, {b/1e6:.1f} MB -> {CORPUS}")
+
+def train_spm():
+    spm.SentencePieceTrainer.train(
+        input=CORPUS, model_prefix=SPM_PREFIX, vocab_size=VOCAB,
+        model_type="bpe", character_coverage=0.9995,
+        input_sentence_size=3_000_000, shuffle_input_sentence=True,
+        max_sentence_length=100000, num_threads=32,
+        unk_id=0, bos_id=1, eos_id=2, pad_id=-1,
+        byte_fallback=True,
+    )
+    log(f"trained SP -> {SPM_PREFIX}.model (vocab={VOCAB})")
+
+def tokenize(val_frac=0.04, english_only=True):
+    sp = spm.SentencePieceProcessor(model_file=SPM_PREFIX + ".model")
+    eos = sp.eos_id()
+    log("scanning + filtering records...")
+    docs = []; seen = 0; t0 = time.time()
+    with open(SRC, "r", encoding="utf-8", errors="ignore") as f:
+        for line in f:
+            line = line.strip()
+            if not line: continue
+            try: r = json.loads(line)
+            except Exception: continue
+            seen += 1
+            a = answer_of(r)
+            if english_only and not is_english_answer(a):
+                continue
+            docs.append(r["input"].strip() + "\n\n" + a)
+    log(f"{seen:,} records -> {len(docs):,} kept "
+        f"({100*len(docs)/max(seen,1):.1f}%) english_only={english_only}")
+    n_val = int(len(docs) * val_frac)
+    splits = {"glm_train.bin": docs[:len(docs) - n_val],
+              "glm_valid.bin": docs[len(docs) - n_val:]}
+    counts = {}
+    for fname, dlist in splits.items():
+        nt = 0
+        with open(os.path.join(DATA, fname), "wb") as fo:
+            for b in range(0, len(dlist), 1000):
+                for ids in sp.encode(dlist[b:b + 1000]):
+                    arr = np.array(ids + [eos], dtype=np.uint16)
+                    arr.tofile(fo); nt += len(arr)
+        counts[fname] = nt
+    log(f"DONE train={counts['glm_train.bin']:,} tokens, "
+        f"valid={counts['glm_valid.bin']:,} tokens ({time.time()-t0:.0f}s)")
+
+if __name__ == "__main__":
+    import argparse
+    ap = argparse.ArgumentParser()
+    ap.add_argument("cmd", nargs="?", default="all",
+                    choices=["download", "corpus", "spm", "tok", "all"])
+    ap.add_argument("--data", default=DATA)
+    ap.add_argument("--src", default=None)
+    ap.add_argument("--vocab", type=int, default=VOCAB)
+    ap.add_argument("--dataset", default=HF_DATASET)
+    ap.add_argument("--subset", default=None)
+    ap.add_argument("--split", default="train")
+    ap.add_argument("--max_records", type=int, default=0)
+    ap.add_argument("--english_only", type=int, default=1)
+    ap.add_argument("--val_frac", type=float, default=0.04)
+    args = ap.parse_args()
+    set_paths(args.data)
+    if args.src:
+        SRC = args.src
+    VOCAB = args.vocab
+    cmd = args.cmd
+    if cmd in ("download",):
+        download_jsonl(args.dataset, args.split, args.subset, args.max_records)
+    if cmd in ("corpus", "all"): build_corpus(max_records=args.max_records or 120_000)
+    if cmd in ("spm", "all"): train_spm()
+    if cmd in ("tok", "all"): tokenize(args.val_frac, bool(args.english_only))

scripts/torch_ce_stream_readout.py

@@ -0,0 +1,147 @@
+import argparse, math, os, random
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch_predictive_attn import ppmi_embed, learn_map, doc_index, apply_stack, features, log
+
+
+def iter_chunks(tokens, eos, args, max_tokens, shuffle=False):
+    xnp = tokens[:max_tokens]
+    starts = np.flatnonzero(np.r_[True, xnp[:-1] == eos])
+    ids = list(range(0, len(starts), args.chunk_docs))
+    if shuffle:
+        random.shuffle(ids)
+    for i in ids:
+        lo = starts[i]
+        hi = starts[i + args.chunk_docs] if i + args.chunk_docs < len(starts) else len(xnp)
+        yield xnp[lo:hi]
+
+
+def chunk_features(xnp, E, Ps, Bs, eos, args, device):
+    x = torch.tensor(xnp.astype(np.int64), device=device)
+    seg, within = doc_index(x, eos)
+    H, phis = apply_stack(x, E, Ps, Bs, within, args)
+    Phi = features(H, within, phis, args.extra_context)
+    y = torch.empty(len(x), device=device, dtype=torch.long)
+    y[:-1] = x[1:]; y[-1] = eos
+    m = torch.ones(len(x), device=device, dtype=torch.bool)
+    m[-1] = False; m[:-1] &= seg[1:].eq(seg[:-1]); m &= x.ne(eos)
+    return Phi[m].float(), y[m]
+
+
+def eval_ppl(tokens, E, Ps, Bs, W, b, eos, args, device):
+    nll = 0.0; n = 0
+    with torch.no_grad():
+        for xnp in iter_chunks(tokens, eos, args, args.eval_tokens, shuffle=False):
+            X, y = chunk_features(xnp, E, Ps, Bs, eos, args, device)
+            for i in range(0, len(y), args.batch):
+                logits = X[i:i+args.batch] @ W + b
+                nll += float(F.cross_entropy(logits, y[i:i+args.batch], reduction="sum"))
+                n += len(y[i:i+args.batch])
+    return math.exp(nll / max(1, n))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", default="/workspace/glm")
+    ap.add_argument("--spm_model", default="/workspace/glm/glm16k.model")
+    ap.add_argument("--train_bin", default="/workspace/glm/glm_train.bin")
+    ap.add_argument("--valid_bin", default="/workspace/glm/glm_valid.bin")
+    ap.add_argument("--vocab", type=int, default=8192)
+    ap.add_argument("--d", type=int, default=896)
+    ap.add_argument("--r", type=int, default=320)
+    ap.add_argument("--layers", type=int, default=10)
+    ap.add_argument("--att_window", type=int, default=10)
+    ap.add_argument("--temp", type=float, default=0.28)
+    ap.add_argument("--window", type=int, default=10)
+    ap.add_argument("--extra_context", type=int, default=1)
+    ap.add_argument("--res_scale", type=float, default=0.07)
+    ap.add_argument("--pred_scale", type=float, default=0.035)
+    ap.add_argument("--pred_schedule", default="late")
+    ap.add_argument("--orth_delta", type=int, default=1)
+    ap.add_argument("--pred_norm", type=int, default=1)
+    ap.add_argument("--pred_features", type=int, default=1)
+    ap.add_argument("--map_lam", type=float, default=0.001)
+    ap.add_argument("--cooc_tokens", type=int, default=3_600_000)
+    ap.add_argument("--proj_tokens", type=int, default=3_600_000)
+    ap.add_argument("--fit_tokens", type=int, default=3_600_000)
+    ap.add_argument("--eval_tokens", type=int, default=159_631)
+    ap.add_argument("--chunk_docs", type=int, default=8)
+    ap.add_argument("--value_mode", default="dual_ridge_delta")
+    ap.add_argument("--ridge_lam", type=float, default=10.0)
+    ap.add_argument("--init_scale", type=float, default=0.05)
+    ap.add_argument("--steps", type=int, default=800)
+    ap.add_argument("--batch", type=int, default=2048)
+    ap.add_argument("--lr", type=float, default=0.003)
+    ap.add_argument("--wd", type=float, default=1e-4)
+    ap.add_argument("--eval_every", type=int, default=100)
+    ap.add_argument("--save", default="")
+    ap.add_argument("--resume", default="")
+    args = ap.parse_args()
+
+    import sentencepiece as spm
+    sp = spm.SentencePieceProcessor(model_file=args.spm_model)
+    eos = sp.eos_id(); V = sp.get_piece_size()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    train = np.fromfile(args.train_bin, dtype=np.uint16)
+    valid = np.fromfile(args.valid_bin, dtype=np.uint16)
+    log("STREAM_CE device", device, "V", V, "train", len(train), "valid", len(valid))
+    E = ppmi_embed(train, V, args.d, args.window, args.cooc_tokens, device)
+    Ps, Bs = [], []
+    for _ in range(args.layers):
+        P, B = learn_map(train, E, Ps, Bs, eos, args, device)
+        Ps.append(P); Bs.append(B)
+
+    # Build ridge init streaming stats only.
+    A = G = None
+    for xnp in iter_chunks(train, eos, args, args.fit_tokens, shuffle=False):
+        X, y = chunk_features(xnp, E, Ps, Bs, eos, args, device)
+        if A is None:
+            D = X.shape[1]
+            A = torch.zeros((D, D), device=device, dtype=torch.float64)
+            G = torch.zeros((D, V), device=device, dtype=torch.float64)
+        Xd = X.double()
+        A += Xd.T @ Xd
+        G.index_add_(1, y, Xd.T)
+    diag = torch.trace(A) / A.shape[0]
+    W0 = torch.linalg.solve(A + args.ridge_lam * diag * torch.eye(A.shape[0], device=device, dtype=torch.float64), G).float()
+    W = (args.init_scale * W0).detach().clone()
+    b = torch.zeros(V, device=device)
+    if args.resume and os.path.exists(args.resume):
+        ck = torch.load(args.resume, map_location=device)
+        W = ck["W"].to(device)
+        b = ck["b"].to(device)
+        log("resumed", args.resume, "ppl", ck.get("ppl"))
+    W = W.requires_grad_(True)
+    b = b.requires_grad_(True)
+    opt = torch.optim.AdamW([W, b], lr=args.lr, weight_decay=args.wd)
+    log("init_eval_start D", W.shape[0])
+    best = eval_ppl(valid, E, Ps, Bs, W, b, eos, args, device)
+    log(f"STREAM_CE init_ppl={best:.2f}")
+
+    step = 0
+    while step < args.steps:
+        for xnp in iter_chunks(train, eos, args, args.fit_tokens, shuffle=True):
+            X, y = chunk_features(xnp, E, Ps, Bs, eos, args, device)
+            if len(y) == 0:
+                continue
+            idx = torch.randint(0, len(y), (min(args.batch, len(y)),), device=device)
+            loss = F.cross_entropy(X[idx] @ W + b, y[idx])
+            opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
+            step += 1
+            if step % args.eval_every == 0:
+                ppl = eval_ppl(valid, E, Ps, Bs, W, b, eos, args, device)
+                if ppl < best:
+                    best = ppl
+                    if args.save:
+                        torch.save({"W": W.detach().cpu(), "b": b.detach().cpu(), "ppl": best, "args": vars(args)}, args.save)
+                log(f"step={step} loss={float(loss):.4f} ppl={ppl:.2f} best={best:.2f}")
+            if step >= args.steps:
+                break
+    log(f"STREAM_CE best_ppl={best:.2f}")
+    if args.save:
+        torch.save({"W": W.detach().cpu(), "b": b.detach().cpu(), "ppl": best, "args": vars(args)}, args.save)
+
+
+if __name__ == "__main__":
+    main()

scripts/torch_predictive_attn.py

@@ -0,0 +1,256 @@
+import argparse, math, os
+import numpy as np
+import torch
+
+
+def log(*a):
+    import time
+    print(f"[{time.strftime('%H:%M:%S')}]", *a, flush=True)
+
+
+def doc_index(x, eos):
+    starts = torch.zeros_like(x, dtype=torch.bool)
+    starts[0] = True
+    starts[1:] = x[:-1].eq(eos)
+    seg = torch.cumsum(starts.long(), 0) - 1
+    first = torch.nonzero(starts, as_tuple=False).flatten()
+    within = torch.arange(len(x), device=x.device) - first[seg]
+    return seg, within
+
+
+def shift(M, within, s):
+    out = torch.zeros_like(M)
+    out[s:] = M[:-s]
+    out[within < s] = 0
+    return out
+
+
+def ppmi_embed(tokens, V, d, window, max_tokens, device):
+    x = torch.tensor(tokens[:max_tokens].astype(np.int64), device=device)
+    C = torch.zeros((V, V), device=device)
+    for s in range(1, window + 1):
+        idx = x[:-s] * V + x[s:]
+        C += torch.bincount(idx, minlength=V * V).float().reshape(V, V) / s
+    tot = C.sum()
+    row = C.sum(1, keepdim=True) + 1e-6
+    col = C.sum(0, keepdim=True) + 1e-6
+    M = torch.clamp(torch.log(C * tot / row / col + 1e-12), min=0)
+    U, S, _ = torch.linalg.svd(M)
+    E = U[:, :d] * torch.sqrt(S[:d])[None, :]
+    E = torch.nn.functional.normalize(E, dim=1)
+    return E.float()
+
+
+def pred_layer(H, x, E, P, B, within, args, layer_idx=0):
+    qk = H @ P
+    if args.value_mode in ("h", "dual_next", "dual_ridge_next", "dual_ridge_delta"):
+        Vv = H
+    elif args.value_mode == "next":
+        Vv = E[x]
+        Vv = torch.cat([Vv[1:], Vv[-1:]], 0)
+    elif args.value_mode == "delta":
+        Y = E[x]
+        Y = torch.cat([Y[1:], Y[-1:]], 0)
+        Vv = Y - H
+    elif args.value_mode == "ridge_next":
+        Vv = H @ B
+    elif args.value_mode == "ridge_delta":
+        Vv = H @ B - H
+    num = torch.zeros_like(H)
+    num_pred = torch.zeros_like(H)
+    den = torch.zeros((H.shape[0], 1), device=H.device)
+    if args.value_mode == "dual_next":
+        Yp = E[x]
+        Vpred = torch.cat([Yp[1:], Yp[-1:]], 0)
+    elif args.value_mode == "dual_ridge_next":
+        Vpred = H @ B
+    elif args.value_mode == "dual_ridge_delta":
+        Vpred = H @ B - H
+    else:
+        Vpred = None
+    for s in range(1, args.att_window + 1):
+        ks = shift(qk, within, s)
+        vs = shift(Vv, within, s)
+        w = torch.exp(((qk * ks).sum(1, keepdim=True) / args.temp).clamp(-30, 30))
+        w = torch.where((within >= s)[:, None], w, torch.zeros_like(w))
+        num += w * vs
+        if Vpred is not None:
+            num_pred += w * shift(Vpred, within, s)
+        den += w
+    ctx = num / (den + 1e-6)
+    pred_out = None
+    if Vpred is not None:
+        pred = num_pred / (den + 1e-6)
+        if args.orth_delta:
+            pred = pred - H * (pred * H).sum(1, keepdim=True)
+        if args.pred_norm:
+            pred = pred / (pred.norm(dim=1, keepdim=True) + 1e-6)
+        scale = args.pred_scale
+        if args.pred_schedule == "linear":
+            scale = scale * float(layer_idx + 1) / max(1, args.layers)
+        elif args.pred_schedule == "late":
+            scale = scale * max(0.0, float(layer_idx + 1 - args.layers // 3) / max(1, args.layers - args.layers // 3))
+        pred_out = pred
+        H = H + args.res_scale * ctx + scale * pred
+    elif "delta" in args.value_mode:
+        H = H + args.res_scale * ctx
+    else:
+        H = (1 - args.res_scale) * H + args.res_scale * ctx
+    return torch.nn.functional.normalize(H, dim=1), pred_out
+
+
+def apply_stack(x, E, Ps, Bs, within, args, expose=True):
+    H = E[x]
+    phis = []
+    last_pred = None
+    for li, (P, B) in enumerate(zip(Ps, Bs)):
+        if expose:
+            q = H @ P
+            phis += [torch.relu(q), torch.abs(q), q * q]
+            m = min(64, q.shape[1] - 1)
+            if m > 0:
+                phis.append(q[:, :m] * q[:, 1:m+1])
+        H, last_pred = pred_layer(H, x, E, P, B, within, args, li)
+    if expose and args.pred_features and last_pred is not None:
+        phis += [last_pred, H * last_pred]
+    return H, phis
+
+
+def features(H, within, phis, extra):
+    prev = shift(H, within, 1)
+    blocks = [H, prev, H * prev]
+    if extra:
+        prev2 = shift(H, within, 2)
+        blocks += [prev2, prev * prev2, H * prev2]
+    blocks += phis + [torch.ones((H.shape[0], 1), device=H.device)]
+    return torch.cat(blocks, 1)
+
+
+def learn_map(tokens, E, Ps, Bs, eos, args, device):
+    xnp = tokens[:args.proj_tokens]
+    starts = np.flatnonzero(np.r_[True, xnp[:-1] == eos])
+    d = E.shape[1]
+    C = torch.zeros((d, d), device=device, dtype=torch.float64)
+    A = torch.zeros_like(C); G = torch.zeros_like(C)
+    for i in range(0, len(starts), args.chunk_docs):
+        lo = starts[i]; hi = starts[i+args.chunk_docs] if i+args.chunk_docs < len(starts) else len(xnp)
+        x = torch.tensor(xnp[lo:hi].astype(np.int64), device=device)
+        seg, within = doc_index(x, eos)
+        H, _ = apply_stack(x, E, Ps, Bs, within, args, expose=False)
+        Y = torch.cat([E[x][1:], E[x][-1:]], 0)
+        valid = torch.ones(len(x), device=device, dtype=torch.bool)
+        valid[-1] = False; valid[:-1] &= seg[1:].eq(seg[:-1]); valid &= x.ne(eos)
+        X = H[valid].double(); T = Y[valid].double()
+        C += X.T @ T; A += X.T @ X; G += X.T @ T
+    U, S, _ = torch.linalg.svd(C)
+    P = U[:, :args.r].float()
+    diag = torch.trace(A) / d
+    B = torch.linalg.solve(A + args.map_lam * diag * torch.eye(d, device=device, dtype=torch.float64), G).float()
+    log("sv", S[:4].detach().cpu().numpy().round(1))
+    return P, B
+
+
+def fit_eval(train, valid, E, Ps, Bs, eos, V, args, device):
+    def stats(tokens, max_tokens):
+        xnp = tokens[:max_tokens]
+        starts = np.flatnonzero(np.r_[True, xnp[:-1] == eos])
+        A = G = None
+        for i in range(0, len(starts), args.chunk_docs):
+            lo = starts[i]; hi = starts[i+args.chunk_docs] if i+args.chunk_docs < len(starts) else len(xnp)
+            x = torch.tensor(xnp[lo:hi].astype(np.int64), device=device)
+            seg, within = doc_index(x, eos)
+            H, phis = apply_stack(x, E, Ps, Bs, within, args)
+            Phi = features(H, within, phis, args.extra_context)
+            y = torch.empty(len(x), device=device, dtype=torch.long); y[:-1] = x[1:]; y[-1] = eos
+            valid_m = torch.ones(len(x), device=device, dtype=torch.bool)
+            valid_m[-1] = False; valid_m[:-1] &= seg[1:].eq(seg[:-1]); valid_m &= x.ne(eos)
+            Phi = Phi[valid_m]; y = y[valid_m]
+            if A is None:
+                D = Phi.shape[1]
+                A = torch.zeros((D, D), device=device, dtype=torch.float64)
+                G = torch.zeros((D, V), device=device, dtype=torch.float64)
+            A += Phi.double().T @ Phi.double()
+            G.index_add_(1, y, Phi.double().T)
+        return A, G, A.shape[0]
+    A, G, D = stats(train, args.fit_tokens)
+    uni = torch.tensor((np.bincount(train.astype(np.int64), minlength=V)+1), device=device).float()
+    uni = uni / uni.sum()
+    best = None
+    for lam in [float(x) for x in args.lams.split(",")]:
+        diag = torch.trace(A) / D
+        W = torch.linalg.solve(A + lam * diag * torch.eye(D, device=device, dtype=torch.float64), G).float()
+        nll = 0.0; n = 0
+        xnp = valid[:args.eval_tokens]
+        starts = np.flatnonzero(np.r_[True, xnp[:-1] == eos])
+        for i in range(0, len(starts), args.chunk_docs):
+            lo = starts[i]; hi = starts[i+args.chunk_docs] if i+args.chunk_docs < len(starts) else len(xnp)
+            x = torch.tensor(xnp[lo:hi].astype(np.int64), device=device)
+            seg, within = doc_index(x, eos)
+            H, phis = apply_stack(x, E, Ps, Bs, within, args)
+            Phi = features(H, within, phis, args.extra_context)
+            y = torch.empty(len(x), device=device, dtype=torch.long); y[:-1] = x[1:]; y[-1] = eos
+            valid_m = torch.ones(len(x), device=device, dtype=torch.bool)
+            valid_m[-1] = False; valid_m[:-1] &= seg[1:].eq(seg[:-1]); valid_m &= x.ne(eos)
+            Phi = Phi[valid_m]; y = y[valid_m]
+            S = Phi @ W
+            Pp = torch.relu(S) + args.floor * uni[None, :]
+            Pp = Pp / Pp.sum(1, keepdim=True)
+            nll += float(-torch.log(Pp[torch.arange(len(y), device=device), y] + 1e-12).sum())
+            n += len(y)
+        ppl = math.exp(nll / n)
+        log("lam", lam, "ppl", round(ppl, 2))
+        if best is None or ppl < best[1]: best = (lam, ppl)
+    log(f"TORCH_PRED mode={args.value_mode} ppl={best[1]:.2f} lam={best[0]} D={D}")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", default="/workspace/ts_mini")
+    ap.add_argument("--spm_model", default=None)
+    ap.add_argument("--train_bin", default=None)
+    ap.add_argument("--valid_bin", default=None)
+    ap.add_argument("--vocab", type=int, default=1024)
+    ap.add_argument("--d", type=int, default=192)
+    ap.add_argument("--r", type=int, default=64)
+    ap.add_argument("--layers", type=int, default=2)
+    ap.add_argument("--att_window", type=int, default=8)
+    ap.add_argument("--temp", type=float, default=0.3)
+    ap.add_argument("--window", type=int, default=8)
+    ap.add_argument("--extra_context", type=int, default=1)
+    ap.add_argument("--res_scale", type=float, default=0.12)
+    ap.add_argument("--pred_scale", type=float, default=0.04)
+    ap.add_argument("--pred_schedule", choices=["flat", "linear", "late"], default="flat")
+    ap.add_argument("--orth_delta", type=int, default=0)
+    ap.add_argument("--pred_norm", type=int, default=0)
+    ap.add_argument("--pred_features", type=int, default=0)
+    ap.add_argument("--map_lam", type=float, default=0.001)
+    ap.add_argument("--cooc_tokens", type=int, default=1_000_000)
+    ap.add_argument("--proj_tokens", type=int, default=500_000)
+    ap.add_argument("--fit_tokens", type=int, default=800_000)
+    ap.add_argument("--eval_tokens", type=int, default=100_000)
+    ap.add_argument("--chunk_docs", type=int, default=40)
+    ap.add_argument("--lams", default="0.003,0.01,0.03,0.1")
+    ap.add_argument("--floor", type=float, default=1e-4)
+    ap.add_argument("--value_mode", choices=["h","next","delta","ridge_next","ridge_delta",
+                                             "dual_next","dual_ridge_next","dual_ridge_delta"], default="h")
+    args = ap.parse_args()
+    import sentencepiece as spm
+    spm_model = args.spm_model or os.path.join(args.data, f"sp{args.vocab}.model")
+    train_bin = args.train_bin or os.path.join(args.data, "train.bin")
+    valid_bin = args.valid_bin or os.path.join(args.data, "valid.bin")
+    sp = spm.SentencePieceProcessor(model_file=spm_model)
+    eos = sp.eos_id(); V = sp.get_piece_size()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    train = np.fromfile(train_bin, dtype=np.uint16)
+    valid = np.fromfile(valid_bin, dtype=np.uint16)
+    log("mode", args.value_mode, "device", device, "V", V, "train", len(train), "valid", len(valid))
+    E = ppmi_embed(train, V, args.d, args.window, args.cooc_tokens, device)
+    Ps, Bs = [], []
+    for _ in range(args.layers):
+        P, B = learn_map(train, E, Ps, Bs, eos, args, device)
+        Ps.append(P); Bs.append(B)
+    fit_eval(train, valid, E, Ps, Bs, eos, V, args, device)
+
+
+if __name__ == "__main__":
+    main()

stream_ce_lam10.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7addeb9f7d9815da4127edf10f82139f9ea54ee7047c2b21b10f40b5da4a108
+size 570493669

stream_ce_lam10_long.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5ee6c3d43c0bd51d4696b15b2706e644d81535a14b1aeefd7dcc1c2b8e7f9709
+size 570493773

tokenizer/glm16k.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67468b3552cc095e8c25e2bae56dbf4b192ce340d704763129becaf184cf60e5
+size 366707