tinydoc-5m v4: OmniDocBench EN text edit 0.117 (v3 0.127), inline maths as text, invoice summary tables

e24bf15 verified 5 days ago

22.1 kB

"""Train one component.
    python -m tinydoc.train recognizer            # CNN + Transformer + CTC line reader
    python -m tinydoc.train detector              # text-segment detector
    python -m tinydoc.train orientation           # 0/90/180/270 page classifier
    python -m tinydoc.train extractor             # field / relation / doc-type model
Checkpoints go to runs/<component>/{last,best}.pt, metrics to runs/<component>/log.jsonl.
Runs on CUDA (with fp16 mixed precision), Apple MPS, or CPU. --max-minutes saves and exits cleanly before
a time limit (e.g. a Kaggle session); --resume continues from last.pt; --init <runs dir> fine-tunes from an
earlier best.pt (heads that grew, e.g. a new label, keep their old rows and initialise the new ones).
"""

import argparse
import json
import math
import os
import time
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

from .charset import BLANK, CHARS
from .data import datasets as D
from .models import build, count_params
from .geometry import rect_quad, transform_quads
from .pipeline.detect import detector_input, kernels_to_quads, match_quads
from .synth.augment import geometric, photometric
from .schema import FIELD_LABELS, REL_TYPES

DEFAULTS = {
    "recognizer": dict(steps=60_000, batch=96, lr=8e-4, eval_every=2_000, warmup=1_500),
    "detector": dict(steps=40_000, batch=16, lr=2e-3, eval_every=2_000, warmup=1_000),
    "orientation": dict(steps=6_000, batch=16, lr=2e-3, eval_every=1_000, warmup=300),   # batch = pages x 8 patches
    "extractor": dict(steps=25_000, batch=32, lr=1e-3, eval_every=1_000, warmup=1_000),
}


# ------------------------------------------------------------------------------------------- losses
def detector_loss(logits, kernel, mask, neg_ratio=3.0):
    """logits/kernel/mask: (B, C, H, W). Channel 0 = line kernels, channel 1 = block kernels (weighted 0.5;
    its mask is zero on pages without block labels)."""
    loss = _kernel_loss(logits[:, :1], kernel[:, :1], mask[:, :1], neg_ratio)
    if logits.shape[1] > 1 and kernel.shape[1] > 1:
        loss = loss + 0.5 * _kernel_loss(logits[:, 1:2], kernel[:, 1:2], mask[:, 1:2], neg_ratio)
    return loss


def _kernel_loss(logits, kernel, mask, neg_ratio=3.0):
    bce = F.binary_cross_entropy_with_logits(logits, kernel, reduction="none") * mask
    pos = (kernel * mask).flatten()
    neg = ((1 - kernel) * mask).flatten()
    n_pos = int(pos.sum().item())
    n_neg = int(min(neg.sum().item(), max(neg_ratio * n_pos, 4096)))
    flat = bce.flatten()
    hard_neg = torch.topk(flat * neg, n_neg).values if n_neg > 0 else flat.new_zeros(1)
    bce_loss = ((flat * pos).sum() + hard_neg.sum()) / (n_pos + n_neg + 1e-6)
    prob = torch.sigmoid(logits) * mask
    inter = (prob * kernel).sum()
    dice = 1 - 2 * inter / (prob.sum() + (kernel * mask).sum() + 1e-6)
    return bce_loss + dice


def extractor_loss(out, b):
    field_logits, rel_logits, doc_logits = out
    lf = F.cross_entropy(field_logits.flatten(0, 1), b["labels"].flatten(), ignore_index=-100, label_smoothing=0.05)
    valid = b["valid"]
    pair = valid[:, None, :, None] * valid[:, None, None, :] * b["rel_known"][:, :, None, None]
    eye = torch.eye(valid.shape[1], device=valid.device)[None, None]
    pair = pair * (1 - eye)
    lr_ = F.binary_cross_entropy_with_logits(rel_logits, b["rel"], pos_weight=torch.tensor(4.0, device=valid.device), reduction="none")
    lr_ = (lr_ * pair).sum() / pair.sum().clamp(min=1)
    ld = F.cross_entropy(doc_logits, b["doc"])
    return lf + 2.0 * lr_ + 0.3 * ld


# -------------------------------------------------------------------------------------------- evals
def cer(pred, gt):
    if not gt:
        return float(len(pred) > 0)
    prev = list(range(len(gt) + 1))
    for i, pc in enumerate(pred, 1):
        cur = [i] + [0] * len(gt)
        for j, gc in enumerate(gt, 1):
            cur[j] = min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (pc != gc))
        prev = cur
    return prev[-1] / len(gt)


def decode_batch(logits):
    ids = logits.argmax(-1).cpu().numpy()
    out = []
    for row in ids:
        s, prev = [], BLANK
        for k in row:
            if k != BLANK and k != prev:
                s.append(CHARS[k - 1])
            prev = k
        out.append("".join(s))
    return out


@torch.no_grad()
def eval_recognizer(model, device, args):
    model.eval()
    ds = D.LineDataset("val", length=args.eval_size, seed=12345)
    dl = torch.utils.data.DataLoader(ds, batch_size=64, num_workers=args.workers, collate_fn=D.collate_lines)
    errs, chars, exact, n = defaultdict(float), defaultdict(int), defaultdict(int), defaultdict(int)
    for x, mask, steps, _, _, srcs, targets in dl:
        preds = decode_batch(model(x.to(device), mask.to(device)))
        for p, t, s in zip(preds, targets, srcs):
            gt = "".join(CHARS[i - 1] for i in t.tolist())
            errs[s] += cer(p, gt) * len(gt)
            chars[s] += len(gt)
            exact[s] += p == gt
            n[s] += 1
    m = {f"cer/{s}": errs[s] / max(chars[s], 1) for s in errs}
    m.update({f"line_acc/{s}": exact[s] / n[s] for s in n})
    m["cer"] = sum(errs.values()) / max(sum(chars.values()), 1)
    model.train()
    return m, -m["cer"]


