scripts/train_t5_rewrite.py

from __future__ import annotations

import argparse
import json
import os
import time
import pickle
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Set

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    DataCollatorForSeq2Seq,
    Seq2SeqTrainer,
    TrainerCallback,
    Seq2SeqTrainingArguments,
)
from transformers.trainer_utils import get_last_checkpoint

REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_SPLIT_DIR = REPO_ROOT / "data" / "external" / "caption_emporium" / "t5_rewrite_splits"
DEFAULT_BASE_MODEL = REPO_ROOT / "models" / "t5-small"
DEFAULT_OUT_DIR = REPO_ROOT / "models" / "finetune" / "t5-rewrite"


class TokenizedListDataset(Dataset):
    def __init__(self, records: List[Dict[str, List[int]]]):
        self.records = records

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, idx: int) -> Dict[str, List[int]]:
        return self.records[idx]


def _canon_tag(tag: str) -> str:
    t = " ".join(str(tag or "").strip().split()).lower()
    return t.replace(" ", "_").replace("\\(", "(").replace("\\)", ")")


def _parse_tag_set(text: str) -> Set[str]:
    out: Set[str] = set()
    for raw in (text or "").split(","):
        t = _canon_tag(raw)
        if t:
            out.add(t)
    return out


def _set_metrics(pred_texts: List[str], gold_texts: List[str]) -> Dict[str, float]:
    if not pred_texts:
        return {
            "set_precision": 0.0,
            "set_recall": 0.0,
            "set_f1": 0.0,
            "exact_set_match": 0.0,
            "avg_pred_tags": 0.0,
            "avg_gold_tags": 0.0,
        }

    p_vals: List[float] = []
    r_vals: List[float] = []
    f_vals: List[float] = []
    exact = 0
    pred_sizes: List[int] = []
    gold_sizes: List[int] = []

    for ptxt, gtxt in zip(pred_texts, gold_texts):
        pset = _parse_tag_set(ptxt)
        gset = _parse_tag_set(gtxt)
        pred_sizes.append(len(pset))
        gold_sizes.append(len(gset))
        if pset == gset:
            exact += 1
        if not pset and not gset:
            p_vals.append(1.0)
            r_vals.append(1.0)
            f_vals.append(1.0)
            continue
        if not pset or not gset:
            p_vals.append(0.0)
            r_vals.append(0.0)
            f_vals.append(0.0)
            continue
        tp = len(pset & gset)
        p = tp / len(pset)
        r = tp / len(gset)
        f = (2 * p * r / (p + r)) if (p + r) > 0 else 0.0
        p_vals.append(p)
        r_vals.append(r)
        f_vals.append(f)

    n = len(pred_texts)
    return {
        "set_precision": float(np.mean(p_vals)),
        "set_recall": float(np.mean(r_vals)),
        "set_f1": float(np.mean(f_vals)),
        "exact_set_match": exact / n,
        "avg_pred_tags": float(np.mean(pred_sizes)),
        "avg_gold_tags": float(np.mean(gold_sizes)),
    }


