feat: remove lightning (#337)

* feat: remove lightning

* additional test coverage

* reduce steps

* weight metrics by batch size

Author: stephantul

model2vec/train/README.md

@@ -1,6 +1,6 @@
 # Training
 
-Aside from [distillation](../../README.md#distillation), `model2vec` also supports training simple classifiers on top of static models, using [pytorch](https://pytorch.org/), [lightning](https://lightning.ai/) and [scikit-learn](https://scikit-learn.org/stable/index.html).
+Aside from [distillation](../../README.md#distillation), `model2vec` also supports training simple classifiers on top of static models, using [pytorch](https://pytorch.org/) and [scikit-learn](https://scikit-learn.org/stable/index.html).
 
 We support both single and multi-label classification, which work seamlessly based on the labels you provide.
 
@@ -53,7 +53,7 @@ print(classification_report)
 
 As you can see, we got a pretty nice 91% accuracy, with only 81 seconds of training.
 
-The training loop is handled by [`lightning`](https://pypi.org/project/lightning/). By default the training loop splits the data into a train and validation split, with 90% of the data being used for training and 10% for validation. By default, it runs with early stopping on the validation set accuracy, with a patience of 5.
+The training loop is a plain PyTorch loop (see [`model2vec/train/trainer.py`](trainer.py)). By default the training loop splits the data into a train and validation split, with 90% of the data being used for training and 10% for validation. By default, it runs with early stopping on the validation set accuracy, with a patience of 5.
 
 Note that this model is as fast as you're used to from us:
 
@@ -142,9 +142,9 @@ The core functionality of the `StaticModelForClassification` is contained in a c
 * `train_test_split`: governs the train test split before classification.
 * `prepare_dataset`: Selects the `torch.Dataset` that will be used in the `Dataloader` during training.
 * `_encode`: The encoding function used in the model.
-* `fit`: contains all the lightning-related fitting logic.
+* `fit`: contains all the fitting logic.
 
-The training of the model is done in a `lighting.LightningModule`, which can be modified but is very basic.
+The training loop itself lives in `model2vec.train.trainer.run_training_loop`, a plain torch loop that is fairly basic and easy to modify. Each task passes in its own loss function (and, for classification, a small function that computes extra validation metrics like accuracy).
 
 # Results
 

model2vec/train/__init__.py

@@ -1,5 +1,3 @@
-import logging
-
 from model2vec.utils import get_package_extras, importable
 
 _REQUIRED_EXTRA = "train"
@@ -10,9 +8,5 @@
 from model2vec.train.classifier import StaticModelForClassification
 from model2vec.train.regression import StaticModelForRegression
 from model2vec.train.similarity import StaticModelForSimilarity
-from model2vec.train.utils import TipFilter
 
 __all__ = ["StaticModelForClassification", "StaticModelForSimilarity", "StaticModelForRegression"]
-
-
-logging.getLogger("lightning.pytorch.utilities.rank_zero").addFilter(TipFilter())

model2vec/train/base.py

@@ -2,14 +2,10 @@
 
 import logging
 from collections.abc import Sequence
-from tempfile import TemporaryDirectory
 from typing import Any, TypeVar
 
-import lightning.pytorch as pl
 import numpy as np
 import torch
-from lightning.pytorch import LightningModule
-from lightning.pytorch.callbacks import Callback, EarlyStopping
 from tokenizers import Encoding, Tokenizer
 from torch import nn
 from torch.nn.utils.rnn import pad_sequence
@@ -18,10 +14,10 @@
 from model2vec.inference import StaticModelPipeline
 from model2vec.model import PathLike, StaticModel
 from model2vec.train.dataset import TextDataset
+from model2vec.train.trainer import MetricsFn, default_metrics, resolve_device, run_training_loop
 from model2vec.train.utils import (
     get_probable_pad_token_id,
     logit,
-    suppress_lightning_warnings,
     to_pipeline,
     train_test_split,
 )
@@ -222,10 +218,9 @@ def encode(self, X: list[str], batch_size: int = 1024, show_progress_bar: bool =
 
         return np.concatenate(pred, axis=0)
 
-    def forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
         """Forward pass through the mean, and a classifier layer after."""
-        encoded = self._encode(input_ids)
-        return self.head(encoded), encoded
+        return self.head(self._encode(input_ids))
 
     def tokenize(self, texts: list[str], max_length: int | None = 512) -> torch.Tensor:
         """Tokenize a bunch of strings into a single padded 2D tensor.
@@ -308,10 +303,10 @@ def _check_val_split(
 
         return train_texts, validation_texts, train_labels, validation_labels
 
-    @suppress_lightning_warnings
     def _train(
         self,
-        module: LightningModule,
+        loss_function: nn.Module,
+        learning_rate: float,
         train_dataset: TextDataset,
         val_dataset: TextDataset,
         batch_size: int,
@@ -320,48 +315,31 @@ def _train(
         max_epochs: int | None,
         device: str,
         validation_steps: int | None,
+        compute_metrics: MetricsFn = default_metrics,
     ) -> None:
-        callbacks: list[Callback] = []
-        if early_stopping_patience is not None:
-            callback = EarlyStopping(
-                monitor=self.val_metric,
-                mode=self.early_stopping_direction,
-                patience=early_stopping_patience,
-                min_delta=0.001,
-            )
-            callbacks.append(callback)
-
         val_check_interval, check_val_every_epoch = self._determine_val_check_interval(
             validation_steps, len(train_dataset), batch_size
         )
 
-        with TemporaryDirectory() as tempdir:
-            trainer = pl.Trainer(
-                min_epochs=min_epochs,
-                max_epochs=max_epochs,
-                callbacks=callbacks,
-                val_check_interval=val_check_interval,
-                check_val_every_n_epoch=check_val_every_epoch,
-                accelerator=device,
-                default_root_dir=tempdir,
-            )
-
-            trainer.fit(
-                module,
-                train_dataloaders=train_dataset.to_dataloader(shuffle=True, batch_size=batch_size),
-                val_dataloaders=val_dataset.to_dataloader(shuffle=False, batch_size=batch_size),
-            )
-            best_model_path = trainer.checkpoint_callback.best_model_path  # type: ignore
-            best_model_weights = torch.load(best_model_path, weights_only=True)
-
-        state_dict = {}
-        for weight_name, weight in best_model_weights["state_dict"].items():
-            if "loss_function" in weight_name:
-                # Skip the loss function class weight as its not needed for predictions
-                continue
-            state_dict[weight_name.removeprefix("model.")] = weight
+        state_dict = run_training_loop(
+            model=self,
+            loss_function=loss_function,
+            learning_rate=learning_rate,
+            val_metric=self.val_metric,
+            early_stopping_direction=self.early_stopping_direction,
+            train_loader=train_dataset.to_dataloader(shuffle=True, batch_size=batch_size),
+            val_loader=val_dataset.to_dataloader(shuffle=False, batch_size=batch_size),
+            early_stopping_patience=early_stopping_patience,
+            min_epochs=min_epochs,
+            max_epochs=max_epochs,
+            device=resolve_device(device),
+            val_check_interval=val_check_interval,
+            check_val_every_epoch=check_val_every_epoch,
+            compute_metrics=compute_metrics,
+        )
 
         self.load_state_dict(state_dict)
+        self.to("cpu")
         self.eval()
 
     @staticmethod

model2vec/train/classifier.py

@@ -5,22 +5,35 @@
 from itertools import chain
 from typing import Any, Literal, cast
 
-import lightning as pl
 import numpy as np
 import torch
+from sklearn.metrics import jaccard_score
 from tokenizers import Tokenizer
+from torch import nn
 from tqdm import trange
 
 from model2vec.inference import evaluate_single_or_multi_label
 from model2vec.train.base import BaseFinetuneable
-from model2vec.train.lightning_modules import ClassifierLightningModule, MultiLabelClassifierLightningModule
-from model2vec.train.utils import DEFAULT_RANDOM_SEED
+from model2vec.train.utils import DEFAULT_RANDOM_SEED, seed_everything
 
 logger = logging.getLogger(__name__)
 
 LabelType = list[str] | list[list[str]]
 
 
+def _classifier_metrics(head_out: torch.Tensor, y: torch.Tensor, loss: torch.Tensor) -> dict[str, float]:
+    """Validation metrics for single-label classification: loss and accuracy."""
+    accuracy = (head_out.argmax(dim=1) == y).float().mean()
+    return {"val_loss": loss.item(), "val_accuracy": accuracy.item()}
+
+
+def _multilabel_classifier_metrics(head_out: torch.Tensor, y: torch.Tensor, loss: torch.Tensor) -> dict[str, float]:
+    """Validation metrics for multi-label classification: loss and Jaccard accuracy."""
+    preds = (torch.sigmoid(head_out) > 0.5).float()
+    accuracy = cast(float, jaccard_score(y.cpu(), preds.cpu(), average="samples"))
+    return {"val_loss": loss.item(), "val_accuracy": accuracy}
+
+
 class StaticModelForClassification(BaseFinetuneable):
     val_metric = "val_accuracy"
     early_stopping_direction = "max"
@@ -127,7 +140,7 @@ def fit(
     ) -> StaticModelForClassification:
         """Fit a model.
 
-        This function creates a Lightning Trainer object and fits the model to the data.
+        This function trains the model with a plain torch training loop.
         It supports both single-label and multi-label classification.
         We use early stopping. After training, the weights of the best model are loaded back into the model.
 
@@ -157,7 +170,7 @@ def fit(
         :return: The fitted model.
         :raises ValueError: If either X_val or y_val are provided, but not both.
         """
-        pl.seed_everything(random_seed)
+        seed_everything(random_seed)
         logger.info("Re-initializing model.")
 
         # Determine whether the task is multilabel based on the type of y.
@@ -177,16 +190,16 @@ def fit(
         train_dataset, val_dataset = self._create_datasets(X, y, X_val, y_val, test_size)
         batch_size = self._determine_batch_size(batch_size, len(train_dataset))
 
-        c: pl.LightningModule
         if self.multilabel:
-            c = MultiLabelClassifierLightningModule(
-                self, learning_rate=learning_rate, class_weight=resolved_class_weight
-            )
+            loss_function: nn.Module = nn.BCEWithLogitsLoss(pos_weight=resolved_class_weight)
+            compute_metrics = _multilabel_classifier_metrics
         else:
-            c = ClassifierLightningModule(self, learning_rate=learning_rate, class_weight=resolved_class_weight)
+            loss_function = nn.CrossEntropyLoss(weight=resolved_class_weight)
+            compute_metrics = _classifier_metrics
 
         self._train(
-            module=c,
+            loss_function=loss_function,
+            learning_rate=learning_rate,
             train_dataset=train_dataset,
             val_dataset=val_dataset,
             batch_size=batch_size,
@@ -195,6 +208,7 @@ def fit(
             max_epochs=max_epochs,
             device=device,
             validation_steps=validation_steps,
+            compute_metrics=compute_metrics,
         )
 
         return self

model2vec/train/lightning_modules.py

@@ -2,11 +2,15 @@
 
 import logging
 
-from model2vec.train.lightning_modules import RegressionLightningModule
+from torch import nn
+
 from model2vec.train.similarity import StaticModelForSimilarity
 
 logger = logging.getLogger(__name__)
 
 
 class StaticModelForRegression(StaticModelForSimilarity):
-    _lightning_class = RegressionLightningModule
+    @staticmethod
+    def _build_loss_function() -> nn.Module:
+        """Construct the loss function used to train this model."""
+        return nn.MSELoss()