@torch.no_grad()
def eval_detector(model, device, args):
    model.eval()
    m = {}
    f1s = []
    inter, denom = {}, {}
    for src in ("synth", "doclaynet", "wildreceipt"):
        pool = D.PagePool("val", "detector", {src: 1.0})
        if src not in pool.weights:
            continue
        tp = fp = fn = 0
        for i in range(args.eval_pages):
            rng = np.random.default_rng([999, i])
            img, quads, ign, regions, _, geo, strength, thermal, blocks = pool.sample(rng)
            img, hm = geometric(img, rng, geo)
            quads, regions = transform_quads(quads, hm), transform_quads(regions, hm)
            img = photometric(img, rng, strength, thermal)
            x, scale = detector_input(np.ascontiguousarray(img))
            out = torch.sigmoid(model(torch.from_numpy(x)[None].to(device)))[0].cpu().numpy()
            prob = out[0]
            if blocks is not None and len(blocks[0]) and out.shape[0] > 1:
                # block-kernel dice on this page (targets built like the training ones, in model coordinates)
                bq = [np.asarray(q) * scale for q in transform_quads([rect_quad(*r) for r in blocks[0]], hm)]
                pq = [np.asarray(q) * scale for q in transform_quads([rect_quad(r[0], r[1], r[0] + lh, r[1] + lh)
                                                                      for r, lh in zip(*blocks)], hm)]
                bk, _ = D.block_targets(out.shape[1], out.shape[2], bq, pq, True, [])
                pb = out[1] > 0.5
                inter[src] = inter.get(src, 0) + float((pb & (bk > 0)).sum())
                denom[src] = denom.get(src, 0) + float(pb.sum() + (bk > 0).sum())
            pred = [q / scale for q, _ in kernels_to_quads(prob)]
            gt = [q for q, g in zip(quads, ign) if not g]
            # predictions overlapping ignore regions/boxes are neither right nor wrong
            ign_q = [q for q, g in zip(quads, ign) if g] + list(regions)
            matches, up, ug = match_quads(pred, gt)
            if ign_q:
                _, up_after, _ = match_quads([pred[i] for i in up], ign_q, 0.1)
                up = up_after
            tp, fp, fn = tp + len(matches), fp + len(up), fn + len(ug)
        p, r = tp / max(tp + fp, 1), tp / max(tp + fn, 1)
        f1 = 2 * p * r / max(p + r, 1e-6)
        m.update({f"precision/{src}": p, f"recall/{src}": r, f"f1/{src}": f1})
        f1s.append(f1)
    m["f1"] = float(np.mean(f1s))
    for src in inter:
        m[f"block_dice/{src}"] = 2 * inter[src] / max(denom[src], 1)
    score = m["f1"]
    if inter:
        m["block_dice"] = float(np.mean([m[f"block_dice/{s_}"] for s_ in inter]))
        score = m["f1"] + 0.25 * m["block_dice"]   # lines first; blocks break ties between line-equal checkpoints
    model.train()
    return m, score