class ProgressFileCallback(TrainerCallback):
    def __init__(self, progress_path: Path, history_path: Optional[Path] = None):
        self.progress_path = progress_path
        self.history_path = history_path
        self.start_time: Optional[float] = None

    def _write(self, payload: Dict[str, object]) -> None:
        self.progress_path.parent.mkdir(parents=True, exist_ok=True)
        with self.progress_path.open("w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False, indent=2)

    def _append_history(self, payload: Dict[str, object]) -> None:
        if self.history_path is None:
            return
        self.history_path.parent.mkdir(parents=True, exist_ok=True)
        with self.history_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(payload, ensure_ascii=False) + "\n")

    def _pct_eta(self, global_step: int, max_steps: int) -> Dict[str, Optional[float]]:
        max_steps = max(0, int(max_steps))
        global_step = max(0, int(global_step))
        pct = (100.0 * global_step / max_steps) if max_steps > 0 else None
        eta = None
        elapsed = None
        if self.start_time is not None:
            elapsed = time.time() - self.start_time
            if max_steps > 0 and global_step > 0:
                eta = (elapsed / global_step) * (max_steps - global_step)
        return {"pct": pct, "eta_sec": eta, "elapsed_sec": elapsed}

    def on_train_begin(self, args, state, control, **kwargs):
        self.start_time = time.time()
        info = self._pct_eta(state.global_step, state.max_steps)
        self._write(
            {
                "status": "running",
                "global_step": int(state.global_step),
                "max_steps": int(state.max_steps),
                "pct_complete": info["pct"],
                "elapsed_sec": info["elapsed_sec"],
                "eta_sec": info["eta_sec"],
                "last_log": {},
                "updated_at_epoch_sec": time.time(),
            }
        )

    def on_log(self, args, state, control, logs=None, **kwargs):
        info = self._pct_eta(state.global_step, state.max_steps)
        payload = {
            "status": "running",
            "global_step": int(state.global_step),
            "max_steps": int(state.max_steps),
            "pct_complete": info["pct"],
            "elapsed_sec": info["elapsed_sec"],
            "eta_sec": info["eta_sec"],
            "last_log": logs or {},
            "updated_at_epoch_sec": time.time(),
        }
        self._write(payload)
        pct_text = f"{info['pct']:.1f}%" if info["pct"] is not None else "n/a"
        eta_text = f"{info['eta_sec']:.0f}s" if info["eta_sec"] is not None else "n/a"
        print(f"[train] step {state.global_step}/{state.max_steps} ({pct_text}) eta={eta_text} logs={logs or {}}")

    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        row = {
            "event": "evaluate",
            "global_step": int(state.global_step),
            "max_steps": int(state.max_steps),
            "metrics": metrics or {},
            "updated_at_epoch_sec": time.time(),
        }
        self._append_history(row)

    def on_train_end(self, args, state, control, **kwargs):
        info = self._pct_eta(state.global_step, state.max_steps)
        self._write(
            {
                "status": "completed",
                "global_step": int(state.global_step),
                "max_steps": int(state.max_steps),
                "pct_complete": info["pct"],
                "elapsed_sec": info["elapsed_sec"],
                "eta_sec": 0.0,
                "last_log": {},
                "updated_at_epoch_sec": time.time(),
            }
        )


class PeriodicTestEvalCallback(TrainerCallback):
    def __init__(self, test_dataset: Dataset, every_steps: int):
        self.test_dataset = test_dataset
        self.every_steps = max(0, int(every_steps))
        self._trainer: Optional[Seq2SeqTrainer] = None
        self._in_test_eval = False

    def bind_trainer(self, trainer: Seq2SeqTrainer) -> None:
        self._trainer = trainer

    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        if self.every_steps <= 0 or self._trainer is None or self._in_test_eval:
            return
        if state.global_step <= 0 or (state.global_step % self.every_steps) != 0:
            return
        if isinstance(metrics, dict) and any(k.startswith("test_") for k in metrics.keys()):
            return

        self._in_test_eval = True
        try:
            m = self._trainer.evaluate(eval_dataset=self.test_dataset, metric_key_prefix="test")
            test_recall = m.get("test_set_recall")
            if test_recall is None:
                print(f"[test-eval] step={state.global_step} ran periodic held-out test evaluation")
            else:
                print(f"[test-eval] step={state.global_step} test_set_recall={float(test_recall):.4f}")
        finally:
            self._in_test_eval = False


