fix: trainer uses best checkpoint

model2vec/train/trainer.py

@@ -88,7 +88,7 @@ def _run_validation(
     return {key: total / total_samples for key, total in weighted_sums.items()}
 
 
-def run_training_loop(
+def run_training_loop(  # noqa: C901
     model: nn.Module,
     loss_function: nn.Module,
     learning_rate: float,
@@ -121,7 +121,7 @@ def run_training_loop(
     :param val_check_interval: If set, validate every this many training steps.
     :param check_val_every_epoch: If set, validate every this many epochs.
     :param compute_metrics: Computes validation metrics from `(head_out, y, loss)`. Defaults to just `val_loss`.
-    :return: The model's state dict as of the last validation check before training stopped.
+    :return: The model's state dict from the validation check with the best `val_metric`.
     """
     model.to(device)
     loss_function.to(device)
@@ -144,17 +144,22 @@ def run_training_loop(
 
     max_epochs = _resolve_max_epochs(max_epochs)
 
-    last_checkpoint = copy.deepcopy(model.state_dict())
+    best_checkpoint = copy.deepcopy(model.state_dict())
+    best_val_metric = float("inf") if early_stopping_direction == "min" else float("-inf")
     current_epoch = 0
     global_step = 0
     postfix: dict[str, str] = {}
     latest_val_loss: float | None = None
 
     def validate_and_checkpoint() -> bool:
-        nonlocal last_checkpoint, latest_val_loss
+        nonlocal best_checkpoint, best_val_metric, latest_val_loss
         metrics = _run_validation(model, loss_function, compute_metrics, val_loader, device)
         latest_val_loss = metrics["val_loss"]
-        last_checkpoint = copy.deepcopy(model.state_dict())
+        current = metrics[val_metric]
+        improved = current < best_val_metric if early_stopping_direction == "min" else current > best_val_metric
+        if improved:
+            best_val_metric = current
+            best_checkpoint = copy.deepcopy(model.state_dict())
         postfix.update({key: f"{value:.4f}" for key, value in metrics.items()})
         if early_stopper is None:
             return False
@@ -178,17 +183,17 @@ def validate_and_checkpoint() -> bool:
                     should_stop = validate_and_checkpoint()
                     pbar.set_postfix(postfix)
                     if should_stop and (min_epochs is None or current_epoch >= min_epochs):
-                        return last_checkpoint
+                        return best_checkpoint
 
             current_epoch += 1
 
             if check_val_every_epoch is not None and current_epoch % check_val_every_epoch == 0:
                 should_stop = validate_and_checkpoint()
                 pbar.set_postfix(postfix)
                 if should_stop and (min_epochs is None or current_epoch >= min_epochs):
-                    return last_checkpoint
+                    return best_checkpoint
 
             _step_plateau_scheduler(scheduler, latest_val_loss)
 
             if current_epoch >= max_epochs:
-                return last_checkpoint
+                return best_checkpoint

tests/test_inference.py

@@ -23,9 +23,9 @@ def test_init_predict(mock_inference_pipeline: StaticModelPipeline) -> None:
     if mock_inference_pipeline.head.activation == Activation.SIGMOID:
         assert mock_inference_pipeline.classes_ is not None
         if isinstance(mock_inference_pipeline.classes_[0], str):
-            target = [["a", "b"]]
+            target = [["b"]]
         else:
-            target = [[0, 1]]  # type: ignore
+            target = [[1]]  # type: ignore
     else:
         assert mock_inference_pipeline.classes_ is not None
         if isinstance(mock_inference_pipeline.classes_[0], str):
@@ -82,9 +82,9 @@ def test_roundtrip_save(mock_inference_pipeline: StaticModelPipeline) -> None:
         if mock_inference_pipeline.head.activation == Activation.SIGMOID:
             assert mock_inference_pipeline.classes_ is not None
             if isinstance(mock_inference_pipeline.classes_[0], str):
-                target = [["a", "b"]]
+                target = [["b"]]
             else:
-                target = [[0, 1]]  # type: ignore
+                target = [[1]]  # type: ignore
         else:
             assert mock_inference_pipeline.classes_ is not None
             if isinstance(mock_inference_pipeline.classes_[0], str):

tests/test_trainable.py

@@ -152,9 +152,9 @@ def test_predict(mock_trained_pipeline: StaticModelForClassification) -> None:
     result = mock_trained_pipeline.predict(["dog cat", "dog"]).tolist()
     if mock_trained_pipeline.multilabel:
         if type(mock_trained_pipeline.classes_[0]) == str:
-            assert result == [["a", "b"], ["a", "b"]]
+            assert result == [["b"], ["b"]]
         else:
-            assert result == [[0, 1], [0, 1]]
+            assert result == [[1], [1]]
     else:
         if type(mock_trained_pipeline.classes_[0]) == str:
             assert result == ["b", "b"]