@torch.no_grad()
def eval_orientation(model, device, args):
    model.eval()
    ds = D.OrientationDataset("val", length=max(1, args.eval_size // 8), seed=4242)
    dl = torch.utils.data.DataLoader(ds, batch_size=16, num_workers=args.workers)
    correct = total = 0
    for x, y in dl:
        x, y = x.flatten(0, 1), y.flatten()
        pred = model(x.to(device)).argmax(-1).cpu()
        correct += (pred == y).sum().item()
        total += len(y)
    model.train()
    acc = correct / total
    return {"patch_acc": acc}, acc


@torch.no_grad()
def eval_extractor(model, device, args):
    model.eval()
    ds = D.LayoutDataset("val", length=args.eval_size, seed=777)
    dl = torch.utils.data.DataLoader(ds, batch_size=32, num_workers=args.workers, collate_fn=D.collate_layouts)
    R = len(REL_TYPES)
    by_src = defaultdict(lambda: {"tp": np.zeros(len(FIELD_LABELS)), "fp": np.zeros(len(FIELD_LABELS)), "fn": np.zeros(len(FIELD_LABELS)),
                                  "rel_tp": np.zeros(R), "rel_fp": np.zeros(R), "rel_fn": np.zeros(R), "doc": 0, "n": 0})
    for b in dl:
        t = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in b.items()}
        fl, rl, dl_ = model(t["char_ids"], t["feats"], t["rel_boxes"], t["valid"])
        pred = fl.argmax(-1).cpu().numpy()
        rel = (torch.sigmoid(rl) > 0.5).cpu().numpy()
        doc = dl_.argmax(-1).cpu().numpy()
        for i, src in enumerate(b["src"]):
            s = by_src[src]
            n = int(b["valid"][i].sum())
            lab = b["labels"][i, :n].numpy()
            pr = pred[i, :n]
            known = lab >= 0
            for c in range(len(FIELD_LABELS)):
                s["tp"][c] += ((pr == c) & (lab == c) & known).sum()
                s["fp"][c] += ((pr == c) & (lab != c) & known).sum()
                s["fn"][c] += ((pr != c) & (lab == c) & known).sum()
            off_diag = ~np.eye(n, dtype=bool)
            for r in range(R):
                if b["rel_known"][i, r] > 0:
                    g = b["rel"][i, r, :n, :n].numpy() > 0.5
                    p = rel[i, r, :n, :n] & off_diag
                    s["rel_tp"][r] += int((p & g).sum())
                    s["rel_fp"][r] += int((p & ~g).sum())
                    s["rel_fn"][r] += int((~p & g).sum())
            s["doc"] += int(doc[i] == b["doc"][i])
            s["n"] += 1
    m, f1s = {}, []
    for src, s in by_src.items():
        present = (s["tp"] + s["fn"]) > 0
        present[0] = False  # "other" is not a field
        f1 = 2 * s["tp"] / np.maximum(2 * s["tp"] + s["fp"] + s["fn"], 1)
        mf1 = float(f1[present].mean()) if present.any() else 0.0
        m[f"field_macro_f1/{src}"] = mf1
        f1s.append(mf1)
        m[f"doc_acc/{src}"] = s["doc"] / s["n"]
        for r, name in enumerate(REL_TYPES):
            if s["rel_tp"][r] + s["rel_fn"][r] > 0:
                m[f"{name}_f1/{src}"] = float(2 * s["rel_tp"][r] / max(2 * s["rel_tp"][r] + s["rel_fp"][r] + s["rel_fn"][r], 1))
    m["field_macro_f1"] = float(np.mean(f1s))
    blocks = [v for k, v in m.items() if k.startswith("same_block_f1/")]
    score = m["field_macro_f1"]
    if blocks:
        m["same_block_f1"] = float(np.mean(blocks))
        score = 0.5 * (score + m["same_block_f1"])   # fields and page structure count equally
    model.train()
    return m, score


# --------------------------------------------------------------------------------------------- loop
BUCKET_GROUP = 4  # batches drawn together and regrouped by length


def identity_collate(samples):
    return samples


def bucketed(loader, size_of, collate, batch, rng):
    """Regroup each mega-batch so that samples of similar length share a batch (less padding)."""
    for mega in loader:
        mega = sorted(mega, key=size_of)
        chunks = [mega[i:i + batch] for i in range(0, len(mega) - batch + 1, batch)]
        for k in rng.permutation(len(chunks)):
            yield collate(chunks[k])


def make_loader(component, args):
    n = args.steps * args.batch
    if component in ("recognizer", "extractor"):   # whole mega-batches, so the stream never ends before the last step
        mega = args.batch * BUCKET_GROUP
        n = -(-n // mega) * mega
    if component == "recognizer":
        ds, collate = D.LineDataset("train", length=n, seed=args.seed), D.collate_lines
    elif component == "detector":
        ds, collate = D.DetectorDataset("train", length=n, seed=args.seed), None
    elif component == "orientation":
        ds, collate = D.OrientationDataset("train", length=n, seed=args.seed), None
    else:
        ds, collate = D.LayoutDataset("train", length=n, seed=args.seed), D.collate_layouts
    sampler = range(args.start_step * args.batch, n)  # sequential indices: resuming continues the stream
    size_of = {"recognizer": lambda s: s[0].shape[-1], "extractor": lambda s: s["feats"].shape[0]}.get(component)
    if size_of is None:
        return torch.utils.data.DataLoader(ds, batch_size=args.batch, sampler=sampler, num_workers=args.workers,
                                           collate_fn=collate, persistent_workers=args.workers > 0,
                                           prefetch_factor=4 if args.workers > 0 else None, drop_last=True)
    mega = torch.utils.data.DataLoader(ds, batch_size=args.batch * BUCKET_GROUP, sampler=sampler, num_workers=args.workers,
                                       collate_fn=identity_collate, persistent_workers=args.workers > 0,
                                       prefetch_factor=2 if args.workers > 0 else None, drop_last=True)
    return bucketed(mega, size_of, collate, args.batch, np.random.default_rng(args.seed + args.start_step))


def step_loss(component, model, batch, device):
    if component == "recognizer":
        x, mask, steps, targets, tlens, _, _ = batch
        logits = model(x.to(device), mask.to(device))
        logp = logits.float().log_softmax(-1).transpose(0, 1)
        return F.ctc_loss(logp, targets.to(device), steps.to(device), tlens.to(device), blank=BLANK, zero_infinity=True)
    if component == "detector":
        x, k, m = (t.to(device) for t in batch)
        return detector_loss(model(x).float(), k, m)
    if component == "orientation":
        x, y = batch
        x, y = x.flatten(0, 1), y.flatten()
        return F.cross_entropy(model(x.to(device)).float(), y.to(device), label_smoothing=0.05)
    b = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()}
    return extractor_loss(model(b["char_ids"], b["feats"], b["rel_boxes"], b["valid"]), b)


def init_from(model, path):
    """Load matching weights from a checkpoint. A tensor that grew (more labels or relations, appended at
    the end) gets its old values in the leading slice; the remainder keeps its fresh initialisation."""
    old = torch.load(path, map_location="cpu", weights_only=False)["model"]
    state = model.state_dict()
    full = partial = 0
    for k, v in state.items():
        o = old.get(k)
        if o is None or o.ndim != v.ndim:
            continue
        if o.shape == v.shape:
            state[k] = o
            full += 1
        elif all(a <= b for a, b in zip(o.shape, v.shape)):
            v = v.clone()
            v[tuple(slice(0, a) for a in o.shape)] = o
            state[k] = v
            partial += 1
    model.load_state_dict(state)
    print(f"initialised from {path}: {full} tensors copied, {partial} grown, {len(state) - full - partial} new")


EVALS = {"recognizer": eval_recognizer, "detector": eval_detector, "orientation": eval_orientation, "extractor": eval_extractor}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("component", choices=list(DEFAULTS))
    ap.add_argument("--steps", type=int)
    ap.add_argument("--batch", type=int)
    ap.add_argument("--lr", type=float)
    ap.add_argument("--eval-every", type=int)
    ap.add_argument("--warmup", type=int)
    ap.add_argument("--eval-size", type=int, default=2000, help="samples per eval (lines / patches / docs)")
    ap.add_argument("--eval-pages", type=int, default=40, help="detector eval pages per source")
    ap.add_argument("--workers", type=int, default=min(6, max(1, (os.cpu_count() or 8) - 3)),
                    help="each worker holds its own PyTorch import (~300 MB); keep modest on 24 GB machines")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
    ap.add_argument("--amp", action=argparse.BooleanOptionalAction, default=None, help="fp16 mixed precision (default: on for CUDA)")
    ap.add_argument("--max-minutes", type=float, default=None, help="save a checkpoint and stop after this long")
    ap.add_argument("--real-root", help="override data/real (sets TINYDOC_REAL_ROOT)")
    ap.add_argument("--synth-root", help="override data/synth (sets TINYDOC_SYNTH_ROOT)")
    ap.add_argument("--font-root", help="override fonts/ (sets TINYDOC_FONT_ROOT)")
    ap.add_argument("--out", default="runs")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--resume", action="store_true")
    ap.add_argument("--init", help="runs dir whose <component>/best.pt initialises the model (fine-tuning)")
    ap.add_argument("--arch", default="", help="model options saved with the checkpoint, e.g. cnn=separable")
    args = ap.parse_args()
    arch = {k: int(v) if v.isdigit() else v for k, v in (kv.split("=", 1) for kv in args.arch.split(",") if kv)}
    for flag, env in (("real_root", "TINYDOC_REAL_ROOT"), ("synth_root", "TINYDOC_SYNTH_ROOT"), ("font_root", "TINYDOC_FONT_ROOT")):
        if getattr(args, flag):
            os.environ[env] = getattr(args, flag)  # inherited by DataLoader workers
    for k, v in DEFAULTS[args.component].items():
        if getattr(args, k) is None:
            setattr(args, k, v * 0.5 if (k == "lr" and args.init) else v)   # gentler when fine-tuning

    device = torch.device(args.device)
    use_amp = (device.type == "cuda") if args.amp is None else args.amp
    if device.type == "cuda":
        torch.backends.cudnn.benchmark = args.component in ("detector", "orientation")  # fixed input shapes
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp and device.type == "cuda")
    deadline = time.time() + args.max_minutes * 60 if args.max_minutes else None
    out = Path(args.out) / args.component
    out.mkdir(parents=True, exist_ok=True)
    model = build(args.component, **arch).to(device)
    decay = [p for n, p in model.named_parameters() if p.ndim >= 2]
    no_decay = [p for n, p in model.named_parameters() if p.ndim < 2]
    opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.01}, {"params": no_decay, "weight_decay": 0.0}], lr=args.lr)
    args.start_step, best = 0, -math.inf
    if args.resume and (out / "last.pt").exists():
        ck = torch.load(out / "last.pt", map_location="cpu", weights_only=False)  # our own checkpoint
        model.load_state_dict(ck["model"])
        opt.load_state_dict(ck["optim"])
        args.start_step, best = ck["step"], ck.get("best", -math.inf)
        if "scaler" in ck and scaler.is_enabled():
            scaler.load_state_dict(ck["scaler"])
        print(f"resumed from step {args.start_step}")
    elif args.init:
        init_from(model, Path(args.init) / args.component / "best.pt")
    if args.start_step >= args.steps:
        print("already finished")
        return
    print(f"{args.component}: {count_params(model):,} params, device={device}, workers={args.workers}, "
          f"steps={args.steps}, batch={args.batch}, lr={args.lr}, amp={use_amp}")

    def lr_at(s):
        if s < args.warmup:
            return args.lr * (s + 1) / args.warmup
        t = (s - args.warmup) / max(1, args.steps - args.warmup)
        return args.lr * (0.05 + 0.95 * 0.5 * (1 + math.cos(math.pi * t)))

    log = open(out / "log.jsonl", "a")
    loader = make_loader(args.component, args)
    model.train()
    step, t0, running = args.start_step, time.time(), []
    for batch in loader:
        for g in opt.param_groups:
            g["lr"] = lr_at(step)
        with torch.autocast(device.type, dtype=torch.float16, enabled=use_amp):
            loss = step_loss(args.component, model, batch, device)
        opt.zero_grad(set_to_none=True)
        scaler.scale(loss).backward()
        scaler.unscale_(opt)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        scaler.step(opt)
        scaler.update()
        running.append(loss.item())
        step += 1
        if step % 50 == 0:
            dt = time.time() - t0
            print(f"step {step:6d}  loss {np.mean(running):.4f}  lr {lr_at(step):.2e}  {50 * args.batch / dt:.0f} samples/s", flush=True)
            log.write(json.dumps({"step": step, "loss": float(np.mean(running)), "lr": lr_at(step)}) + "\n")
            running, t0 = [], time.time()
        out_of_time = deadline is not None and time.time() > deadline
        if step % args.eval_every == 0 or step == args.steps or out_of_time:
            metrics, score = EVALS[args.component](model, device, args)
            is_best = score > best
            best = max(best, score)
            ck = {"model": model.state_dict(), "optim": opt.state_dict(), "step": step, "best": best,
                  "component": args.component, "metrics": metrics, "scaler": scaler.state_dict(), "arch": arch}
            torch.save(ck, out / "last.pt")
            if is_best:
                torch.save({k: v for k, v in ck.items() if k != "optim"}, out / "best.pt")
            print(f"eval @ {step}: " + "  ".join(f"{k}={v:.4f}" for k, v in sorted(metrics.items())) + ("  *best*" if is_best else ""), flush=True)
            log.write(json.dumps({"step": step, "eval": metrics, "best": is_best}) + "\n")
            log.flush()
            t0 = time.time()
        if out_of_time:
            print(f"time budget reached at step {step}; resume with --resume", flush=True)
            break
        if step >= args.steps:
            break


