Fix pre-commit hooks🐙 (#384)

src/otc/models/callback.py

@@ -15,9 +15,9 @@
 from typing import Any
 
 import optuna
+import wandb
 from catboost import CatBoostClassifier
 
-import wandb
 from otc.config.config import settings
 from otc.data.fs import fs
 from otc.models.transformer_classifier import TransformerClassifier

src/otc/models/fttransformer.py

@@ -1,5 +1,5 @@
 """
-Implementation of FTTraansformer model.
+FT-Transformer.
 
 Adapted from:
 https://github.com/Yura52/rtdl/
@@ -54,20 +54,20 @@ def _all_or_none(values: list[Any]) -> bool:
 
 class CLSHead(nn.Module):
     """
-    2 Layer MLP projection head
+    2 Layer MLP projection head.
     """
 
     def __init__(self, *, d_in: int, d_hidden: int):
         """
-        tbd
+        Initialize the module.
         """
         super().__init__()
         self.first = nn.Linear(d_in, d_hidden)
         self.out = nn.Linear(d_hidden, 1)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
-        tbd
+        Forward pass.
         """
         x = x[:, 1:]
         x = self.out(F.relu(self.first(x))).squeeze(2)

src/otc/models/train_model.py

@@ -14,11 +14,11 @@
 import click
 import optuna
 import pandas as pd
+import wandb
 import yaml
 from optuna.exceptions import ExperimentalWarning
 from optuna.integration.wandb import WeightsAndBiasesCallback
 
-import wandb
 from otc.config.config import settings
 from otc.features.build_features import (
     features_categorical,

src/otc/models/transformer_classifier.py

@@ -210,7 +210,7 @@ def array_to_dataloader_pretrain(
         del data
         return tab_dl
 
-    def fit(
+    def fit(  # noqa: C901
         self,
         X: npt.NDArray | pd.DataFrame,
         y: npt.NDArray | pd.Series,
@@ -331,7 +331,7 @@ def fit(
 
                     with torch.autocast(device_type="cuda", dtype=torch.float16):
                         logits = self.clf(x_cat, x_cont)
-                        train_loss = criterion(logits, masks.float())
+                        train_loss = criterion(logits, masks.float())  # type: ignore[union-attr] # noqa: E501
 
                     scaler.scale(train_loss).backward()
                     scaler.step(optimizer)
@@ -357,7 +357,7 @@ def fit(
 
                         # for my implementation
                         logits = self.clf(x_cat, x_cont)
-                        val_loss = criterion(logits, masks.float())
+                        val_loss = criterion(logits, masks.float())  # type: ignore[union-attr] # noqa: E501
 
                         loss_in_epoch_val += val_loss.item()
 

src/otc/optim/scheduler.py

@@ -1,22 +1,54 @@
+"""
+Learnin rate scheduler with linear warmup phase and cosine decay.
+"""
 from typing import List
 
 import numpy as np
 from torch import optim
 
 
 class CosineWarmupScheduler(optim.lr_scheduler._LRScheduler):
+    """
+    Cosine learning rate scheduler with linear warmup.
+
+    Args:
+        optim (optim): learning rate scheduler
+    """
+
     def __init__(self, optimizer: optim.Optimizer, warmup: int, max_iters: int):
+        """
+        Cosine learning rate scheduler with linear warmup.
+
+        Args:
+            optimizer (optim.Optimizer): _description_
+            warmup (int): number of warmup iterations
+            max_iters (int): maximum number of iterations
+        """
         self.warmup = warmup
         self.max_num_iters = max_iters
         super().__init__(optimizer)
 
     def get_lr(self) -> List[float]:
-        lr_factor = self.get_lr_factor(epoch=self.last_epoch)
+        """
+        Get the learning rate.
+
+        Returns:
+            List[float]: List of learning rates.
+        """
+        lr_factor = self.get_lr_factor(iteration=self.last_epoch)
         return [base_lr * lr_factor for base_lr in self.base_lrs]
 
-    def get_lr_factor(self, epoch: int) -> float:
+    def get_lr_factor(self, iteration: int) -> float:
+        """
+        Get the learning rate factor for the given epoch.
+
+        Args:
+            epoch (int): epoch number
 
-        lr_factor = 0.5 * (1 + np.cos(np.pi * epoch / self.max_num_iters))
-        if epoch <= self.warmup:
-            lr_factor *= epoch * 1.0 / self.warmup
+        Returns:
+            float: learning rate factor
+        """
+        lr_factor = 0.5 * (1 + np.cos(np.pi * iteration / self.max_num_iters))
+        if iteration <= self.warmup:
+            lr_factor *= iteration * 1.0 / self.warmup
         return lr_factor

tests/test_transformer_classifier.py

@@ -16,55 +16,55 @@
 #         unittest (_type_): unittest module
 #     """
 
-    # def setup(self) -> None:
-    #     """
-    #     Set up basic classifier and data.
+# def setup(self) -> None:
+#     """
+#     Set up basic classifier and data.
 
-    #     Prepares inputs and expected outputs for testing.
-    #     """
-    #     self.x_train = pd.DataFrame(
-    #         [[1, 2], [3, 4], [1, 2], [3, 4]], columns=["BEST_ASK", "BEST_BID"]
-    #     )
-    #     self.y_train = pd.Series([1, 1, -1, -1])
-    #     self.x_test = pd.DataFrame(
-    #         [[1, 2], [3, 4], [1, 2], [3, 4]], columns=["BEST_ASK", "BEST_BID"]
-    #     )
-    #     self.y_test = pd.Series([1, -1, 1, -1])
+#     Prepares inputs and expected outputs for testing.
+#     """
+#     self.x_train = pd.DataFrame(
+#         [[1, 2], [3, 4], [1, 2], [3, 4]], columns=["BEST_ASK", "BEST_BID"]
+#     )
+#     self.y_train = pd.Series([1, 1, -1, -1])
+#     self.x_test = pd.DataFrame(
+#         [[1, 2], [3, 4], [1, 2], [3, 4]], columns=["BEST_ASK", "BEST_BID"]
+#     )
+#     self.y_test = pd.Series([1, -1, 1, -1])
 
-    #     dl_params = {
-    #         "batch_size": 8,
-    #         "shuffle": False,
-    #         "device": "cpu",
-    #     }
+#     dl_params = {
+#         "batch_size": 8,
+#         "shuffle": False,
+#         "device": "cpu",
+#     }
 
-    #     module_params = {
-    #         "depth": 1,
-    #         "heads": 2,
-    #         "dim": 2,
-    #         "dim_out": 1,
-    #         "mlp_act": nn.ReLU,
-    #         "mlp_hidden_mults": (4, 2),
-    #         "attn_dropout": 0.5,
-    #         "ff_dropout": 0.5,
-    #         "cat_features": [],
-    #         "cat_cardinalities": (),
-    #         "num_continuous": self.x_train.shape[1],
-    #     }
+#     module_params = {
+#         "depth": 1,
+#         "heads": 2,
+#         "dim": 2,
+#         "dim_out": 1,
+#         "mlp_act": nn.ReLU,
+#         "mlp_hidden_mults": (4, 2),
+#         "attn_dropout": 0.5,
+#         "ff_dropout": 0.5,
+#         "cat_features": [],
+#         "cat_cardinalities": (),
+#         "num_continuous": self.x_train.shape[1],
+#     }
 
-    #     optim_params = {"lr": 0.1, "weight_decay": 1e-3}
+#     optim_params = {"lr": 0.1, "weight_decay": 1e-3}
 
-    #     with patch.object(TransformerClassifier, "epochs", 5):
-    #         self.clf = TransformerClassifier(
-    #             module=FTTransformer,  # type: ignore
-    #             module_params=module_params,
-    #             optim_params=optim_params,
-    #             dl_params=dl_params,
-    #             callbacks=CallbackContainer([]),
-    #         ).fit(self.x_train, self.y_train)
+#     with patch.object(TransformerClassifier, "epochs", 5):
+#         self.clf = TransformerClassifier(
+#             module=FTTransformer,  # type: ignore
+#             module_params=module_params,
+#             optim_params=optim_params,
+#             dl_params=dl_params,
+#             callbacks=CallbackContainer([]),
+#         ).fit(self.x_train, self.y_train)
 
-    # def test_sklearn_compatibility(self) -> None:
-    #     """
-    #     Test, if classifier is compatible with sklearn.
-    #     """
-    #     with patch.object(TransformerClassifier, "epochs", 1):
-    #         check_estimator(self.clf)
+# def test_sklearn_compatibility(self) -> None:
+#     """
+#     Test, if classifier is compatible with sklearn.
+#     """
+#     with patch.object(TransformerClassifier, "epochs", 1):
+#         check_estimator(self.clf)