docs: add method typing

Author: marcpinet

neuralnetlib/callbacks.py

@@ -1,48 +1,48 @@
 from neuralnetlib.metrics import Metric
+import numpy as np
 
 
 class Callback:
-    def on_train_begin(self, logs=None):
+    def on_train_begin(self, logs: dict | None = None) -> None:
         pass
 
-    def on_train_end(self, logs=None):
+    def on_train_end(self, logs: dict | None = None) -> None:
         pass
 
-    def on_epoch_begin(self, epoch, logs=None):
+    def on_epoch_begin(self, epoch: int, logs: dict | None = None) -> None:
         pass
 
-    def on_epoch_end(self, epoch, logs=None):
+    def on_epoch_end(self, epoch: int, logs: dict | None = None) -> None:
         pass
 
-    def on_batch_begin(self, batch, logs=None):
+    def on_batch_begin(self, batch: int, logs: dict | None = None) -> None:
         pass
 
-    def on_batch_end(self, batch, logs=None):
+    def on_batch_end(self, batch: int, logs: dict | None = None) -> None:
         pass
 
 
 class EarlyStopping(Callback):
-    def __init__(self, patience: int = 5, min_delta: float = 0.001, restore_best_weights: bool = True,
-                 start_from_epoch: int = 0, monitor: str = 'loss', mode: str = 'auto', baseline: float = None):
+    def __init__(self, patience: int = 5, min_delta: float = 0.001, restore_best_weights: bool = True,start_from_epoch: int = 0, monitor: str = 'loss', mode: str = 'auto', baseline: float | None = None) -> None:
         super().__init__()
-        self.patience = patience
-        self.min_delta = min_delta
-        self.restore_best_weights = restore_best_weights
-        self.start_from_epoch = start_from_epoch
-        self.monitor = Metric(monitor) if monitor != 'loss' else 'loss'
-        self.mode = mode
-        self.baseline = baseline
-        self.best_weights = None
-        self.best_metric = None
-        self.patience_counter = 0
-        self.stop_training = False
-
-    def on_train_begin(self, logs=None):
+        self.patience: int = patience
+        self.min_delta: float = min_delta
+        self.restore_best_weights: bool = restore_best_weights
+        self.start_from_epoch: int = start_from_epoch
+        self.monitor: Metric | str = Metric(monitor) if monitor != 'loss' else 'loss'
+        self.mode: str = mode
+        self.baseline: float | None = baseline
+        self.best_weights: list | None = None
+        self.best_metric: float | None = None
+        self.patience_counter: int = 0
+        self.stop_training: bool = False
+
+    def on_train_begin(self, logs: dict | None = None) -> None:
         self.patience_counter = 0
         self.best_metric = None
         self.stop_training = False
 
-    def on_epoch_end(self, epoch, logs=None):
+    def on_epoch_end(self, epoch: int, logs: dict | None = None) -> bool:
         logs = logs or {}
         model = logs.get('model')
         if epoch < self.start_from_epoch or model is None:
@@ -72,14 +72,14 @@ def on_epoch_end(self, epoch, logs=None):
             self.stop_training = True
             if self.restore_best_weights and self.best_weights is not None:
                 for layer, best_weights in zip([layer for layer in model.layers if hasattr(layer, 'weights')],
-                                               self.best_weights):
+                                             self.best_weights):
                     layer.weights = best_weights
             print(f"\nEarly stopping triggered after epoch {epoch + 1}")
             return True
 
         return False
 
-    def _get_monitor_value(self, logs):
+    def _get_monitor_value(self, logs: dict) -> float:
         logs = logs or {}
         if isinstance(self.monitor, Metric):
             monitor_value = logs.get(self.monitor.name)
@@ -93,5 +93,5 @@ def _get_monitor_value(self, logs):
                 monitor_value = logs
             else:
                 raise ValueError(f"Monitored metric '{self.monitor}' is not available. "
-                                 f"Available metrics are: {','.join(logs.keys())}")
-        return monitor_value
+                               f"Available metrics are: {','.join(logs.keys())}")
+        return float(monitor_value)