sysone 0.2.0__py3-none-any.whl

sysone/questions.py

@@ -0,0 +1,186 @@
+"""Typed questions: validation, options per type, and training targets from labels.
+
+A question is a plain dict in the Jev `/v1/systemone` shape::
+
+    {"type": "choice", "instructions": "...", "criteria": ["a", "b"] or {"a": "desc", ...}}
+    {"type": "score",  "instructions": "...", "criteria": ["level 0", "level 1", ...]}
+    {"type": "noul",   "instructions": "...", "criteria": {"false": "...", "true": "..."}}  # criteria optional
+
+Every type is reduced to a list of options, so one decision head can serve all of them.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import dataclass
+from typing import Any
+
+QUESTION_TYPES = ("choice", "score", "noul")
+TYPE_INDEX = {name: i for i, name in enumerate(QUESTION_TYPES)}
+
+_TRUE_STRINGS = {"true", "yes", "1"}
+_FALSE_STRINGS = {"false", "no", "0"}
+
+
+@dataclass(frozen=True)
+class Question:
+    """A validated question.
+
+    `keys` are the option identifiers answers are reported with: the caller's choice labels, the
+    level indices for a score, and `False`/`True` for a noul. `texts` are the option descriptions
+    (`None` where none was given): a choice label's description, a score level's text, or a noul
+    criterion.
+    """
+
+    type: str
+    instructions: str
+    keys: tuple
+    texts: tuple
+
+    @property
+    def k(self) -> int:
+        return len(self.keys)
+
+    @property
+    def type_index(self) -> int:
+        return TYPE_INDEX[self.type]
+
+
+def _as_text(value: Any) -> str:
+    if isinstance(value, str):
+        return value
+    return json.dumps(value, ensure_ascii=False, default=str)
+
+
+def parse_question(qid: Any, qdef: Any) -> Question:
+    """Validate a question dict and return a `Question`. Errors name the question and the problem."""
+    where = "question %r" % (qid,)
+    if not isinstance(qdef, dict):
+        raise ValueError("%s: definition must be a dict, got %s" % (where, type(qdef).__name__))
+    qtype = qdef.get("type")
+    if qtype not in QUESTION_TYPES:
+        raise ValueError("%s: unknown type %r; use one of %s" % (where, qtype, ", ".join(QUESTION_TYPES)))
+    instructions = qdef.get("instructions")
+    if instructions is None or (isinstance(instructions, (str, list, dict)) and not instructions):
+        raise ValueError("%s: 'instructions' is missing or empty" % where)
+    instructions = _as_text(instructions).strip()
+    if not instructions:
+        raise ValueError("%s: 'instructions' is missing or empty" % where)
+    criteria = qdef.get("criteria")
+
+    if qtype == "choice":
+        if isinstance(criteria, dict):
+            keys = tuple(criteria.keys())
+            texts = tuple(None if v is None or v == "" else _as_text(v) for v in criteria.values())
+        elif isinstance(criteria, (list, tuple)):
+            keys = tuple(criteria)
+            texts = (None,) * len(keys)
+        else:
+            raise ValueError("%s: choice 'criteria' must be a list of labels or a {label: description} dict"
+                             % where)
+        if len(keys) < 2:
+            raise ValueError("%s: a choice needs at least 2 options, got %d" % (where, len(keys)))
+        if len({str(k) for k in keys}) != len(keys):
+            raise ValueError("%s: choice labels must be distinct" % where)
+        if any(str(k).strip() == "" for k in keys):
+            raise ValueError("%s: choice labels must not be empty" % where)
+        return Question("choice", instructions, keys, texts)
+
+    if qtype == "score":
+        if not isinstance(criteria, (list, tuple)):
+            raise ValueError("%s: score 'criteria' must be a list of level descriptions, low to high" % where)
+        if len(criteria) < 2:
+            raise ValueError("%s: a score needs at least 2 levels, got %d" % (where, len(criteria)))
+        return Question("score", instructions, tuple(range(len(criteria))), tuple(_as_text(c) for c in criteria))
+
+    # noul
+    if criteria is None:
+        criteria = {}
+    if not isinstance(criteria, dict):
+        raise ValueError("%s: noul 'criteria' must be a {\"false\": ..., \"true\": ...} dict" % where)
+    normalized = {str(k).lower(): v for k, v in criteria.items()}
+    unknown = set(normalized) - {"false", "true"}
+    if unknown:
+        raise ValueError("%s: noul 'criteria' keys must be 'false' and 'true', got %s"
+                         % (where, sorted(unknown)))
+    texts = tuple(None if normalized.get(k) in (None, "") else _as_text(normalized[k]) for k in ("false", "true"))
+    return Question("noul", instructions, (False, True), texts)
+
+
+def _label_index(question: Question, value: Any, where: str) -> int:
+    if question.type == "choice":
+        for i, key in enumerate(question.keys):
+            if value == key or str(value) == str(key):
+                return i
+        raise ValueError("%s: label %r is not one of the options %r" % (where, value, list(question.keys)))
+    if question.type == "noul":
+        if isinstance(value, bool):
+            return int(value)
+        if isinstance(value, (int, float)) and value in (0, 1):
+            return int(value)
+        if isinstance(value, str) and value.strip().lower() in _TRUE_STRINGS | _FALSE_STRINGS:
+            return int(value.strip().lower() in _TRUE_STRINGS)
+        raise ValueError("%s: label %r cannot be read as true/false" % (where, value))
+    for i, text in enumerate(question.texts):
+        if isinstance(value, str) and value == text:
+            return i
+    raise ValueError("%s: label %r is not a level index or level text" % (where, value))
+
+
+def target_from_label(question: Question, value: Any, label_smoothing: float = 0.0,
+                      where: str = "label") -> list[float]:
+    """A target distribution over the question's options from one hard label.
+
+    A score label may be fractional: `2.5` puts half its mass on level 2 and half on level 3.
+    `label_smoothing` mixes in a uniform distribution: `(1 - eps) * target + eps / k`.
+    """
+    if not 0.0 <= label_smoothing < 1.0:
+        raise ValueError("label_smoothing must be in [0, 1), got %r" % (label_smoothing,))
+    k = question.k
+    target = [0.0] * k
+    is_number = isinstance(value, (int, float)) and not isinstance(value, bool)
+    if question.type == "score" and is_number:
+        if not math.isfinite(value) or not 0 <= value <= k - 1:
+            raise ValueError("%s: score %r is outside the levels 0..%d" % (where, value, k - 1))
+        low = int(math.floor(value))
+        frac = value - low
+        target[low] = 1.0 - frac
+        if frac > 0:
+            target[low + 1] = frac
+    else:
+        target[_label_index(question, value, where)] = 1.0
+    if label_smoothing:
+        target = [(1.0 - label_smoothing) * t + label_smoothing / k for t in target]
+    return target
+
+
+def target_from_gold(question: Question, gold: Any, where: str = "gold") -> list[float]:
+    """A target distribution from soft probabilities.
+
+    `gold` is `{"probabilities": {...}}` or the probabilities dict itself, keyed by choice label,
+    `"false"`/`"true"` for a noul, or the level index (as a string or int) for a score. Missing
+    keys count as 0. The result is normalised; an all-zero target becomes uniform.
+    """
+    probs = gold.get("probabilities", gold) if isinstance(gold, dict) else None
+    if not isinstance(probs, dict):
+        raise ValueError("%s: expected a dict of probabilities, got %r" % (where, gold))
+    lookup = {str(k).lower() if isinstance(k, bool) else str(k): v for k, v in probs.items()}
+    if question.type == "noul":
+        names = ["false", "true"]
+    else:
+        names = [str(key) for key in question.keys]
+    unknown = set(lookup) - set(names)
+    if unknown:
+        raise ValueError("%s: probabilities name unknown options %s" % (where, sorted(unknown)))
+    target = []
+    for name in names:
+        value = lookup.get(name, 0.0)
+        if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value < 0:
+            raise ValueError("%s: probability for %r must be a finite non-negative number, got %r"
+                             % (where, name, value))
+        target.append(float(value))
+    total = sum(target)
+    if total <= 0:
+        return [1.0 / len(target)] * len(target)
+    return [t / total for t in target]

sysone/training.py

@@ -0,0 +1,399 @@
+"""Validation split and the training loop.
+
+The loop trains a `Network`: the decision head, plus the backbone when any of it learns. Items
+are cached option and summary vectors (`Features`) for a frozen backbone, or token layouts
+(`Encoded`) that go through the backbone on every step.
+
+All progress lives in a `TrainState`, so a run that is interrupted, or saved and loaded, resumes
+exactly where it stopped.
+"""
+
+from __future__ import annotations
+
+import copy
+import math
+import random
+from collections import Counter, defaultdict
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+
+import torch
+from torch import nn
+
+from ._util import autocast, empty_device_cache, is_out_of_memory
+from .data import Example
+from .features import Features, collate_features, read_vectors, run_in_batches
+from .layout import Encoded, collate
+from .losses import soft_ce
+
+OPTIMIZERS = ("adamw", "adamw_8bit", "sgd")
+SCHEDULERS = ("cosine", "linear", "constant")
+#: Validation checks per epoch with validation_interval="auto" when the backbone learns.
+CHECKS_PER_EPOCH = 10
+
+
+@dataclass
+class TrainItem:
+    inputs: Features | Encoded
+    qtype: int
+    target: torch.Tensor  # [k], in slot order
+
+
+@dataclass
+class TrainSettings:
+    epochs: int
+    batch_size: int
+    lrs: dict  # learning rate per parameter group: "head", "adapter", "full"
+    weight_decay: float
+    optimizer: str
+    scheduler: str
+    warmup_ratio: float
+    grad_clip: float
+    early_stopping: bool
+    n_iter_no_change: int  # epochs without a lower validation loss before stopping
+    target_loss: float | None
+    validation_interval: int  # optimizer steps between validation checks
+    micro_batch: int  # examples per forward pass to start with; halved when memory runs out
+
+
+@dataclass
+class TrainState:
+    """Everything needed to continue a run as if it had not stopped."""
+
+    fingerprint: str  # of the training data and the settings that shape the run
+    epoch_generator: torch.Tensor  # sampling generator state at the start of the current epoch
+    step: int = 0  # optimizer steps done
+    last_check: int = 0  # step of the last validation check
+    best_loss: float = math.inf
+    best_step: int = 0  # step of the check with the lowest validation loss
+    running_loss: float = 0.0  # training loss summed since the last check
+    running_n: int = 0
+    rng: dict = field(default_factory=dict)  # torch generator states after the last step
+    optimizer: dict | None = None  # optimizer state after the last step, on the CPU
+    current: dict | None = None  # trained weights after the last step, fp32 on the CPU
+
+    @classmethod
+    def fresh(cls, fingerprint: str, seed: int) -> TrainState:
+        return cls(fingerprint, torch.Generator().manual_seed(seed).get_state())
+
+    def nbytes(self) -> int:
+        tensors = list((self.current or {}).values())
+        for slot in (self.optimizer or {}).get("state", {}).values():
+            tensors.extend(v for v in slot.values() if torch.is_tensor(v))
+        return sum(t.numel() * t.element_size() for t in tensors)
+
+
+class Network:
+    """The trained parts of the model: the head, and the backbone when it learns.
+
+    With `backbone=None` items carry cached `Features`; otherwise they carry `Encoded` layouts,
+    run through the backbone under autocast to `autocast_dtype`.
+    """
+
+    def __init__(self, head: nn.Module, device: torch.device, backbone: nn.Module | None = None, arch=None,
+                 pad_id: int = 0, backbone_groups: dict | None = None, autocast_dtype: torch.dtype | None = None):
+        self.head = head
+        self.device = device
+        self.backbone = backbone
+        self.arch = arch
+        self.pad_id = pad_id
+        self.autocast_dtype = autocast_dtype
+        self.param_groups = {"head": [("head." + n, p) for n, p in head.named_parameters()]}
+        for group, named in (backbone_groups or {}).items():
+            if named:
+                self.param_groups[group] = [("backbone." + n, p) for n, p in named]
+
+    def named_trainable(self) -> list[tuple[str, nn.Parameter]]:
+        return [item for named in self.param_groups.values() for item in named]
+
+    def train(self, mode: bool = True) -> None:
+        self.head.train(mode)
+        if self.backbone is not None:
+            self.backbone.train(mode)
+
+    def forward(self, items: Sequence[TrainItem]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """`(logits [B, K], option_mask [B, K], target [B, K])` for a batch of items."""
+        qtype = torch.tensor([it.qtype for it in items], dtype=torch.long, device=self.device)
+        if self.backbone is None:
+            options, mask, summary = collate_features([it.inputs for it in items], self.device)
+        else:
+            batch = {k: v.to(self.device) for k, v in collate([it.inputs for it in items], self.pad_id).items()}
+            with autocast(self.device, self.autocast_dtype):
+                hidden = self.arch.hidden_states(self.backbone, batch)
+            options, summary = read_vectors(hidden, batch)
+            mask = batch["option_mask"]
+        target = torch.zeros(mask.shape, dtype=torch.float32)
+        for i, it in enumerate(items):
+            target[i, : it.target.shape[0]] = it.target
+        return self.head(options, mask, summary, qtype), mask, target.to(self.device)
+
+    def snapshot(self) -> dict[str, torch.Tensor]:
+        return {n: p.detach().to("cpu", torch.float32, copy=True) for n, p in self.named_trainable()}
+
+    @torch.no_grad()
+    def restore(self, weights: dict[str, torch.Tensor]) -> None:
+        for name, p in self.named_trainable():
+            p.copy_(weights[name])
+
+
+def _length(item: TrainItem) -> int:
+    return len(item.inputs.input_ids) if isinstance(item.inputs, Encoded) else item.inputs.options.shape[0]
+
+
+def split_validation(examples: Sequence[Example], fraction: float, seed: int) -> tuple[list[int], list[int]]:
+    """Indices `(train, validation)`.
+
+    Questions about the same state stay on the same side, so the validation slice never shares a
+    state with training. The split is stratified by category: each category gives about
+    `fraction` of its states. Returns an empty validation list when the data is too small to give
+    both sides at least one example.
+    """
+    if fraction <= 0:
+        return list(range(len(examples))), []
+    groups: dict[str, list[int]] = defaultdict(list)
+    for i, ex in enumerate(examples):
+        groups[ex.state].append(i)
+    strata: dict[str | None, list[str]] = defaultdict(list)
+    for state, members in groups.items():
+        strata[examples[members[0]].category].append(state)
+    rng = random.Random(seed)
+    held: set[str] = set()
+    for category in sorted(strata, key=lambda c: (c is None, str(c))):
+        states = sorted(strata[category])
+        rng.shuffle(states)
+        held.update(states[: int(round(fraction * len(states)))])
+    val = [i for i, ex in enumerate(examples) if ex.state in held]
+    train = [i for i, ex in enumerate(examples) if ex.state not in held]
+    if not val or not train:
+        return list(range(len(examples))), []
+    return train, val
+
+
+def sample_weights(examples: Sequence[Example], balance: str | None) -> torch.Tensor | None:
+    """Per-example sampling weights that even out categories or questions."""
+    if balance is None:
+        return None
+    if balance == "category":
+        keys = [ex.category for ex in examples]
+    elif balance == "question":
+        keys = [(ex.qid, ex.question.instructions) for ex in examples]
+    else:
+        raise ValueError("balance must be None, 'category' or 'question', got %r" % (balance,))
+    counts = Counter(keys)
+    return torch.tensor([1.0 / counts[k] for k in keys], dtype=torch.double)
+
+
+def make_optimizer(name: str, groups: list[dict], weight_decay: float) -> torch.optim.Optimizer:
+    if name == "adamw":
+        return torch.optim.AdamW(groups, weight_decay=weight_decay)
+    if name == "sgd":
+        return torch.optim.SGD(groups, weight_decay=weight_decay, momentum=0.9)
+    if name == "adamw_8bit":
+        try:
+            import bitsandbytes as bnb
+        except ImportError as exc:
+            raise ImportError("optimizer='adamw_8bit' needs bitsandbytes, which is not installed on this "
+                              "platform; use optimizer='adamw'") from exc
+        return bnb.optim.AdamW8bit(groups, weight_decay=weight_decay)
+    raise ValueError("optimizer must be one of %s, got %r" % (", ".join(OPTIMIZERS), name))
+
+
+def make_scheduler(optimizer, kind: str, warmup_ratio: float, total_steps: int, start: int = 0):
+    """A warmup-then-decay schedule, positioned at optimizer step `start`."""
+    warmup = int(warmup_ratio * total_steps)
+
+    def factor(step: int) -> float:
+        if step < warmup:
+            return (step + 1) / (warmup + 1)
+        if kind == "constant":
+            return 1.0
+        progress = (step - warmup) / max(1, total_steps - warmup)
+        if kind == "linear":
+            return max(0.0, 1.0 - progress)
+        return 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress)))
+
+    if kind not in SCHEDULERS:
+        raise ValueError("scheduler must be one of %s, got %r" % (", ".join(SCHEDULERS), kind))
+    for group in optimizer.param_groups:
+        group.setdefault("initial_lr", group["lr"])
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, factor, last_epoch=start - 1)
+
+
+@torch.no_grad()
+def evaluate_items(net: Network, items: Sequence[TrainItem], batch_size: int = 64) -> dict:
+    """Mean loss and accuracy (argmax of the prediction vs argmax of the target)."""
+    net.train(False)
+
+    def batch_fn(chunk: list[TrainItem]) -> list[tuple[float, bool]]:
+        logits, mask, target = net.forward(chunk)
+        losses = soft_ce(logits, target, mask, reduction="none")
+        hits = logits.argmax(-1) == target.argmax(-1)
+        return list(zip(losses.tolist(), hits.tolist()))
+
+    results = run_in_batches(items, batch_fn, start=batch_size, length=_length)
+    n = max(1, len(items))
+    return {"loss": sum(r[0] for r in results) / n, "accuracy": sum(r[1] for r in results) / n, "n": len(items)}
+
+
+def _rng_states() -> dict:
+    states = {"cpu": torch.get_rng_state()}
+    if torch.cuda.is_available():
+        states["cuda"] = torch.cuda.get_rng_state()
+    if torch.backends.mps.is_available():
+        states["mps"] = torch.mps.get_rng_state()
+    return states
+
+
+def _set_rng_states(states: dict) -> None:
+    if "cpu" in states:
+        torch.set_rng_state(states["cpu"])
+    if "cuda" in states and torch.cuda.is_available():
+        torch.cuda.set_rng_state(states["cuda"])
+    if "mps" in states and torch.backends.mps.is_available():
+        torch.mps.set_rng_state(states["mps"])
+
+
+def _to_cpu(value):
+    if torch.is_tensor(value):
+        return value.detach().to("cpu", copy=True)
+    if isinstance(value, dict):
+        return {k: _to_cpu(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return [_to_cpu(v) for v in value]
+    return copy.deepcopy(value)
+
+
+def train(net: Network, train_items: Callable[[int], Sequence[TrainItem]], n_train: int,
+          val_items: Sequence[TrainItem], weights: torch.Tensor | None, settings: TrainSettings,
+          state: TrainState, history: list, log: Callable[[str], None], resume: bool = False) -> TrainState:
+    """Train `net` in place, continuing from `state`; returns the updated state.
+
+    `train_items(epoch)` returns that epoch's training items: the same cached list every epoch, or
+    fresh ones when options are shuffled. Validation runs every `validation_interval` optimizer
+    steps and at the end; each check appends a record to `history`. Training stops after
+    `epochs`, after `n_iter_no_change` epochs without a lower validation loss (with early
+    stopping; counted in epochs however often validation runs), or once the monitored loss
+    reaches `target_loss`. The weights from the best check
+    are then restored, and `state` keeps the last weights so a later call can resume.
+
+    A `KeyboardInterrupt` stops training the same way and is then raised again.
+    """
+    steps_per_epoch = math.ceil(n_train / settings.batch_size)
+    total_steps = steps_per_epoch * settings.epochs
+    have_val = len(val_items) > 0
+    stop_early = settings.early_stopping and have_val
+    if settings.target_loss is not None and not have_val:
+        log("no validation slice: target_loss applies to the training loss")
+
+    optimizer = make_optimizer(settings.optimizer, [
+        {"params": [p for _, p in named], "lr": settings.lrs[group]} for group, named in net.param_groups.items()
+    ], settings.weight_decay)
+    best = None
+    if resume:
+        optimizer.load_state_dict(state.optimizer)
+        if math.isfinite(state.best_loss):
+            best = net.snapshot()  # the model holds the best weights of the saved run
+        net.restore(state.current)
+        _set_rng_states(state.rng)
+    scheduler = make_scheduler(optimizer, settings.scheduler, settings.warmup_ratio, total_steps, state.step)
+    params = [p for _, p in net.named_trainable()]
+    generator = torch.Generator()
+    micro = max(1, min(settings.micro_batch, settings.batch_size))
+
+    def optimizer_step(chunk: list[TrainItem]) -> float:
+        nonlocal micro
+        chunk = sorted(chunk, key=_length, reverse=True)  # similar lengths share a forward pass
+        while True:
+            try:
+                optimizer.zero_grad(set_to_none=True)
+                total = 0.0
+                for start in range(0, len(chunk), micro):
+                    part = chunk[start: start + micro]
+                    logits, mask, target = net.forward(part)
+                    loss = soft_ce(logits, target, mask)
+                    (loss * (len(part) / len(chunk))).backward()
+                    total += loss.item() * len(part)
+                break
+            except (RuntimeError, torch.OutOfMemoryError) as exc:
+                if not is_out_of_memory(exc) or micro == 1:
+                    raise
+                optimizer.zero_grad(set_to_none=True)
+                empty_device_cache()
+                micro = max(1, micro // 2)
+                log("out of memory: running %d examples per forward pass" % micro)
+        if settings.grad_clip:
+            nn.utils.clip_grad_norm_(params, settings.grad_clip)
+        optimizer.step()
+        scheduler.step()
+        return total
+
+    def check() -> bool:
+        """Validate at `state.step`; True when a stopping rule says to stop."""
+        nonlocal best
+        record = {"step": state.step, "epoch": round(state.step / steps_per_epoch, 3),
+                  "train_loss": state.running_loss / max(1, state.running_n)}
+        if have_val:
+            metrics = evaluate_items(net, val_items)
+            record.update(val_loss=metrics["loss"], val_accuracy=metrics["accuracy"])
+            if record["val_loss"] < state.best_loss - 1e-6:
+                state.best_loss, state.best_step, best = record["val_loss"], state.step, net.snapshot()
+        state.running_loss, state.running_n, state.last_check = 0.0, 0, state.step
+        history.append(record)
+        log("epoch %g/%d  step %d  " % (record["epoch"], settings.epochs, state.step)
+            + "  ".join("%s=%.4f" % (k, v) for k, v in record.items() if k not in ("epoch", "step")))
+        monitored = record["val_loss"] if have_val else record["train_loss"]
+        if settings.target_loss is not None and monitored <= settings.target_loss:
+            log("stopping: %s loss %.4f reached target_loss=%g"
+                % ("validation" if have_val else "training", monitored, settings.target_loss))
+            return True
+        if stop_early and state.step - state.best_step >= settings.n_iter_no_change * steps_per_epoch:
+            log("early stopping: no lower validation loss in %g epochs"
+                % round((state.step - state.best_step) / steps_per_epoch, 3))
+            return True
+        return False
+
+    def check_due() -> bool:
+        return state.step > state.last_check and (state.step % settings.validation_interval == 0
+                                                  or state.step == total_steps)
+
+    interrupted = False
+    try:
+        stop = check_due() and check()  # a check interrupted before it finished
+        if state.step >= total_steps and not stop:
+            log("all %d epochs are already trained; raise epochs to train more" % settings.epochs)
+        while not stop and state.step < total_steps:
+            epoch, first = divmod(state.step, steps_per_epoch)
+            generator.set_state(state.epoch_generator)
+            if weights is not None:
+                order = torch.multinomial(weights, n_train, replacement=True, generator=generator).tolist()
+            else:
+                order = torch.randperm(n_train, generator=generator).tolist()
+            items = train_items(epoch)
+            net.train(True)
+            for s in range(first, steps_per_epoch):
+                chunk = [items[i] for i in order[s * settings.batch_size: (s + 1) * settings.batch_size]]
+                loss_sum = optimizer_step(chunk)
+                state.step += 1
+                state.running_loss += loss_sum
+                state.running_n += len(chunk)
+                state.rng = _rng_states()
+                if s == steps_per_epoch - 1:
+                    state.epoch_generator = generator.get_state()
+                if check_due():
+                    stop = check()
+                    if stop:
+                        break
+                    net.train(True)
+    except KeyboardInterrupt:
+        interrupted = True
+        log("interrupted at step %d (epoch %g); fit() again to resume" % (state.step, state.step / steps_per_epoch))
+    finally:
+        optimizer.zero_grad(set_to_none=True)
+        state.current = net.snapshot()
+        state.optimizer = _to_cpu(optimizer.state_dict())
+        if best is not None:
+            net.restore(best)
+        net.train(False)
+    if interrupted:
+        raise KeyboardInterrupt
+    return state