class RecallWeightedSeq2SeqTrainer(Seq2SeqTrainer):
    def __init__(
        self,
        *args,
        eos_token_id: int,
        comma_token_ids: Sequence[int],
        eos_loss_weight: float,
        comma_loss_weight: float,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.eos_token_id = int(eos_token_id)
        self.comma_token_ids = [int(x) for x in comma_token_ids]
        self.eos_loss_weight = float(eos_loss_weight)
        self.comma_loss_weight = float(comma_loss_weight)
        self.use_weighted_loss = (self.eos_loss_weight != 1.0) or (self.comma_loss_weight != 1.0)

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        if not self.use_weighted_loss:
            return super().compute_loss(
                model,
                inputs,
                return_outputs=return_outputs,
                num_items_in_batch=num_items_in_batch,
            )

        labels = inputs.get("labels")
        if labels is None:
            return super().compute_loss(
                model,
                inputs,
                return_outputs=return_outputs,
                num_items_in_batch=num_items_in_batch,
            )

        outputs = model(**inputs)
        logits = outputs.get("logits") if isinstance(outputs, dict) else outputs.logits
        vocab = logits.size(-1)

        token_loss = F.cross_entropy(
            logits.view(-1, vocab),
            labels.view(-1),
            ignore_index=-100,
            reduction="none",
        ).view_as(labels)

        valid = (labels != -100).to(logits.dtype)
        weights = torch.ones_like(labels, dtype=logits.dtype)
        if self.eos_loss_weight != 1.0:
            weights = torch.where(labels == self.eos_token_id, torch.full_like(weights, self.eos_loss_weight), weights)
        if self.comma_loss_weight != 1.0 and self.comma_token_ids:
            for cid in self.comma_token_ids:
                weights = torch.where(labels == cid, torch.full_like(weights, self.comma_loss_weight), weights)

        denom = (weights * valid).sum().clamp(min=1.0)
        loss = (token_loss * weights * valid).sum() / denom
        return (loss, outputs) if return_outputs else loss


def _read_jsonl(path: Path) -> List[Dict[str, str]]:
    rows: List[Dict[str, str]] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            src = str(obj.get("source_text", "")).strip()
            tgt = str(obj.get("target_text", "")).strip()
            if not src or not tgt:
                continue
            rows.append({"source_text": src, "target_text": tgt})
    return rows


def _cap(rows: List[Dict[str, str]], n: int) -> List[Dict[str, str]]:
    if n <= 0:
        return rows
    return rows[: min(n, len(rows))]


def _tokenize_rows(
    rows: Sequence[Dict[str, str]],
    tokenizer,
    source_max_len: int,
    target_max_len: int,
) -> TokenizedListDataset:
    srcs = [r["source_text"] for r in rows]
    tgts = [r["target_text"] for r in rows]
    src_tok = tokenizer(srcs, truncation=True, max_length=source_max_len)
    tgt_tok = tokenizer(text_target=tgts, truncation=True, max_length=target_max_len)

    recs: List[Dict[str, List[int]]] = []
    for i in range(len(rows)):
        recs.append(
            {
                "input_ids": src_tok["input_ids"][i],
                "attention_mask": src_tok["attention_mask"][i],
                "labels": tgt_tok["input_ids"][i],
            }
        )
    return TokenizedListDataset(recs)


def _maybe_disable_bad_rng_state(checkpoint_dir: Optional[str]) -> Optional[str]:
    if not checkpoint_dir:
        return checkpoint_dir
    rng_path = Path(checkpoint_dir) / "rng_state.pth"
    if not rng_path.is_file():
        return checkpoint_dir
    try:
        torch.load(str(rng_path))
        return checkpoint_dir
    except pickle.UnpicklingError:
        # Newer torch defaults to weights_only=True and can reject older rng blobs.
        # If RNG state is unreadable, training can still resume from weights/optimizer;
        # we just skip restoring exact RNG stream.
        bad_path = rng_path.with_suffix(".pth.unusable")
        try:
            if bad_path.exists():
                bad_path.unlink()
            rng_path.replace(bad_path)
            print(f"Disabled unreadable RNG state file: {rng_path} -> {bad_path}")
        except Exception as ex:
            print(f"Warning: could not move unreadable RNG state file {rng_path}: {ex}")
        return checkpoint_dir


def main() -> int:
    ap = argparse.ArgumentParser(description="Fine-tune local T5 for caption -> comma-separated canonical tags")
    ap.add_argument("--split-dir", type=Path, default=DEFAULT_SPLIT_DIR)
    ap.add_argument("--base-model-dir", type=Path, default=DEFAULT_BASE_MODEL)
    ap.add_argument("--output-dir", type=Path, default=DEFAULT_OUT_DIR)
    ap.add_argument("--source-max-len", type=int, default=160)
    ap.add_argument("--target-max-len", type=int, default=256)
    ap.add_argument("--num-beams", type=int, default=4)
    ap.add_argument("--generation-length-penalty", type=float, default=0.8,
                    help="<1 encourages longer outputs (recall-leaning), >1 encourages shorter outputs")
    ap.add_argument("--lr", type=float, default=2e-4)
    ap.add_argument("--weight-decay", type=float, default=0.01)
    ap.add_argument("--warmup-ratio", type=float, default=0.05)
    ap.add_argument("--label-smoothing", type=float, default=0.1)
    ap.add_argument("--eos-loss-weight", type=float, default=1.0,
                    help="When <1, penalize EOS-token mismatch less to reduce over-short outputs")
    ap.add_argument("--comma-loss-weight", type=float, default=1.0,
                    help="When <1, penalize comma-token mismatch less to focus loss on tag tokens")
    ap.add_argument("--epochs", type=float, default=1.0)
    ap.add_argument("--max-steps", type=int, default=3000, help="<=0 uses full epoch schedule")
    ap.add_argument("--train-batch-size", type=int, default=2)
    ap.add_argument("--eval-batch-size", type=int, default=2)
    ap.add_argument("--grad-accum", type=int, default=8)
    ap.add_argument("--logging-steps", type=int, default=25)
    ap.add_argument("--eval-steps", type=int, default=500, help="<=0 evaluates each epoch")
    ap.add_argument("--save-steps", type=int, default=250, help="<=0 saves each epoch")
    ap.add_argument("--max-train-samples", type=int, default=0, help="Cap train samples after loading (0 disables)")
    ap.add_argument("--max-val-samples", type=int, default=300, help="Cap validation samples for eval (0 disables)")
    ap.add_argument("--max-test-samples", type=int, default=300, help="Cap test samples for eval (0 disables)")
    ap.add_argument("--eval-during-train", action="store_true", default=False,
                    help="Enable periodic evaluation/checkpoint selection during training")
    ap.add_argument("--periodic-test-eval", action="store_true", default=False,
                    help="When eval-during-train is enabled, evaluate both val and test each eval pass")
    ap.add_argument("--test-eval-every-steps", type=int, default=0,
                    help="If >0, run held-out test eval every N global steps (after val eval)")
    ap.add_argument("--save-total-limit", type=int, default=3,
                    help="Max number of checkpoints kept on disk")
    ap.add_argument("--best-model-metric", type=str, default="recall",
                    choices=["recall", "f1", "precision", "loss"],
                    help="Metric used to select best checkpoint when load_best_model_at_end is active")
    ap.add_argument("--require-cuda", action="store_true", default=False,
                    help="Fail immediately if CUDA is not available")
    ap.add_argument("--progress-file", type=Path,
                    default=REPO_ROOT / "data" / "runtime_metrics" / "t5_rewrite_train_progress.json",
                    help="JSON file updated at each logging step with percent/ETA/progress")
    ap.add_argument("--progress-history-file", type=Path,
                    default=REPO_ROOT / "data" / "runtime_metrics" / "t5_rewrite_train_progress_history.jsonl",
                    help="JSONL history file for periodic evaluation events")
    ap.add_argument("--fp16", action="store_true", default=False, help="Enable fp16 mixed precision training")
    ap.add_argument("--bf16", action="store_true", default=False, help="Enable bf16 mixed precision training")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--report-to", type=str, default="none", help="none or tensorboard/wandb")
    ap.add_argument(
        "--resume-if-available",
        action="store_true",
        default=False,
        help="Resume from latest checkpoint in output-dir when present",
    )
    args = ap.parse_args()

    split_dir = args.split_dir if args.split_dir.is_absolute() else (REPO_ROOT / args.split_dir).resolve()
    model_dir = args.base_model_dir if args.base_model_dir.is_absolute() else (REPO_ROOT / args.base_model_dir).resolve()
    out_dir = args.output_dir if args.output_dir.is_absolute() else (REPO_ROOT / args.output_dir).resolve()
    progress_path = args.progress_file if args.progress_file.is_absolute() else (REPO_ROOT / args.progress_file).resolve()
    progress_history_path = (
        args.progress_history_file
        if args.progress_history_file.is_absolute()
        else (REPO_ROOT / args.progress_history_file).resolve()
    )
    out_dir.mkdir(parents=True, exist_ok=True)

    cuda_available = torch.cuda.is_available()
    cuda_name = torch.cuda.get_device_name(0) if cuda_available else ""
    if args.require_cuda and not cuda_available:
        raise RuntimeError(
            "CUDA is required for this run but not available. "
            "Use a CUDA-enabled PyTorch environment (GPU wheel) and retry."
        )
    print(
        f"torch={torch.__version__} cuda_available={cuda_available}"
        + (f" device='{cuda_name}'" if cuda_name else "")
    )
    if args.eos_loss_weight != 1.0 or args.comma_loss_weight != 1.0:
        if args.label_smoothing != 0.0:
            raise ValueError(
                "Weighted loss is enabled via eos/comma loss weights, but label smoothing is non-zero. "
                "Set --label-smoothing 0 when using weighted loss."
            )

    def _write_stage_status(status: str, extra: Optional[Dict[str, object]] = None) -> None:
        payload: Dict[str, object] = {
            "status": status,
            "updated_at_epoch_sec": time.time(),
        }
        if extra:
            payload.update(extra)
        progress_path.parent.mkdir(parents=True, exist_ok=True)
        with progress_path.open("w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False, indent=2)

    train_path = split_dir / "train.jsonl"
    val_path = split_dir / "val.jsonl"
    test_path = split_dir / "test.jsonl"
    for p in (train_path, val_path, test_path):
        if not p.is_file():
            raise FileNotFoundError(f"Missing split file: {p}")
    if not model_dir.is_dir():
        raise FileNotFoundError(f"Missing base model dir: {model_dir}")

    _write_stage_status("loading_dataset")
    train_rows = _cap(_read_jsonl(train_path), args.max_train_samples)
    val_rows = _cap(_read_jsonl(val_path), args.max_val_samples)
    test_rows = _cap(_read_jsonl(test_path), args.max_test_samples)
    print("dataset_rows:", {"train": len(train_rows), "validation": len(val_rows), "test": len(test_rows)})

    _write_stage_status(
        "loading_model",
        {
            "dataset_rows": {
                "train": len(train_rows),
                "validation": len(val_rows),
                "test": len(test_rows),
            }
        },
    )
    tokenizer = AutoTokenizer.from_pretrained(str(model_dir), local_files_only=True, use_fast=False)
    model = AutoModelForSeq2SeqLM.from_pretrained(str(model_dir), local_files_only=True)
    model.generation_config.length_penalty = float(args.generation_length_penalty)
    model.generation_config.num_beams = max(1, int(args.num_beams))
    _write_stage_status("tokenizing")
    train_ds = _tokenize_rows(train_rows, tokenizer, args.source_max_len, args.target_max_len)
    val_ds = _tokenize_rows(val_rows, tokenizer, args.source_max_len, args.target_max_len)
    test_ds = _tokenize_rows(test_rows, tokenizer, args.source_max_len, args.target_max_len)

    collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)

    report_to = [] if args.report_to == "none" else [args.report_to]
    if args.eval_during_train:
        eval_strategy = "steps" if args.eval_steps > 0 else "epoch"
        save_strategy = "steps" if args.save_steps > 0 else "epoch"
        load_best = not (args.periodic_test_eval and args.test_eval_every_steps <= 0)
    else:
        eval_strategy = "no"
        save_strategy = "no"
        load_best = False
    metric_for_best_map = {
        "recall": "eval_set_recall",
        "f1": "eval_set_f1",
        "precision": "eval_set_precision",
        "loss": "eval_loss",
    }
    metric_for_best_model = metric_for_best_map[args.best_model_metric]
    greater_is_better = args.best_model_metric != "loss"

    targs = Seq2SeqTrainingArguments(
        output_dir=str(out_dir),
        learning_rate=args.lr,
        weight_decay=args.weight_decay,
        warmup_ratio=args.warmup_ratio,
        label_smoothing_factor=args.label_smoothing,
        per_device_train_batch_size=args.train_batch_size,
        per_device_eval_batch_size=args.eval_batch_size,
        gradient_accumulation_steps=args.grad_accum,
        predict_with_generate=True,
        generation_num_beams=args.num_beams,
        generation_max_length=args.target_max_len,
        num_train_epochs=args.epochs,
        max_steps=args.max_steps if args.max_steps > 0 else -1,
        evaluation_strategy=eval_strategy,
        eval_steps=args.eval_steps if (args.eval_during_train and args.eval_steps > 0) else None,
        save_strategy=save_strategy,
        save_steps=args.save_steps if (args.eval_during_train and args.save_steps > 0) else None,
        logging_steps=args.logging_steps,
        logging_strategy="steps",
        save_total_limit=args.save_total_limit,
        load_best_model_at_end=load_best,
        metric_for_best_model=metric_for_best_model if load_best else None,
        greater_is_better=greater_is_better if load_best else None,
        seed=args.seed,
        dataloader_num_workers=0,
        report_to=report_to,
        fp16=args.fp16,
        bf16=args.bf16,
    )

    def _compute_metrics(eval_pred):
        preds, labels = eval_pred
        if isinstance(preds, tuple):
            preds = preds[0]
        preds = np.asarray(preds)
        if preds.ndim == 3:
            preds = np.argmax(preds, axis=-1)
        preds = np.where(
            (preds < 0) | (preds >= tokenizer.vocab_size),
            tokenizer.pad_token_id,
            preds,
        )
        labels = np.asarray(labels)
        labels = np.where(
            (labels < 0) | (labels >= tokenizer.vocab_size),
            tokenizer.pad_token_id,
            labels,
        )
        pred_texts = tokenizer.batch_decode(preds.tolist(), skip_special_tokens=True)
        gold_texts = tokenizer.batch_decode(labels.tolist(), skip_special_tokens=True)
        return _set_metrics(pred_texts, gold_texts)

    eval_dataset_obj = (
        {"val": val_ds, "test": test_ds}
        if (args.eval_during_train and args.periodic_test_eval and args.test_eval_every_steps <= 0)
        else val_ds
    )

    comma_token_ids = tokenizer.encode(",", add_special_tokens=False)
    callbacks: List[TrainerCallback] = [ProgressFileCallback(progress_path, history_path=progress_history_path)]
    periodic_test_cb: Optional[PeriodicTestEvalCallback] = None
    if args.eval_during_train and args.test_eval_every_steps > 0:
        periodic_test_cb = PeriodicTestEvalCallback(test_dataset=test_ds, every_steps=args.test_eval_every_steps)
        callbacks.append(periodic_test_cb)

    trainer = RecallWeightedSeq2SeqTrainer(
        model=model,
        args=targs,
        train_dataset=train_ds,
        eval_dataset=eval_dataset_obj,
        tokenizer=tokenizer,
        data_collator=collator,
        compute_metrics=_compute_metrics,
        callbacks=callbacks,
        eos_token_id=tokenizer.eos_token_id,
        comma_token_ids=comma_token_ids,
        eos_loss_weight=args.eos_loss_weight,
        comma_loss_weight=args.comma_loss_weight,
    )
    if periodic_test_cb is not None:
        periodic_test_cb.bind_trainer(trainer)

    resume_checkpoint = None
    if args.resume_if_available:
        resume_checkpoint = get_last_checkpoint(str(out_dir))
        if resume_checkpoint:
            print(f"Resuming from checkpoint: {resume_checkpoint}")
            resume_checkpoint = _maybe_disable_bad_rng_state(resume_checkpoint)

    _write_stage_status(
        "training_starting",
        {
            "dataset_rows_tokenized": {
                "train": len(train_ds),
                "validation": len(val_ds),
                "test": len(test_ds),
            },
            "max_steps": int(targs.max_steps),
            "num_train_epochs": float(targs.num_train_epochs),
            "resume_from_checkpoint": resume_checkpoint,
        },
    )
    train_result = trainer.train(resume_from_checkpoint=resume_checkpoint)
    trainer.save_model(str(out_dir))
    tokenizer.save_pretrained(str(out_dir))

    val_metrics = trainer.evaluate(eval_dataset=val_ds)
    test_metrics = trainer.evaluate(eval_dataset=test_ds, metric_key_prefix="test")

    metrics = {
        "train": train_result.metrics,
        "val": val_metrics,
        "test": test_metrics,
        "config": {
            "split_dir": str(split_dir),
            "base_model_dir": str(model_dir),
            "output_dir": str(out_dir),
            "source_max_len": args.source_max_len,
            "target_max_len": args.target_max_len,
            "num_beams": args.num_beams,
            "generation_length_penalty": args.generation_length_penalty,
            "learning_rate": args.lr,
            "weight_decay": args.weight_decay,
            "warmup_ratio": args.warmup_ratio,
            "label_smoothing": args.label_smoothing,
            "eos_loss_weight": args.eos_loss_weight,
            "comma_loss_weight": args.comma_loss_weight,
            "epochs": args.epochs,
            "max_steps": args.max_steps,
            "train_batch_size": args.train_batch_size,
            "eval_batch_size": args.eval_batch_size,
            "grad_accum": args.grad_accum,
            "seed": args.seed,
            "max_train_samples": args.max_train_samples,
            "max_val_samples": args.max_val_samples,
            "max_test_samples": args.max_test_samples,
            "eval_during_train": args.eval_during_train,
            "periodic_test_eval": args.periodic_test_eval,
            "test_eval_every_steps": args.test_eval_every_steps,
            "save_total_limit": args.save_total_limit,
            "best_model_metric": args.best_model_metric,
            "require_cuda": args.require_cuda,
            "progress_file": str(progress_path),
            "progress_history_file": str(progress_history_path),
            "cuda_available": cuda_available,
            "cuda_device": cuda_name,
            "fp16": args.fp16,
            "bf16": args.bf16,
            "resume_if_available": args.resume_if_available,
        },
    }
    _write_stage_status(
        "completed",
        {
            "global_step": int(train_result.metrics.get("global_step", targs.max_steps)),
            "max_steps": int(targs.max_steps),
            "pct_complete": 100.0,
            "elapsed_sec": train_result.metrics.get("train_runtime"),
            "eta_sec": 0.0,
            "last_log": {
                "train_loss": train_result.metrics.get("train_loss"),
                "eval_set_f1": val_metrics.get("eval_set_f1"),
                "test_set_f1": test_metrics.get("test_set_f1"),
            },
        },
    )
    with (out_dir / "train_metrics.json").open("w", encoding="utf-8") as f:
        json.dump(metrics, f, ensure_ascii=False, indent=2)

    print(json.dumps(metrics, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    os.chdir(REPO_ROOT)
    raise SystemExit(main())