Merge branch '3313-enable-auto-early-stopping' of https://github.com/…

…ClaudioSalvatoreArcidiacono/LightGBM into 3313-enable-auto-early-stopping
microsoft · Jun 5, 2023 · f2ccd52 · f2ccd52
2 parents 2ca8cc1 + 1910076
commit f2ccd52
Showing 1 changed file with 38 additions and 23 deletions.
diff --git a/python-package/lightgbm/sklearn.py b/python-package/lightgbm/sklearn.py
@@ -85,16 +85,16 @@
 ]
 _LGBM_ScikitCustomEvalSetSplitter = Union[
     Callable[
-        [np.ndarray, np.ndarray],
-        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
+        [_LGBM_ScikitMatrixLike, _LGBM_LabelType],
+        Tuple[_LGBM_ScikitMatrixLike, _LGBM_ScikitMatrixLike, _LGBM_LabelType, _LGBM_LabelType]
     ],
     Callable[
-        [np.ndarray, np.ndarray, np.ndarray],
-        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]
+        [_LGBM_ScikitMatrixLike, _LGBM_LabelType, Optional[np.ndarray]],
+        Tuple[_LGBM_ScikitMatrixLike, _LGBM_ScikitMatrixLike, _LGBM_LabelType, _LGBM_LabelType, Optional[np.ndarray], Optional[np.ndarray]]
     ],
     Callable[
-        [np.ndarray, np.ndarray, np.ndarray],
-        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]
+        [_LGBM_ScikitMatrixLike, _LGBM_LabelType, Optional[np.ndarray], _LGBM_GroupType],
+        Tuple[_LGBM_ScikitMatrixLike, _LGBM_ScikitMatrixLike, _LGBM_LabelType, _LGBM_LabelType, Optional[np.ndarray], Optional[np.ndarray], _LGBM_GroupType, _LGBM_GroupType]
     ],
 ]
 _LGBM_ScikitValidSet = Tuple[_LGBM_ScikitMatrixLike, _LGBM_LabelType]
@@ -256,17 +256,17 @@ def __call__(
 
 
 def _train_test_split(
-    X,
-    y,
+    X: _LGBM_ScikitMatrixLike,
+    y: _LGBM_LabelType,
     weight,
     test_size: float,
     random_state: Optional[Union[int, np.random.RandomState]],
     stratified: bool,
 ) -> Tuple[
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
+    _LGBM_ScikitMatrixLike,
+    _LGBM_ScikitMatrixLike,
+    _LGBM_LabelType,
+    _LGBM_LabelType,
     Optional[np.ndarray],
     Optional[np.ndarray],
 ]:
@@ -319,7 +319,22 @@ def _train_test_split(
         return X_train, X_val, y_train, y_val, None, None
 
 
-def _train_test_group_split(X, y, weight, group, n_splits: int):
+def _train_test_group_split(
+    X: _LGBM_ScikitMatrixLike,
+    y: _LGBM_LabelType,
+    weight,
+    group: _LGBM_GroupType,
+    n_splits: int
+) -> Tuple[
+    _LGBM_ScikitMatrixLike,
+    _LGBM_ScikitMatrixLike,
+    _LGBM_LabelType,
+    _LGBM_LabelType,
+    Optional[np.ndarray],
+    Optional[np.ndarray],
+    _LGBM_GroupType,
+    _LGBM_GroupType,
+]:
     """Split X, y, weights and group into train and test subsets.
 
     Parameters
@@ -390,20 +405,20 @@ def _train_test_group_split(X, y, weight, group, n_splits: int):
 
 
 def _train_test_split_custom_splitter(
-    custom_splitter,
-    X,
-    y,
+    custom_splitter: _LGBM_ScikitCustomEvalSetSplitter,
+    X: _LGBM_ScikitMatrixLike,
+    y: _LGBM_LabelType,
     weight,
-    group
+    group: Optional[_LGBM_GroupType]
 ) -> Tuple[
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    Optional[np.ndarray],
-    Optional[np.ndarray],
+    _LGBM_ScikitMatrixLike,
+    _LGBM_ScikitMatrixLike,
+    _LGBM_LabelType,
+    _LGBM_LabelType,
     Optional[np.ndarray],
     Optional[np.ndarray],
+    Optional[_LGBM_GroupType],
+    Optional[_LGBM_GroupType],
 ]:
     """Call passed custom_splitter with appropriate arguments.