add heavyball

Author: stephantul

model2vec/train/classifier.py

@@ -9,6 +9,7 @@
 import lightning as pl
 import numpy as np
 import torch
+from heavyball import AdamW, Muon
 from lightning.pytorch.callbacks import Callback, EarlyStopping
 from lightning.pytorch.utilities.types import OptimizerLRScheduler
 from sklearn.metrics import jaccard_score
@@ -429,7 +430,7 @@ def validation_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: i
 
     def configure_optimizers(self) -> OptimizerLRScheduler:
         """Configure optimizer and learning rate scheduler."""
-        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
+        optimizer = Muon(self.model.parameters(), lr=self.learning_rate)
         scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
             optimizer,
             mode="min",

pyproject.toml

@@ -63,7 +63,7 @@ dev = [
 distill = ["torch", "transformers", "scikit-learn"]
 onnx = ["onnx", "torch"]
 # train also installs inference
-train = ["torch", "lightning", "scikit-learn", "skops"]
+train = ["torch", "lightning", "scikit-learn", "skops", "heavyball"]
 inference = ["scikit-learn", "skops"]
 tokenizer = ["transformers"]