typedecide.training.trainer

finetune, TrainResult and the manifest writer. Needs the train extra at call time, not at import time.

typedecide.training.trainer

LoRA fine-tuning on the answer token, and the manifest that makes the run traceable.

torch, transformers and peft are optional extras and are imported inside the functions that need them, so import typedecide stays fast and a base install can still load, validate and split data. Anything missing raises TrainingError naming the extra.

Nothing here re-implements the prompt or the masking: both come from prompt.py and dataset.py. What this module owns is the surrounding run -- seeding, the adapter, the schedule, and manifest.json.

MANIFEST_NAME module-attribute

MANIFEST_NAME = 'manifest.json'

MANIFEST_VERSION module-attribute

MANIFEST_VERSION = 1

REDACTED module-attribute

REDACTED = '[redacted]'

TrainResult dataclass

TrainResult(
    adapter_dir: Path,
    manifest: dict[str, Any],
    train_loss: float,
    eval_loss: float | None,
)

Where the adapter landed, what produced it, and how it went.

adapter_dir instance-attribute

adapter_dir: Path

manifest instance-attribute

manifest: dict[str, Any]

train_loss instance-attribute

train_loss: float

eval_loss instance-attribute

eval_loss: float | None

library_versions

library_versions() -> dict[str, str]

Versions of everything that can change a weight, for the manifest.

Reported as "not installed" rather than omitted, because a manifest that is silent about a library cannot be told apart from one written before it mattered.

Source code in src/typedecide/training/trainer.py

def library_versions() -> dict[str, str]:
    """Versions of everything that can change a weight, for the manifest.

    Reported as "not installed" rather than omitted, because a manifest that is silent
    about a library cannot be told apart from one written before it mattered.
    """
    versions: dict[str, str] = {
        "python": platform.python_version(),
        "platform": platform.platform(),
    }
    for name in ("torch", "transformers", "peft", "accelerate", "numpy"):
        try:
            module = __import__(name)
        except ImportError:
            versions[name] = "not installed"
        else:
            versions[name] = str(getattr(module, "__version__", "unknown"))
    try:
        from importlib.metadata import version

        versions["typedecide"] = version("typedecide")
    except Exception:  # noqa: BLE001 - an uninstalled source checkout is normal here
        versions["typedecide"] = "source"
    return versions

seed_everything

seed_everything(seed: int) -> None

Seed python, numpy, torch and transformers from the one number in the config.

Called before the model is built, because LoRA's B matrices and the dataloader shuffle are both drawn from these generators.

Source code in src/typedecide/training/trainer.py

def seed_everything(seed: int) -> None:
    """Seed python, numpy, torch and transformers from the one number in the config.

    Called before the model is built, because LoRA's B matrices and the dataloader
    shuffle are both drawn from these generators.
    """
    random.seed(seed)
    try:
        import numpy

        numpy.random.seed(seed % (2**32))
    except ImportError:  # pragma: no cover - numpy is a base dependency
        log.warning("numpy is not installed, so its generator was not seeded.")
    try:
        import torch
        from transformers import set_seed
    except ImportError:
        log.debug("torch and transformers are absent; seeded python and numpy only.")
        return
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    set_seed(seed)

build_manifest

build_manifest(
    config: TrainConfig,
    *,
    train: Sequence[Decision],
    eval_set: Sequence[Decision] | None,
    answer_slots: dict[str, dict[str, int]],
    started_at: float,
    finished_at: float,
    extra: dict[str, Any] | None = None
) -> dict[str, Any]

Everything needed to reproduce, or to disbelieve, this run.

A result you cannot trace to its input is not a result, so this records the base model, the fully resolved config, the dataset fingerprints and row counts, the seed, the library versions, wall-clock start and end, and the resolved answer-slot token ids per criterion.

Source code in src/typedecide/training/trainer.py

