Buckets:
| """WO38 — urd-g2p-2.0 candidate: byt5-base on the full humair025 split. | |
| urd-g2p-1.0 (byt5-small, 3 epochs): CER 0.1578 on its own held-out | |
| test (WO37 protocol). Arm: same data (610K pairs), same schedule, | |
| byt5-BASE. Gate: test CER (raw char, full 12.7K split) < 0.1578. | |
| Usage (a100): python /ckpt/run-059-urd-g2p-base/train_urd_g2p_hf.py | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import random | |
| import time | |
| from pathlib import Path | |
| import torch | |
| from torch.utils.data import DataLoader, Dataset | |
| DATA = Path("/data-ru/data") | |
| CKPT = Path("/ckpt") / os.environ.get("U_RUN", "run-059-urd-g2p-base") | |
| BACKBONE = os.environ.get("U_BACKBONE", "google/byt5-base") | |
| EPOCHS = int(os.environ.get("U_EPOCHS", "3")) | |
| BS = int(os.environ.get("U_BS", "16")) | |
| ACCUM = int(os.environ.get("U_ACCUM", "2")) | |
| MAXLEN = 256 | |
| def lev(a: str, b: str) -> int: | |
| prev = list(range(len(b) + 1)) | |
| for i, ca in enumerate(a, 1): | |
| cur = [i] | |
| for j, cb in enumerate(b, 1): | |
| cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb))) | |
| prev = cur | |
| return prev[-1] | |
| def main() -> None: | |
| CKPT.mkdir(parents=True, exist_ok=True) | |
| from transformers import (AutoModelForSeq2SeqLM, AutoTokenizer, | |
| get_cosine_schedule_with_warmup) | |
| def load(name): | |
| rows = [json.loads(l) for l in (DATA / name).read_text( | |
| encoding="utf-8").splitlines() if l.strip()] | |
| return [(r["src"].strip(), r["tgt"].strip()) for r in rows | |
| if r.get("src", "").strip() and r.get("tgt", "").strip()] | |
| train = load("train.jsonl") | |
| test = load("test.jsonl") | |
| random.Random(42).shuffle(train) | |
| print(f"[data] train={len(train)} test={len(test)}", flush=True) | |
| tok = AutoTokenizer.from_pretrained(BACKBONE) | |
| model = AutoModelForSeq2SeqLM.from_pretrained(BACKBONE).cuda().train() | |
| model.gradient_checkpointing_enable() | |
| model.config.use_cache = False | |
| class DS(Dataset): | |
| def __len__(self): | |
| return len(train) | |
| def __getitem__(self, i): | |
| return train[i] | |
| def collate(b): | |
| es = [tok(s, truncation=True, max_length=MAXLEN).input_ids for s, _ in b] | |
| ds_ = [tok(t, truncation=True, max_length=MAXLEN).input_ids for _, t in b] | |
| mx = max(len(e) for e in es) | |
| md = max(len(d) for d in ds_) | |
| ids = torch.zeros((len(b), mx), dtype=torch.long) | |
| am = torch.zeros((len(b), mx), dtype=torch.long) | |
| lb = torch.full((len(b), md), -100, dtype=torch.long) | |
| for i, (e, d) in enumerate(zip(es, ds_)): | |
| ids[i, :len(e)] = torch.tensor(e) | |
| am[i, :len(e)] = 1 | |
| lb[i, :len(d)] = torch.tensor(d) | |
| return ids, am, lb | |
| loader = DataLoader(DS(), batch_size=BS, shuffle=True, collate_fn=collate, | |
| num_workers=2, drop_last=True) | |
| total_steps = (len(loader) // ACCUM) * EPOCHS | |
| opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01) | |
| sched = get_cosine_schedule_with_warmup(opt, 200, total_steps) | |
| step = 0 | |
| t0 = time.time() | |
| for _ep in range(EPOCHS): | |
| for ids, am, lb in loader: | |
| ids, am, lb = ids.cuda(), am.cuda(), lb.cuda() | |
| with torch.autocast("cuda", torch.bfloat16): | |
| out = model(input_ids=ids, attention_mask=am, labels=lb) | |
| if not torch.isfinite(out.loss): | |
| print(f"[nan] step {step}", flush=True) | |
| opt.zero_grad() | |
| continue | |
| (out.loss / ACCUM).backward() | |
| if step % ACCUM == ACCUM - 1: | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) | |
| opt.step() | |
| sched.step() | |
| opt.zero_grad() | |
| step += 1 | |
| if step % 200 == 0: | |
| print(f"[train] {step}/{total_steps} loss={float(out.loss):.3f} " | |
| f"({step/(time.time()-t0):.1f} it/s)", flush=True) | |
| model.config.use_cache = True | |
| model.eval() | |
| best = CKPT / "best" | |
| best.mkdir(exist_ok=True) | |
| model.save_pretrained(best) | |
| tok.save_pretrained(best) | |
| # ---- gate: full test split, raw char CER (the WO37 protocol) ---- | |
| preds = [] | |
| with torch.no_grad(): | |
| for i in range(0, len(test), 16): | |
| batch = [s for s, _ in test[i:i + 16]] | |
| enc = tok(batch, return_tensors="pt", padding=True, | |
| truncation=True, max_length=MAXLEN).to("cuda") | |
| with torch.autocast("cuda", torch.bfloat16): | |
| gen = model.generate(**enc, max_new_tokens=MAXLEN, num_beams=1) | |
| preds.extend(tok.batch_decode(gen, skip_special_tokens=True)) | |
| cers, correct = [], 0 | |
| for (s, gold), pred in zip(test, preds): | |
| cers.append(lev(pred, gold) / max(1, len(gold))) | |
| if pred == gold: | |
| correct += 1 | |
| res = {"model": BACKBONE, "pairs": len(train), | |
| "cer": round(sum(cers) / len(cers), 4), | |
| "word_acc": round(correct / len(test), 4), | |
| "gate": "cer < 0.1578 (urd-g2p-1.0)", | |
| "baseline": 0.1578} | |
| (CKPT / "verdict.json").write_text(json.dumps(res, indent=1)) | |
| print(json.dumps(res), flush=True) | |
| (CKPT / "GATE_DONE").write_text("ok\n") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 5.29 kB
- Xet hash:
- 216a5dd6af94face3303fcd501e785f4f3fc019fc3d587ca420a4d80cebc4759
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.