Fix default take value of train and predict methods

related-sciences · Oct 13, 2023 · 3c6cc0f · 3c6cc0f
1 parent bbe83ba
commit 3c6cc0f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/nxontology_ml/model/predict.py b/nxontology_ml/model/predict.py
@@ -49,7 +49,7 @@ def predict(
     training_set: tuple[np.ndarray, np.ndarray],
     nxo: NXOntology[str],
     include_feature_values: bool = True,
-    take: int | None = 0,
+    take: int | None = None,
 ) -> pd.DataFrame:
     target_nodes: list[str] = list(get_disease_nodes(take=take, nxo=nxo))
     assert len(target_nodes) > 0, "No disease node found"
@@ -96,8 +96,8 @@ def train_predict(
     nxo: NXOntology[str] | None = None,
     training_set: tuple[np.ndarray, np.ndarray] | None = None,
     include_feature_values: bool = True,
-    train_take: int | None = 0,
-    predict_take: int | None = 0,
+    train_take: int | None = None,
+    predict_take: int | None = None,
 ) -> pd.DataFrame:
     """
     Run both model training and prediction tasks.

diff --git a/nxontology_ml/model/train.py b/nxontology_ml/model/train.py
@@ -24,7 +24,7 @@ def train_model(
     conf: ModelConfig = DEFAULT_MODEL_CONFIG,
     nxo: NXOntology[str] | None = None,
     training_set: tuple[np.ndarray, np.ndarray] | None = None,
-    take: int | None = 0,
+    take: int | None = None,
 ) -> tuple[Pipeline, CatBoostClassifier]:
     nxo = nxo or get_efo_otar_slim()
     nxo.freeze()