def build_manifest(
    config: TrainConfig,
    *,
    train: Sequence[Decision],
    eval_set: Sequence[Decision] | None,
    answer_slots: dict[str, dict[str, int]],
    started_at: float,
    finished_at: float,
    extra: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Everything needed to reproduce, or to disbelieve, this run.

    A result you cannot trace to its input is not a result, so this records the base
    model, the fully resolved config, the dataset fingerprints and row counts, the
    seed, the library versions, wall-clock start and end, and the resolved answer-slot
    token ids per criterion.
    """
    manifest: dict[str, Any] = {
        "manifest_version": MANIFEST_VERSION,
        "base_model": config.base_model,
        "config": config.to_manifest(),
        "seed": config.seed,
        "dataset": {
            "train": {
                "fingerprint": fingerprint(train),
                "rows": len(train),
                "criteria": sorted({d.criterion.key for d in train}),
                "groups": len({d.grouping for d in train}),
            },
            "eval": None if eval_set is None else {
                "fingerprint": fingerprint(eval_set),
                "rows": len(eval_set),
                "criteria": sorted({d.criterion.key for d in eval_set}),
                "groups": len({d.grouping for d in eval_set}),
            },
        },
        "answer_slots": answer_slots,
        "option_order_randomisation": (
            "per-epoch" if config.randomise_option_order else "off"
        ),
        "timing": {
            "started_at": _iso(started_at),
            "finished_at": _iso(finished_at),
            "wall_seconds": round(finished_at - started_at, 3),
        },
        "versions": library_versions(),
        "argv": redacted_argv(sys.argv),
    }
    if extra:
        manifest.update(extra)
    return manifest

redacted_argv

redacted_argv(argv: Sequence[str]) -> list[str]

argv with anything that looks like a credential blanked out.

The manifest is written beside the adapter and travels with it -- to a model hub, a ticket, a colleague. finetune is a library call, so sys.argv belongs to whatever program imported us and may carry --hf-token hf_.... The command line is worth recording; the secret on it is not.

Source code in src/typedecide/training/trainer.py

def redacted_argv(argv: Sequence[str]) -> list[str]:
    """`argv` with anything that looks like a credential blanked out.

    The manifest is written beside the adapter and travels with it -- to a model hub, a
    ticket, a colleague. `finetune` is a library call, so `sys.argv` belongs to whatever
    program imported us and may carry `--hf-token hf_...`. The command line is worth
    recording; the secret on it is not.
    """
    out: list[str] = []
    hide_next = False
    for argument in argv:
        text = str(argument)
        if hide_next:
            out.append(REDACTED)
            hide_next = False
            continue
        flag, equals, value = text.partition("=")
        if flag.startswith("-") and _SECRET_FLAG.search(flag):
            if equals:
                out.append(f"{flag}={REDACTED}")
            else:
                out.append(text)
                hide_next = True
            continue
        if _SECRET_VALUE.match(text) or (equals and _SECRET_VALUE.match(value)):
            out.append(REDACTED if not equals else f"{flag}={REDACTED}")
            continue
        out.append(text)
    return out

write_manifest

write_manifest(
    manifest: dict[str, Any], output_dir: Path
) -> Path

Write manifest.json beside the adapter and return its path.

Source code in src/typedecide/training/trainer.py

def write_manifest(manifest: dict[str, Any], output_dir: Path) -> Path:
    """Write `manifest.json` beside the adapter and return its path."""
    output_dir.mkdir(parents=True, exist_ok=True)
    path = output_dir / MANIFEST_NAME
    path.write_text(json.dumps(manifest, indent=2, sort_keys=False) + "\n", encoding="utf-8")
    log.info("Manifest written to %s", path)
    return path

finetune

finetune(
    train: Sequence[Decision],
    config: TrainConfig,
    *,
    eval_set: Sequence[Decision] | None = None,
    progress: Callable[[dict[str, Any]], None] | None = None
) -> TrainResult

LoRA-tune config.base_model so that one token carries all the loss.

Every example is encode(prompt) + [answer_token] with every position but the last masked to -100; see dataset.build_example. Option order is re-randomised at the start of each epoch when config.randomise_option_order is set, which is the single most valuable preprocessing step here: without it the model learns "answer A".

progress, if given, receives dicts -- {"event": "preflight" | "log" | "epoch" | "done", ...} -- so a CLI or notebook can show something without this module owning a progress bar.

Source code in src/typedecide/training/trainer.py

def finetune(
    train: Sequence[Decision],
    config: TrainConfig,
    *,
    eval_set: Sequence[Decision] | None = None,
    progress: Callable[[dict[str, Any]], None] | None = None,
) -> TrainResult:
    """LoRA-tune `config.base_model` so that one token carries all the loss.

    Every example is `encode(prompt) + [answer_token]` with every position but the last
    masked to -100; see `dataset.build_example`. Option order is re-randomised at the
    start of each epoch when `config.randomise_option_order` is set, which is the single
    most valuable preprocessing step here: without it the model learns "answer A".

    `progress`, if given, receives dicts -- `{"event": "preflight" | "log" |
    "epoch" | "done", ...}` -- so a CLI or notebook can show something without this
    module owning a progress bar.
    """
    modules = _require_training_deps()
    torch = modules["torch"]

    started_at = time.time()
    labelled, unlabelled = labelled_only(train)
    if unlabelled:
        raise TrainingError(
            f"{len(unlabelled)} of {len(train)} training decisions have no answer_id, "
            f"starting with {', '.join(unlabelled[:5])}. Label them or filter them out; "
            "training silently on a smaller set than you passed is worse than stopping."
        )
    if not labelled:
        raise TrainingError("There are no labelled decisions to train on.")

    labelled_eval: list[Decision] | None = None
    if eval_set is not None:
        labelled_eval, unlabelled_eval = labelled_only(eval_set)
        if unlabelled_eval:
            log.warning(
                "Dropping %d unlabelled eval decisions (%s%s).",
                len(unlabelled_eval), ", ".join(unlabelled_eval[:5]),
                "..." if len(unlabelled_eval) > 5 else "",
            )
        if not labelled_eval:
            raise TrainingError("The eval set has no labelled decisions to score.")

    seed_everything(config.seed)
    _emit(progress, {"event": "start", "train_rows": len(labelled),
                     "eval_rows": 0 if labelled_eval is None else len(labelled_eval)})

    tokenizer = _load_tokenizer(config.base_model)
    train_data = DecisionDataset(
        labelled, tokenizer, config.max_length,
        randomise_option_order=config.randomise_option_order, seed=config.seed,
    )
    # The eval set is never re-permuted: a moving eval set is not a measurement.
    eval_data = None if labelled_eval is None else DecisionDataset(
        labelled_eval, tokenizer, config.max_length,
        randomise_option_order=False, seed=config.seed,
    )

    answer_slots = train_data.preflight()
    if eval_data is not None:
        eval_data.preflight()
    _emit(progress, {"event": "preflight", "criteria": sorted(answer_slots)})
    log.info(
        "%d training decisions, %s held out, %d distinct criteria.",
        len(train_data), "none" if eval_data is None else str(len(eval_data)), len(answer_slots),
    )

    model = _build_peft_model(config, torch, modules["transformers"], modules["peft"])
    trainer = _build_trainer(
        config, model, tokenizer, train_data, eval_data, torch,
        modules["transformers"], progress,
    )

    train_output = trainer.train()
    train_loss = float(train_output.training_loss)

    eval_loss: float | None = None
    if eval_data is not None:
        metrics = trainer.evaluate()
        raw = metrics.get("eval_loss")
        eval_loss = None if raw is None else float(raw)

    config.output_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(config.output_dir))
    tokenizer.save_pretrained(str(config.output_dir))

    manifest = build_manifest(
        config,
        train=labelled,
        eval_set=labelled_eval,
        answer_slots=answer_slots,
        started_at=started_at,
        finished_at=time.time(),
        extra={
            "losses": {"train": train_loss, "eval": eval_loss},
            "steps": int(getattr(train_output, "global_step", 0) or 0),
            "trainable_parameters": _trainable_parameters(model),
        },
    )
    write_manifest(manifest, config.output_dir)
    _emit(progress, {"event": "done", "adapter_dir": str(config.output_dir),
                     "train_loss": train_loss, "eval_loss": eval_loss})
    log.info("Adapter written to %s (train loss %.4f)", config.output_dir, train_loss)

    return TrainResult(
        adapter_dir=config.output_dir,
        manifest=manifest,
        train_loss=train_loss,
        eval_loss=eval_loss,
    )