"""Intrinsic eval — the metrics tracked during pretraining, on a genuinely held-out split (every HOLDOUT_MOD-th record, never trained on; see train/data.py::eligible). bits_per_char masked-char cross-entropy in bits, at a fixed reference mask rate (t=0.30) boundary_f1 word-boundary prediction F1 with the boundary channel forced UNKNOWN (the model must segment scriptio continua from scratch) dia_acc diacritic restoration accuracy, channel forced UNKNOWN punct_acc punctuation restoration accuracy, channel forced UNKNOWN demo a fixed smoke reconstruction, for a human sanity check every checkpoint """ from __future__ import annotations import json, math from pathlib import Path import numpy as np import torch from data.normalize import ALPHABET from model.char_bert import CharBertConfig, CharBertEncoder from train.collate import collate from train.noising import NoiseConfig ALIST = list(ALPHABET) UNK_BND, UNK_DIA, UNK_PUNCT = 3, 48, 6 def held_out_records(shards, n, seed=1234): import pyarrow.parquet as pq d = Path(shards) idx = pq.read_table(d / "index.parquet") offs = idx.column("offset").to_numpy(); lens = idx.column("length").to_numpy() tier = idx.column("tier").to_numpy(zero_copy_only=False) chars = np.memmap(d / "chars.bin", dtype=np.uint8, mode="r") bnd = np.memmap(d / "boundary.bin", dtype=np.uint8, mode="r") dia = np.memmap(d / "dia.bin", dtype=np.uint8, mode="r") cap = np.memmap(d / "cap.bin", dtype=np.uint8, mode="r") punct = np.memmap(d / "punct.bin", dtype=np.uint8, mode="r") from train.data import HOLDOUT_MOD gold = np.flatnonzero((tier == "pristine") & (lens >= 64) & (lens <= 4096)) gold = gold[gold % HOLDOUT_MOD == 0] rng = np.random.default_rng(seed) pick = rng.choice(gold, size=min(n, len(gold)), replace=False) recs = [] for i in pick: o, l = int(offs[i]), int(lens[i]) recs.append(dict(chars=np.array(chars[o:o+l]), boundary=np.array(bnd[o:o+l]), dia=np.array(dia[o:o+l]), cap=np.array(cap[o:o+l]), punct=np.array(punct[o:o+l]))) return recs @torch.no_grad() def evaluate(model, recs, device, T=1024, micro=8, rate=0.30): model.eval() # fixed reference noise: iid at t=rate, boundary/dia/punct channels forced UNKNOWN ncfg = NoiseConfig(w_span=0, w_word=0, w_elastic=0, w_iid=1, w_halfword=0, w_substitute=0, beta_a=1e6, beta_b=1e6 * (1 - rate) / rate, p_bnd_full=0.0, p_bnd_none=1.0, p_dia_full=0.0, p_dia_none=1.0, p_punct_full=0.0, p_punct_none=1.0) g = torch.Generator().manual_seed(0) tot_bits = tot_n = 0 b_tp = b_fp = b_fn = 0 dia_ok = dia_n = 0 punct_ok = punct_n = 0 for i in range(0, len(recs), micro): batch = collate(recs[i:i+micro], ncfg, T, g) batch = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()} with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): out = model(batch) lab = batch["labels"]; m = lab != -100 if m.any(): logp = torch.log_softmax(out["char"][m].float(), -1) nll = -logp.gather(1, lab[m][:, None]).squeeze(1) tot_bits += (nll.sum() / math.log(2)).item(); tot_n += int(m.sum()) bl = batch["bnd_lab"]; bm = bl != -100 if bm.any(): pred = out["boundary"][bm].argmax(-1) gold_fin = bl[bm] >= 1; pred_fin = pred >= 1 b_tp += int((gold_fin & pred_fin).sum()); b_fp += int((~gold_fin & pred_fin).sum()) b_fn += int((gold_fin & ~pred_fin).sum()) dl = batch["dia_lab"]; dm = dl != -100 if dm.any(): dia_ok += int((out["dia"][dm].argmax(-1) == dl[dm]).sum()); dia_n += int(dm.sum()) pl = batch["punct_lab"]; pm = pl != -100 if pm.any(): punct_ok += int((out["punct"][pm].argmax(-1) == pl[pm]).sum()); punct_n += int(pm.sum()) prec = b_tp / (b_tp + b_fp + 1e-9); rec = b_tp / (b_tp + b_fn + 1e-9) bf1 = 2 * prec * rec / (prec + rec + 1e-9) model.train() return dict(bits_per_char=round(tot_bits / max(tot_n, 1), 4), boundary_f1=round(bf1, 4), dia_acc=round(dia_ok / max(dia_n, 1), 4), punct_acc=round(punct_ok / max(punct_n, 1), 4), n_eval=tot_n) @torch.no_grad() def restore_demo(model, device, text="καιολογοσ", mask_slice=(2, 5)): """Mask chars [i:j] and show the model's top-1 reconstruction (the κα---ογοσ demo).""" from data.normalize import Stats, normalize_record st = Stats() r = normalize_record(text, st, with_punct=True) if r is None: return "?" chars, boundary, dia, cap, punct = r ids = torch.tensor(chars, dtype=torch.long) inp = ids.clone(); inp[mask_slice[0]:mask_slice[1]] = 24 # mask_id batch = dict(input_ids=inp[None].to(device), boundary=torch.full((1, len(ids)), UNK_BND, device=device), dia=torch.full((1, len(ids)), UNK_DIA, device=device), punct=torch.full((1, len(ids)), UNK_PUNCT, device=device), seg_id=torch.ones(1, len(ids), dtype=torch.long, device=device)) with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): out = model(batch) pred = out["char"][0].argmax(-1).cpu() recon = "".join(ALIST[c] if c < 24 else "?" for c in pred[mask_slice[0]:mask_slice[1]]) full = "".join(ALIST[inp[k]] if inp[k] < 24 else recon[k-mask_slice[0]] for k in range(len(ids))) return full def load_model(ckpt_path, device): sd = torch.load(ckpt_path, map_location=device) c = sd["cfg"] # metadata-conditioned finetunes (finetune_whole.py) carry e_region/e_century # embeddings; size them from the state dict itself so both kinds of checkpoint load msd = sd["model"] n_region = msd["e_region.weight"].shape[0] if "e_region.weight" in msd else 0 n_century = msd["e_century.weight"].shape[0] if "e_century.weight" in msd else 0 mcfg = CharBertConfig(attn_impl="sdpa", d_model=c["d_model"], n_heads=c["d_model"] // 64, depth=c["depth"], char_window=c["char_window"], qk_norm=c.get("qk_norm", True), n_region=n_region, n_century=n_century) model = CharBertEncoder(mcfg).to(device) model.load_state_dict(msd) return model, c