File size: 2,271 Bytes
ea51a3d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | from __future__ import annotations
import math
from pathlib import Path
import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
ROOT = Path(__file__).resolve().parent
EXPECTED_PARAMS = 96_130_040
print("1/6 tokenizer load + UTF-8 round trip")
tok = AutoTokenizer.from_pretrained(
ROOT,
trust_remote_code=True,
use_fast=False,
)
sample = "TransWave — raw bytes ✓"
ids = tok(sample, return_tensors="pt")["input_ids"]
decoded = tok.decode(ids[0].tolist())
assert decoded == sample, (decoded, sample)
assert tok.vocab_size == 256
print("2/6 config")
cfg = AutoConfig.from_pretrained(ROOT, trust_remote_code=True)
assert cfg.architecture_id == "causal_transwave_blocklifting_meyer_byte_lm_g4_v2"
assert cfg.trained_context_length == 20480
assert cfg.vocab_size == 256
print("3/6 trained checkpoint load")
model = AutoModelForCausalLM.from_pretrained(
ROOT,
trust_remote_code=True,
).eval()
n = sum(p.numel() for p in model.transwave.parameters())
assert n == EXPECTED_PARAMS, n
print("4/6 exact finite NLL")
torch.manual_seed(20261012)
x = torch.randint(0, 256, (1, 32), dtype=torch.long)
with torch.no_grad():
direct = model.transwave.nll(x, checkpoint_layers=False)
wrapped = model(input_ids=x, labels=x).loss
assert torch.isfinite(direct)
assert torch.isfinite(wrapped)
assert abs(float(direct - wrapped)) < 2e-5
print("NLL:", float(wrapped), "PPL:", math.exp(float(wrapped)))
print("5/6 future-perturbation causality")
a = torch.randint(0, 256, (1, 32), dtype=torch.long)
b = a.clone()
b[:, 16:] = torch.randint(0, 256, b[:, 16:].shape, dtype=torch.long)
with torch.no_grad():
ca, _ = model.transwave.contexts(a, checkpoint_layers=False)
cb, _ = model.transwave.contexts(b, checkpoint_layers=False)
diff = float((ca[:, :3, :] - cb[:, :3, :]).abs().max())
assert diff <= 2e-4, diff
print("future->past max diff:", diff)
print("6/6 reference generation")
prompt = tok("abcdefgh", return_tensors="pt")["input_ids"]
with torch.no_grad():
y = model.generate_bytes(prompt, max_new_tokens=2, do_sample=False)
assert y.shape == (1, 10)
assert int(y.min()) >= 0 and int(y.max()) <= 255
print("SMOKE TEST: PASS") |