fix: optional metrics in sl and lp trainers

krishnanlab · Jul 10, 2023 · 771ce41 · 771ce41
1 parent 7f4aea5
commit 771ce41
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/src/obnb/model_trainer/label_propagation.py b/src/obnb/model_trainer/label_propagation.py
@@ -1,13 +1,15 @@
+import numpy as np
+
 from obnb.model_trainer.base import StandardTrainer
-from obnb.typing import LogLevel, Optional
+from obnb.typing import Callable, Dict, LogLevel, Optional
 
 
 class LabelPropagationTrainer(StandardTrainer):
  """Label propagation trainer."""
 
  def __init__(
  self,
- metrics,
+ metrics: Optional[Dict[str, Callable[[np.ndarray, np.ndarray], float]]] = None,
  train_on="train",
  log_level: LogLevel = "WARNING",
  log_path: Optional[str] = None,

diff --git a/src/obnb/model_trainer/supervised_learning.py b/src/obnb/model_trainer/supervised_learning.py
@@ -1,5 +1,7 @@
+import numpy as np
+
 from obnb.model_trainer.base import StandardTrainer
-from obnb.typing import LogLevel, Optional
+from obnb.typing import Callable, Dict, LogLevel, Optional
 
 
 class SupervisedLearningTrainer(StandardTrainer):
@@ -21,7 +23,7 @@ class SupervisedLearningTrainer(StandardTrainer):
 
  def __init__(
  self,
- metrics,
+ metrics: Optional[Dict[str, Callable[[np.ndarray, np.ndarray], float]]] = None,
  train_on="train",
  log_level: LogLevel = "WARNING",
  log_path: Optional[str] = None,