if __name__ == "__main__":
    main()

	"""Train one component.

	python -m tinydoc.train recognizer # CNN + Transformer + CTC line reader
	python -m tinydoc.train detector # text-segment detector
	python -m tinydoc.train orientation # 0/90/180/270 page classifier
	python -m tinydoc.train extractor # field / relation / doc-type model

	Checkpoints go to runs/<component>/{last,best}.pt, metrics to runs/<component>/log.jsonl.
	Runs on CUDA (with fp16 mixed precision), Apple MPS, or CPU. --max-minutes saves and exits cleanly before
	a time limit (e.g. a Kaggle session); --resume continues from last.pt; --init <runs dir> fine-tunes from an
	earlier best.pt (heads that grew, e.g. a new label, keep their old rows and initialise the new ones).
	"""

	import argparse
	import json
	import math
	import os
	import time
	from collections import defaultdict
	from pathlib import Path

	import numpy as np
	import torch
	import torch.nn.functional as F

	from .charset import BLANK, CHARS
	from .data import datasets as D
	from .models import build, count_params
	from .geometry import rect_quad, transform_quads
	from .pipeline.detect import detector_input, kernels_to_quads, match_quads
	from .synth.augment import geometric, photometric
	from .schema import FIELD_LABELS, REL_TYPES

	DEFAULTS = {
	"recognizer": dict(steps=60_000, batch=96, lr=8e-4, eval_every=2_000, warmup=1_500),
	"detector": dict(steps=40_000, batch=16, lr=2e-3, eval_every=2_000, warmup=1_000),
	"orientation": dict(steps=6_000, batch=16, lr=2e-3, eval_every=1_000, warmup=300), # batch = pages x 8 patches
	"extractor": dict(steps=25_000, batch=32, lr=1e-3, eval_every=1_000, warmup=1_000),
	}


	# ------------------------------------------------------------------------------------------- losses
	def detector_loss(logits, kernel, mask, neg_ratio=3.0):
	"""logits/kernel/mask: (B, C, H, W). Channel 0 = line kernels, channel 1 = block kernels (weighted 0.5;
	its mask is zero on pages without block labels)."""
	loss = _kernel_loss(logits[:, :1], kernel[:, :1], mask[:, :1], neg_ratio)
	if logits.shape[1] > 1 and kernel.shape[1] > 1:
	loss = loss + 0.5 * _kernel_loss(logits[:, 1:2], kernel[:, 1:2], mask[:, 1:2], neg_ratio)
	return loss


	def _kernel_loss(logits, kernel, mask, neg_ratio=3.0):
	bce = F.binary_cross_entropy_with_logits(logits, kernel, reduction="none") * mask
	pos = (kernel * mask).flatten()
	neg = ((1 - kernel) * mask).flatten()
	n_pos = int(pos.sum().item())
	n_neg = int(min(neg.sum().item(), max(neg_ratio * n_pos, 4096)))
	flat = bce.flatten()
	hard_neg = torch.topk(flat * neg, n_neg).values if n_neg > 0 else flat.new_zeros(1)
	bce_loss = ((flat * pos).sum() + hard_neg.sum()) / (n_pos + n_neg + 1e-6)
	prob = torch.sigmoid(logits) * mask
	inter = (prob * kernel).sum()
	dice = 1 - 2 * inter / (prob.sum() + (kernel * mask).sum() + 1e-6)
	return bce_loss + dice


	def extractor_loss(out, b):
	field_logits, rel_logits, doc_logits = out
	lf = F.cross_entropy(field_logits.flatten(0, 1), b["labels"].flatten(), ignore_index=-100, label_smoothing=0.05)
	valid = b["valid"]
	pair = valid[:, None, :, None] * valid[:, None, None, :] * b["rel_known"][:, :, None, None]
	eye = torch.eye(valid.shape[1], device=valid.device)[None, None]
	pair = pair * (1 - eye)
	lr_ = F.binary_cross_entropy_with_logits(rel_logits, b["rel"], pos_weight=torch.tensor(4.0, device=valid.device), reduction="none")
	lr_ = (lr_ * pair).sum() / pair.sum().clamp(min=1)
	ld = F.cross_entropy(doc_logits, b["doc"])
	return lf + 2.0 * lr_ + 0.3 * ld


	# -------------------------------------------------------------------------------------------- evals
	def cer(pred, gt):
	if not gt:
	return float(len(pred) > 0)
	prev = list(range(len(gt) + 1))
	for i, pc in enumerate(pred, 1):
	cur = [i] + [0] * len(gt)
	for j, gc in enumerate(gt, 1):
	cur[j] = min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (pc != gc))
	prev = cur
	return prev[-1] / len(gt)


	def decode_batch(logits):
	ids = logits.argmax(-1).cpu().numpy()
	out = []
	for row in ids:
	s, prev = [], BLANK
	for k in row:
	if k != BLANK and k != prev:
	s.append(CHARS[k - 1])
	prev = k
	out.append("".join(s))
	return out


	@torch.no_grad()
	def eval_recognizer(model, device, args):
	model.eval()
	ds = D.LineDataset("val", length=args.eval_size, seed=12345)
	dl = torch.utils.data.DataLoader(ds, batch_size=64, num_workers=args.workers, collate_fn=D.collate_lines)
	errs, chars, exact, n = defaultdict(float), defaultdict(int), defaultdict(int), defaultdict(int)
	for x, mask, steps, _, _, srcs, targets in dl:
	preds = decode_batch(model(x.to(device), mask.to(device)))
	for p, t, s in zip(preds, targets, srcs):
	gt = "".join(CHARS[i - 1] for i in t.tolist())
	errs[s] += cer(p, gt) * len(gt)
	chars[s] += len(gt)
	exact[s] += p == gt
	n[s] += 1
	m = {f"cer/{s}": errs[s] / max(chars[s], 1) for s in errs}
	m.update({f"line_acc/{s}": exact[s] / n[s] for s in n})
	m["cer"] = sum(errs.values()) / max(sum(chars.values()), 1)
	model.train()
	return m, -m["cer"]


	@torch.no_grad()
	def eval_detector(model, device, args):
	model.eval()
	m = {}
	f1s = []
	inter, denom = {}, {}
	for src in ("synth", "doclaynet", "wildreceipt"):
	pool = D.PagePool("val", "detector", {src: 1.0})
	if src not in pool.weights:
	continue
	tp = fp = fn = 0
	for i in range(args.eval_pages):
	rng = np.random.default_rng([999, i])
	img, quads, ign, regions, _, geo, strength, thermal, blocks = pool.sample(rng)
	img, hm = geometric(img, rng, geo)
	quads, regions = transform_quads(quads, hm), transform_quads(regions, hm)
	img = photometric(img, rng, strength, thermal)
	x, scale = detector_input(np.ascontiguousarray(img))
	out = torch.sigmoid(model(torch.from_numpy(x)[None].to(device)))[0].cpu().numpy()
	prob = out[0]
	if blocks is not None and len(blocks[0]) and out.shape[0] > 1:
	# block-kernel dice on this page (targets built like the training ones, in model coordinates)
	bq = [np.asarray(q) * scale for q in transform_quads([rect_quad(*r) for r in blocks[0]], hm)]
	pq = [np.asarray(q) * scale for q in transform_quads([rect_quad(r[0], r[1], r[0] + lh, r[1] + lh)
	for r, lh in zip(*blocks)], hm)]
	bk, _ = D.block_targets(out.shape[1], out.shape[2], bq, pq, True, [])
	pb = out[1] > 0.5
	inter[src] = inter.get(src, 0) + float((pb & (bk > 0)).sum())
	denom[src] = denom.get(src, 0) + float(pb.sum() + (bk > 0).sum())
	pred = [q / scale for q, _ in kernels_to_quads(prob)]
	gt = [q for q, g in zip(quads, ign) if not g]
	# predictions overlapping ignore regions/boxes are neither right nor wrong
	ign_q = [q for q, g in zip(quads, ign) if g] + list(regions)
	matches, up, ug = match_quads(pred, gt)
	if ign_q:
	_, up_after, _ = match_quads([pred[i] for i in up], ign_q, 0.1)
	up = up_after
	tp, fp, fn = tp + len(matches), fp + len(up), fn + len(ug)
	p, r = tp / max(tp + fp, 1), tp / max(tp + fn, 1)
	f1 = 2 * p * r / max(p + r, 1e-6)
	m.update({f"precision/{src}": p, f"recall/{src}": r, f"f1/{src}": f1})
	f1s.append(f1)
	m["f1"] = float(np.mean(f1s))
	for src in inter:
	m[f"block_dice/{src}"] = 2 * inter[src] / max(denom[src], 1)
	score = m["f1"]
	if inter:
	m["block_dice"] = float(np.mean([m[f"block_dice/{s_}"] for s_ in inter]))
	score = m["f1"] + 0.25 * m["block_dice"] # lines first; blocks break ties between line-equal checkpoints
	model.train()
	return m, score


	@torch.no_grad()
	def eval_orientation(model, device, args):
	model.eval()
	ds = D.OrientationDataset("val", length=max(1, args.eval_size // 8), seed=4242)
	dl = torch.utils.data.DataLoader(ds, batch_size=16, num_workers=args.workers)
	correct = total = 0
	for x, y in dl:
	x, y = x.flatten(0, 1), y.flatten()
	pred = model(x.to(device)).argmax(-1).cpu()
	correct += (pred == y).sum().item()
	total += len(y)
	model.train()
	acc = correct / total
	return {"patch_acc": acc}, acc


	@torch.no_grad()
	def eval_extractor(model, device, args):
	model.eval()
	ds = D.LayoutDataset("val", length=args.eval_size, seed=777)
	dl = torch.utils.data.DataLoader(ds, batch_size=32, num_workers=args.workers, collate_fn=D.collate_layouts)
	R = len(REL_TYPES)
	by_src = defaultdict(lambda: {"tp": np.zeros(len(FIELD_LABELS)), "fp": np.zeros(len(FIELD_LABELS)), "fn": np.zeros(len(FIELD_LABELS)),
	"rel_tp": np.zeros(R), "rel_fp": np.zeros(R), "rel_fn": np.zeros(R), "doc": 0, "n": 0})
	for b in dl:
	t = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in b.items()}
	fl, rl, dl_ = model(t["char_ids"], t["feats"], t["rel_boxes"], t["valid"])
	pred = fl.argmax(-1).cpu().numpy()
	rel = (torch.sigmoid(rl) > 0.5).cpu().numpy()
	doc = dl_.argmax(-1).cpu().numpy()
	for i, src in enumerate(b["src"]):
	s = by_src[src]
	n = int(b["valid"][i].sum())
	lab = b["labels"][i, :n].numpy()
	pr = pred[i, :n]
	known = lab >= 0
	for c in range(len(FIELD_LABELS)):
	s["tp"][c] += ((pr == c) & (lab == c) & known).sum()
	s["fp"][c] += ((pr == c) & (lab != c) & known).sum()
	s["fn"][c] += ((pr != c) & (lab == c) & known).sum()
	off_diag = ~np.eye(n, dtype=bool)
	for r in range(R):
	if b["rel_known"][i, r] > 0:
	g = b["rel"][i, r, :n, :n].numpy() > 0.5
	p = rel[i, r, :n, :n] & off_diag
	s["rel_tp"][r] += int((p & g).sum())
	s["rel_fp"][r] += int((p & ~g).sum())
	s["rel_fn"][r] += int((~p & g).sum())
	s["doc"] += int(doc[i] == b["doc"][i])
	s["n"] += 1
	m, f1s = {}, []
	for src, s in by_src.items():
	present = (s["tp"] + s["fn"]) > 0
	present[0] = False # "other" is not a field
	f1 = 2 * s["tp"] / np.maximum(2 * s["tp"] + s["fp"] + s["fn"], 1)
	mf1 = float(f1[present].mean()) if present.any() else 0.0
	m[f"field_macro_f1/{src}"] = mf1
	f1s.append(mf1)
	m[f"doc_acc/{src}"] = s["doc"] / s["n"]
	for r, name in enumerate(REL_TYPES):
	if s["rel_tp"][r] + s["rel_fn"][r] > 0:
	m[f"{name}_f1/{src}"] = float(2 * s["rel_tp"][r] / max(2 * s["rel_tp"][r] + s["rel_fp"][r] + s["rel_fn"][r], 1))
	m["field_macro_f1"] = float(np.mean(f1s))
	blocks = [v for k, v in m.items() if k.startswith("same_block_f1/")]
	score = m["field_macro_f1"]
	if blocks:
	m["same_block_f1"] = float(np.mean(blocks))
	score = 0.5 * (score + m["same_block_f1"]) # fields and page structure count equally
	model.train()
	return m, score


	# --------------------------------------------------------------------------------------------- loop
	BUCKET_GROUP = 4 # batches drawn together and regrouped by length


	def identity_collate(samples):
	return samples


	def bucketed(loader, size_of, collate, batch, rng):
	"""Regroup each mega-batch so that samples of similar length share a batch (less padding)."""
	for mega in loader:
	mega = sorted(mega, key=size_of)
	chunks = [mega[i:i + batch] for i in range(0, len(mega) - batch + 1, batch)]
	for k in rng.permutation(len(chunks)):
	yield collate(chunks[k])


	def make_loader(component, args):
	n = args.steps * args.batch
	if component in ("recognizer", "extractor"): # whole mega-batches, so the stream never ends before the last step
	mega = args.batch * BUCKET_GROUP
	n = -(-n // mega) * mega
	if component == "recognizer":
	ds, collate = D.LineDataset("train", length=n, seed=args.seed), D.collate_lines
	elif component == "detector":
	ds, collate = D.DetectorDataset("train", length=n, seed=args.seed), None
	elif component == "orientation":
	ds, collate = D.OrientationDataset("train", length=n, seed=args.seed), None
	else:
	ds, collate = D.LayoutDataset("train", length=n, seed=args.seed), D.collate_layouts
	sampler = range(args.start_step * args.batch, n) # sequential indices: resuming continues the stream
	size_of = {"recognizer": lambda s: s[0].shape[-1], "extractor": lambda s: s["feats"].shape[0]}.get(component)
	if size_of is None:
	return torch.utils.data.DataLoader(ds, batch_size=args.batch, sampler=sampler, num_workers=args.workers,
	collate_fn=collate, persistent_workers=args.workers > 0,
	prefetch_factor=4 if args.workers > 0 else None, drop_last=True)
	mega = torch.utils.data.DataLoader(ds, batch_size=args.batch * BUCKET_GROUP, sampler=sampler, num_workers=args.workers,
	collate_fn=identity_collate, persistent_workers=args.workers > 0,
	prefetch_factor=2 if args.workers > 0 else None, drop_last=True)
	return bucketed(mega, size_of, collate, args.batch, np.random.default_rng(args.seed + args.start_step))


	def step_loss(component, model, batch, device):
	if component == "recognizer":
	x, mask, steps, targets, tlens, _, _ = batch
	logits = model(x.to(device), mask.to(device))
	logp = logits.float().log_softmax(-1).transpose(0, 1)
	return F.ctc_loss(logp, targets.to(device), steps.to(device), tlens.to(device), blank=BLANK, zero_infinity=True)
	if component == "detector":
	x, k, m = (t.to(device) for t in batch)
	return detector_loss(model(x).float(), k, m)
	if component == "orientation":
	x, y = batch
	x, y = x.flatten(0, 1), y.flatten()
	return F.cross_entropy(model(x.to(device)).float(), y.to(device), label_smoothing=0.05)
	b = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()}
	return extractor_loss(model(b["char_ids"], b["feats"], b["rel_boxes"], b["valid"]), b)


	def init_from(model, path):
	"""Load matching weights from a checkpoint. A tensor that grew (more labels or relations, appended at
	the end) gets its old values in the leading slice; the remainder keeps its fresh initialisation."""
	old = torch.load(path, map_location="cpu", weights_only=False)["model"]
	state = model.state_dict()
	full = partial = 0
	for k, v in state.items():
	o = old.get(k)
	if o is None or o.ndim != v.ndim:
	continue
	if o.shape == v.shape:
	state[k] = o
	full += 1
	elif all(a <= b for a, b in zip(o.shape, v.shape)):
	v = v.clone()
	v[tuple(slice(0, a) for a in o.shape)] = o
	state[k] = v
	partial += 1
	model.load_state_dict(state)
	print(f"initialised from {path}: {full} tensors copied, {partial} grown, {len(state) - full - partial} new")


	EVALS = {"recognizer": eval_recognizer, "detector": eval_detector, "orientation": eval_orientation, "extractor": eval_extractor}


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("component", choices=list(DEFAULTS))
	ap.add_argument("--steps", type=int)
	ap.add_argument("--batch", type=int)
	ap.add_argument("--lr", type=float)
	ap.add_argument("--eval-every", type=int)
	ap.add_argument("--warmup", type=int)
	ap.add_argument("--eval-size", type=int, default=2000, help="samples per eval (lines / patches / docs)")
	ap.add_argument("--eval-pages", type=int, default=40, help="detector eval pages per source")
	ap.add_argument("--workers", type=int, default=min(6, max(1, (os.cpu_count() or 8) - 3)),
	help="each worker holds its own PyTorch import (~300 MB); keep modest on 24 GB machines")
	ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
	ap.add_argument("--amp", action=argparse.BooleanOptionalAction, default=None, help="fp16 mixed precision (default: on for CUDA)")
	ap.add_argument("--max-minutes", type=float, default=None, help="save a checkpoint and stop after this long")
	ap.add_argument("--real-root", help="override data/real (sets TINYDOC_REAL_ROOT)")
	ap.add_argument("--synth-root", help="override data/synth (sets TINYDOC_SYNTH_ROOT)")
	ap.add_argument("--font-root", help="override fonts/ (sets TINYDOC_FONT_ROOT)")
	ap.add_argument("--out", default="runs")
	ap.add_argument("--seed", type=int, default=0)
	ap.add_argument("--resume", action="store_true")
	ap.add_argument("--init", help="runs dir whose <component>/best.pt initialises the model (fine-tuning)")
	ap.add_argument("--arch", default="", help="model options saved with the checkpoint, e.g. cnn=separable")
	args = ap.parse_args()
	arch = {k: int(v) if v.isdigit() else v for k, v in (kv.split("=", 1) for kv in args.arch.split(",") if kv)}
	for flag, env in (("real_root", "TINYDOC_REAL_ROOT"), ("synth_root", "TINYDOC_SYNTH_ROOT"), ("font_root", "TINYDOC_FONT_ROOT")):
	if getattr(args, flag):
	os.environ[env] = getattr(args, flag) # inherited by DataLoader workers
	for k, v in DEFAULTS[args.component].items():
	if getattr(args, k) is None:
	setattr(args, k, v * 0.5 if (k == "lr" and args.init) else v) # gentler when fine-tuning

	device = torch.device(args.device)
	use_amp = (device.type == "cuda") if args.amp is None else args.amp
	if device.type == "cuda":
	torch.backends.cudnn.benchmark = args.component in ("detector", "orientation") # fixed input shapes
	scaler = torch.amp.GradScaler("cuda", enabled=use_amp and device.type == "cuda")
	deadline = time.time() + args.max_minutes * 60 if args.max_minutes else None
	out = Path(args.out) / args.component
	out.mkdir(parents=True, exist_ok=True)
	model = build(args.component, **arch).to(device)
	decay = [p for n, p in model.named_parameters() if p.ndim >= 2]
	no_decay = [p for n, p in model.named_parameters() if p.ndim < 2]
	opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.01}, {"params": no_decay, "weight_decay": 0.0}], lr=args.lr)
	args.start_step, best = 0, -math.inf
	if args.resume and (out / "last.pt").exists():
	ck = torch.load(out / "last.pt", map_location="cpu", weights_only=False) # our own checkpoint
	model.load_state_dict(ck["model"])
	opt.load_state_dict(ck["optim"])
	args.start_step, best = ck["step"], ck.get("best", -math.inf)
	if "scaler" in ck and scaler.is_enabled():
	scaler.load_state_dict(ck["scaler"])
	print(f"resumed from step {args.start_step}")
	elif args.init:
	init_from(model, Path(args.init) / args.component / "best.pt")
	if args.start_step >= args.steps:
	print("already finished")
	return
	print(f"{args.component}: {count_params(model):,} params, device={device}, workers={args.workers}, "
	f"steps={args.steps}, batch={args.batch}, lr={args.lr}, amp={use_amp}")

	def lr_at(s):
	if s < args.warmup:
	return args.lr * (s + 1) / args.warmup
	t = (s - args.warmup) / max(1, args.steps - args.warmup)
	return args.lr * (0.05 + 0.95 * 0.5 * (1 + math.cos(math.pi * t)))

	log = open(out / "log.jsonl", "a")
	loader = make_loader(args.component, args)
	model.train()
	step, t0, running = args.start_step, time.time(), []
	for batch in loader:
	for g in opt.param_groups:
	g["lr"] = lr_at(step)
	with torch.autocast(device.type, dtype=torch.float16, enabled=use_amp):
	loss = step_loss(args.component, model, batch, device)
	opt.zero_grad(set_to_none=True)
	scaler.scale(loss).backward()
	scaler.unscale_(opt)
	torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
	scaler.step(opt)
	scaler.update()
	running.append(loss.item())
	step += 1
	if step % 50 == 0:
	dt = time.time() - t0
	print(f"step {step:6d} loss {np.mean(running):.4f} lr {lr_at(step):.2e} {50 * args.batch / dt:.0f} samples/s", flush=True)
	log.write(json.dumps({"step": step, "loss": float(np.mean(running)), "lr": lr_at(step)}) + "\n")
	running, t0 = [], time.time()
	out_of_time = deadline is not None and time.time() > deadline
	if step % args.eval_every == 0 or step == args.steps or out_of_time:
	metrics, score = EVALS[args.component](model, device, args)
	is_best = score > best
	best = max(best, score)
	ck = {"model": model.state_dict(), "optim": opt.state_dict(), "step": step, "best": best,
	"component": args.component, "metrics": metrics, "scaler": scaler.state_dict(), "arch": arch}
	torch.save(ck, out / "last.pt")
	if is_best:
	torch.save({k: v for k, v in ck.items() if k != "optim"}, out / "best.pt")
	print(f"eval @ {step}: " + " ".join(f"{k}={v:.4f}" for k, v in sorted(metrics.items())) + (" best" if is_best else ""), flush=True)
	log.write(json.dumps({"step": step, "eval": metrics, "best": is_best}) + "\n")
	log.flush()
	t0 = time.time()
	if out_of_time:
	print(f"time budget reached at step {step}; resume with --resume", flush=True)
	break
	if step >= args.steps:
	break


	if __name__ == "__main__":
	main()