sklearn.py 50.3 KB
Newer Older
wxchan's avatar
wxchan committed
1
# coding: utf-8
2
"""Scikit-learn wrapper interface for LightGBM."""
3
import copy
4
from inspect import signature
5
from typing import Callable, Dict, Optional, Union
6

wxchan's avatar
wxchan committed
7
import numpy as np
8

9
10
11
12
from .basic import Dataset, LightGBMError, _choose_param_value, _ConfigAliases, _log_warning
from .compat import (SKLEARN_INSTALLED, LGBMNotFittedError, _LGBMAssertAllFinite, _LGBMCheckArray,
                     _LGBMCheckClassificationTargets, _LGBMCheckSampleWeight, _LGBMCheckXY, _LGBMClassifierBase,
                     _LGBMComputeSampleWeight, _LGBMLabelEncoder, _LGBMModelBase, _LGBMRegressorBase, dt_DataTable,
13
                     pd_DataFrame)
wxchan's avatar
wxchan committed
14
from .engine import train
15

wxchan's avatar
wxchan committed
16

17
class _ObjectiveFunctionWrapper:
18
    """Proxy class for objective function."""
19

20
21
    def __init__(self, func):
        """Construct a proxy class.
22

23
24
        This class transforms objective function to match objective function with signature ``new_func(preds, dataset)``
        as expected by ``lightgbm.engine.train``.
25

26
27
28
        Parameters
        ----------
        func : callable
29
            Expects a callable with signature ``func(y_true, y_pred)`` or ``func(y_true, y_pred, group)``
30
31
32
33
34
35
            and returns (grad, hess):

                y_true : array-like of shape = [n_samples]
                    The target values.
                y_pred : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
                    The predicted values.
36
37
                    Predicted values are returned before any transformation,
                    e.g. they are raw margin instead of probability of positive class for binary task.
38
                group : array-like
39
40
41
                    Group/query data.
                    Only used in the learning-to-rank task.
                    sum(group) = n_samples.
42
43
                    For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
                    where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
44
                grad : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
45
46
                    The value of the first order derivative (gradient) of the loss
                    with respect to the elements of y_pred for each sample point.
47
                hess : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
48
49
                    The value of the second order derivative (Hessian) of the loss
                    with respect to the elements of y_pred for each sample point.
wxchan's avatar
wxchan committed
50

Nikita Titov's avatar
Nikita Titov committed
51
52
53
54
55
        .. note::

            For multi-class task, the y_pred is group by class_id first, then group by row_id.
            If you want to get i-th row y_pred in j-th class, the access way is y_pred[j * num_data + i]
            and you should group grad and hess in this way as well.
56
57
        """
        self.func = func
wxchan's avatar
wxchan committed
58

59
60
61
62
63
64
65
66
67
68
69
70
71
    def __call__(self, preds, dataset):
        """Call passed function with appropriate arguments.

        Parameters
        ----------
        preds : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
            The predicted values.
        dataset : Dataset
            The training dataset.

        Returns
        -------
        grad : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
72
73
            The value of the first order derivative (gradient) of the loss
            with respect to the elements of preds for each sample point.
74
        hess : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
75
76
            The value of the second order derivative (Hessian) of the loss
            with respect to the elements of preds for each sample point.
77
        """
wxchan's avatar
wxchan committed
78
        labels = dataset.get_label()
79
        argc = len(signature(self.func).parameters)
80
        if argc == 2:
81
            grad, hess = self.func(labels, preds)
82
        elif argc == 3:
83
            grad, hess = self.func(labels, preds, dataset.get_group())
84
        else:
85
            raise TypeError(f"Self-defined objective function should have 2 or 3 arguments, got {argc}")
wxchan's avatar
wxchan committed
86
87
88
89
90
91
92
93
94
95
96
        """weighted for objective"""
        weight = dataset.get_weight()
        if weight is not None:
            """only one class"""
            if len(weight) == len(grad):
                grad = np.multiply(grad, weight)
                hess = np.multiply(hess, weight)
            else:
                num_data = len(weight)
                num_class = len(grad) // num_data
                if num_class * num_data != len(grad):
97
                    raise ValueError("Length of grad and hess should equal to num_class * num_data")
98
99
                for k in range(num_class):
                    for i in range(num_data):
wxchan's avatar
wxchan committed
100
101
102
103
104
                        idx = k * num_data + i
                        grad[idx] *= weight[i]
                        hess[idx] *= weight[i]
        return grad, hess

wxchan's avatar
wxchan committed
105

106
class _EvalFunctionWrapper:
107
    """Proxy class for evaluation function."""
108

109
110
    def __init__(self, func):
        """Construct a proxy class.
111

112
113
        This class transforms evaluation function to match evaluation function with signature ``new_func(preds, dataset)``
        as expected by ``lightgbm.engine.train``.
114

115
116
117
118
119
120
121
122
123
124
125
126
127
128
        Parameters
        ----------
        func : callable
            Expects a callable with following signatures:
            ``func(y_true, y_pred)``,
            ``func(y_true, y_pred, weight)``
            or ``func(y_true, y_pred, weight, group)``
            and returns (eval_name, eval_result, is_higher_better) or
            list of (eval_name, eval_result, is_higher_better):

                y_true : array-like of shape = [n_samples]
                    The target values.
                y_pred : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
                    The predicted values.
129
130
                    In case of custom ``objective``, predicted values are returned before any transformation,
                    e.g. they are raw margin instead of probability of positive class for binary task in this case.
131
132
133
                weight : array-like of shape = [n_samples]
                    The weight of samples.
                group : array-like
134
135
136
                    Group/query data.
                    Only used in the learning-to-rank task.
                    sum(group) = n_samples.
137
138
                    For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
                    where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
139
                eval_name : str
Andrew Ziem's avatar
Andrew Ziem committed
140
                    The name of evaluation function (without whitespace).
141
142
143
144
145
                eval_result : float
                    The eval result.
                is_higher_better : bool
                    Is eval result higher better, e.g. AUC is ``is_higher_better``.

Nikita Titov's avatar
Nikita Titov committed
146
147
148
149
        .. note::

            For multi-class task, the y_pred is group by class_id first, then group by row_id.
            If you want to get i-th row y_pred in j-th class, the access way is y_pred[j * num_data + i].
150
151
        """
        self.func = func
152

153
154
    def __call__(self, preds, dataset):
        """Call passed function with appropriate arguments.
155

156
157
158
159
160
161
162
163
164
        Parameters
        ----------
        preds : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
            The predicted values.
        dataset : Dataset
            The training dataset.

        Returns
        -------
165
        eval_name : str
Andrew Ziem's avatar
Andrew Ziem committed
166
            The name of evaluation function (without whitespace).
167
168
169
170
171
        eval_result : float
            The eval result.
        is_higher_better : bool
            Is eval result higher better, e.g. AUC is ``is_higher_better``.
        """
172
        labels = dataset.get_label()
173
        argc = len(signature(self.func).parameters)
174
        if argc == 2:
175
            return self.func(labels, preds)
176
        elif argc == 3:
177
            return self.func(labels, preds, dataset.get_weight())
178
        elif argc == 4:
179
            return self.func(labels, preds, dataset.get_weight(), dataset.get_group())
180
        else:
181
            raise TypeError(f"Self-defined eval function should have 2, 3 or 4 arguments, got {argc}")
182

wxchan's avatar
wxchan committed
183

184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
# documentation templates for LGBMModel methods are shared between the classes in
# this module and those in the ``dask`` module

_lgbmmodel_doc_fit = (
    """
    Build a gradient boosting model from the training set (X, y).

    Parameters
    ----------
    X : {X_shape}
        Input feature matrix.
    y : {y_shape}
        The target values (class labels in classification, real numbers in regression).
    sample_weight : {sample_weight_shape}
        Weights of training data.
199
    init_score : {init_score_shape}
200
201
202
203
204
205
206
207
208
        Init score of training data.
    group : {group_shape}
        Group/query data.
        Only used in the learning-to-rank task.
        sum(group) = n_samples.
        For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
        where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
    eval_set : list or None, optional (default=None)
        A list of (X, y) tuple pairs to use as validation sets.
209
    eval_names : list of str, or None, optional (default=None)
210
        Names of eval_set.
211
    eval_sample_weight : {eval_sample_weight_shape}
212
213
214
        Weights of eval data.
    eval_class_weight : list or None, optional (default=None)
        Class weights of eval data.
215
    eval_init_score : {eval_init_score_shape}
216
        Init score of eval data.
217
    eval_group : {eval_group_shape}
218
        Group data of eval data.
219
220
    eval_metric : str, callable, list or None, optional (default=None)
        If str, it should be a built-in evaluation metric to use.
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
        If callable, it should be a custom evaluation metric, see note below for more details.
        If list, it can be a list of built-in metrics, a list of custom evaluation metrics, or a mix of both.
        In either case, the ``metric`` from the model parameters will be evaluated and used as well.
        Default: 'l2' for LGBMRegressor, 'logloss' for LGBMClassifier, 'ndcg' for LGBMRanker.
    early_stopping_rounds : int or None, optional (default=None)
        Activates early stopping. The model will train until the validation score stops improving.
        Validation score needs to improve at least every ``early_stopping_rounds`` round(s)
        to continue training.
        Requires at least one validation data and one metric.
        If there's more than one, will check all of them. But the training data is ignored anyway.
        To check only the first metric, set the ``first_metric_only`` parameter to ``True``
        in additional parameters ``**kwargs`` of the model constructor.
    verbose : bool or int, optional (default=True)
        Requires at least one evaluation data.
        If True, the eval metric on the eval set is printed at each boosting stage.
        If int, the eval metric on the eval set is printed at every ``verbose`` boosting stage.
        The last boosting stage or the boosting stage found by using ``early_stopping_rounds`` is also printed.

        .. rubric:: Example

        With ``verbose`` = 4 and at least one item in ``eval_set``,
        an evaluation metric is printed every 4 (instead of 1) boosting stages.

244
    feature_name : list of str, or 'auto', optional (default='auto')
245
246
        Feature names.
        If 'auto' and data is pandas DataFrame, data columns names are used.
247
    categorical_feature : list of str or int, or 'auto', optional (default='auto')
248
249
        Categorical features.
        If list of int, interpreted as indices.
250
        If list of str, interpreted as feature names (need to specify ``feature_name`` as well).
251
252
253
254
255
        If 'auto' and data is pandas DataFrame, pandas unordered categorical columns are used.
        All values in categorical features should be less than int32 max value (2147483647).
        Large values could be memory consuming. Consider using consecutive integers starting from zero.
        All negative values in categorical features will be treated as missing values.
        The output cannot be monotonically constrained with respect to a categorical feature.
256
    callbacks : list of callable, or None, optional (default=None)
257
258
        List of callback functions that are applied at each iteration.
        See Callbacks in Python API for more information.
259
    init_model : str, pathlib.Path, Booster, LGBMModel or None, optional (default=None)
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
        Filename of LightGBM model, Booster instance or LGBMModel instance used for continue training.

    Returns
    -------
    self : object
        Returns self.
    """
)

_lgbmmodel_doc_custom_eval_note = """
    Note
    ----
    Custom eval function expects a callable with following signatures:
    ``func(y_true, y_pred)``, ``func(y_true, y_pred, weight)`` or
    ``func(y_true, y_pred, weight, group)``
    and returns (eval_name, eval_result, is_higher_better) or
    list of (eval_name, eval_result, is_higher_better):

        y_true : array-like of shape = [n_samples]
            The target values.
        y_pred : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
            The predicted values.
282
283
            In case of custom ``objective``, predicted values are returned before any transformation,
            e.g. they are raw margin instead of probability of positive class for binary task in this case.
284
285
286
287
288
289
290
291
        weight : array-like of shape = [n_samples]
            The weight of samples.
        group : array-like
            Group/query data.
            Only used in the learning-to-rank task.
            sum(group) = n_samples.
            For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
            where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
292
        eval_name : str
Andrew Ziem's avatar
Andrew Ziem committed
293
            The name of evaluation function (without whitespace).
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
        eval_result : float
            The eval result.
        is_higher_better : bool
            Is eval result higher better, e.g. AUC is ``is_higher_better``.

    For multi-class task, the y_pred is group by class_id first, then group by row_id.
    If you want to get i-th row y_pred in j-th class, the access way is y_pred[j * num_data + i].
"""

_lgbmmodel_doc_predict = (
    """
    {description}

    Parameters
    ----------
    X : {X_shape}
        Input features matrix.
    raw_score : bool, optional (default=False)
        Whether to predict raw scores.
    start_iteration : int, optional (default=0)
        Start index of the iteration to predict.
        If <= 0, starts from the first iteration.
    num_iteration : int or None, optional (default=None)
        Total number of iterations used in the prediction.
        If None, if the best iteration exists and start_iteration <= 0, the best iteration is used;
        otherwise, all iterations from ``start_iteration`` are used (no limits).
        If <= 0, all iterations from ``start_iteration`` are used (no limits).
    pred_leaf : bool, optional (default=False)
        Whether to predict leaf index.
    pred_contrib : bool, optional (default=False)
        Whether to predict feature contributions.

        .. note::

            If you want to get more explanations for your model's predictions using SHAP values,
            like SHAP interaction values,
            you can install the shap package (https://github.com/slundberg/shap).
            Note that unlike the shap package, with ``pred_contrib`` we return a matrix with an extra
            column, where the last column is the expected value.

    **kwargs
        Other parameters for the prediction.

    Returns
    -------
    {output_name} : {predicted_result_shape}
        The predicted values.
    X_leaves : {X_leaves_shape}
        If ``pred_leaf=True``, the predicted leaf of every tree for each sample.
    X_SHAP_values : {X_SHAP_values_shape}
        If ``pred_contrib=True``, the feature contributions for each sample.
    """
)


349
350
class LGBMModel(_LGBMModelBase):
    """Implementation of the scikit-learn API for LightGBM."""
wxchan's avatar
wxchan committed
351

352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
    def __init__(
        self,
        boosting_type: str = 'gbdt',
        num_leaves: int = 31,
        max_depth: int = -1,
        learning_rate: float = 0.1,
        n_estimators: int = 100,
        subsample_for_bin: int = 200000,
        objective: Optional[Union[str, Callable]] = None,
        class_weight: Optional[Union[Dict, str]] = None,
        min_split_gain: float = 0.,
        min_child_weight: float = 1e-3,
        min_child_samples: int = 20,
        subsample: float = 1.,
        subsample_freq: int = 0,
        colsample_bytree: float = 1.,
        reg_alpha: float = 0.,
        reg_lambda: float = 0.,
        random_state: Optional[Union[int, np.random.RandomState]] = None,
        n_jobs: int = -1,
372
        silent: Union[bool, str] = 'warn',
373
374
375
        importance_type: str = 'split',
        **kwargs
    ):
376
        r"""Construct a gradient boosting model.
wxchan's avatar
wxchan committed
377
378
379

        Parameters
        ----------
380
        boosting_type : str, optional (default='gbdt')
381
382
383
384
385
            'gbdt', traditional Gradient Boosting Decision Tree.
            'dart', Dropouts meet Multiple Additive Regression Trees.
            'goss', Gradient-based One-Side Sampling.
            'rf', Random Forest.
        num_leaves : int, optional (default=31)
wxchan's avatar
wxchan committed
386
            Maximum tree leaves for base learners.
387
        max_depth : int, optional (default=-1)
388
            Maximum tree depth for base learners, <=0 means no limit.
389
        learning_rate : float, optional (default=0.1)
390
            Boosting learning rate.
391
392
393
            You can use ``callbacks`` parameter of ``fit`` method to shrink/adapt learning rate
            in training using ``reset_parameter`` callback.
            Note, that this will ignore the ``learning_rate`` argument in training.
394
        n_estimators : int, optional (default=100)
wxchan's avatar
wxchan committed
395
            Number of boosted trees to fit.
396
        subsample_for_bin : int, optional (default=200000)
wxchan's avatar
wxchan committed
397
            Number of samples for constructing bins.
398
        objective : str, callable or None, optional (default=None)
wxchan's avatar
wxchan committed
399
400
            Specify the learning task and the corresponding learning objective or
            a custom objective function to be used (see note below).
401
            Default: 'regression' for LGBMRegressor, 'binary' or 'multiclass' for LGBMClassifier, 'lambdarank' for LGBMRanker.
402
403
404
405
        class_weight : dict, 'balanced' or None, optional (default=None)
            Weights associated with classes in the form ``{class_label: weight}``.
            Use this parameter only for multi-class classification task;
            for binary classification task you may use ``is_unbalance`` or ``scale_pos_weight`` parameters.
406
407
408
            Note, that the usage of all these parameters will result in poor estimates of the individual class probabilities.
            You may want to consider performing probability calibration
            (https://scikit-learn.org/stable/modules/calibration.html) of your model.
409
410
411
            The 'balanced' mode uses the values of y to automatically adjust weights
            inversely proportional to class frequencies in the input data as ``n_samples / (n_classes * np.bincount(y))``.
            If None, all classes are supposed to have weight one.
412
            Note, that these weights will be multiplied with ``sample_weight`` (passed through the ``fit`` method)
413
            if ``sample_weight`` is specified.
414
        min_split_gain : float, optional (default=0.)
wxchan's avatar
wxchan committed
415
            Minimum loss reduction required to make a further partition on a leaf node of the tree.
416
        min_child_weight : float, optional (default=1e-3)
417
            Minimum sum of instance weight (hessian) needed in a child (leaf).
418
        min_child_samples : int, optional (default=20)
419
            Minimum number of data needed in a child (leaf).
420
        subsample : float, optional (default=1.)
wxchan's avatar
wxchan committed
421
            Subsample ratio of the training instance.
422
        subsample_freq : int, optional (default=0)
Andrew Ziem's avatar
Andrew Ziem committed
423
            Frequency of subsample, <=0 means no enable.
424
        colsample_bytree : float, optional (default=1.)
wxchan's avatar
wxchan committed
425
            Subsample ratio of columns when constructing each tree.
426
        reg_alpha : float, optional (default=0.)
427
            L1 regularization term on weights.
428
        reg_lambda : float, optional (default=0.)
429
            L2 regularization term on weights.
430
        random_state : int, RandomState object or None, optional (default=None)
wxchan's avatar
wxchan committed
431
            Random number seed.
432
433
434
            If int, this number is used to seed the C++ code.
            If RandomState object (numpy), a random integer is picked based on its state to seed the C++ code.
            If None, default seeds in C++ code are used.
435
        n_jobs : int, optional (default=-1)
436
            Number of parallel threads.
437
        silent : bool, optional (default=True)
wxchan's avatar
wxchan committed
438
            Whether to print messages while running boosting.
439
        importance_type : str, optional (default='split')
440
            The type of feature importance to be filled into ``feature_importances_``.
441
442
443
444
            If 'split', result contains numbers of times the feature is used in a model.
            If 'gain', result contains total gains of splits which use the feature.
        **kwargs
            Other parameters for the model.
wxchan's avatar
wxchan committed
445
            Check http://lightgbm.readthedocs.io/en/latest/Parameters.html for more parameters.
446

Nikita Titov's avatar
Nikita Titov committed
447
448
449
            .. warning::

                \*\*kwargs is not supported in sklearn, it may cause unexpected issues.
wxchan's avatar
wxchan committed
450
451
452

        Note
        ----
453
454
        A custom objective function can be provided for the ``objective`` parameter.
        In this case, it should have the signature
455
456
        ``objective(y_true, y_pred) -> grad, hess`` or
        ``objective(y_true, y_pred, group) -> grad, hess``:
wxchan's avatar
wxchan committed
457

Nikita Titov's avatar
Nikita Titov committed
458
            y_true : array-like of shape = [n_samples]
459
                The target values.
Nikita Titov's avatar
Nikita Titov committed
460
            y_pred : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
461
                The predicted values.
462
463
                Predicted values are returned before any transformation,
                e.g. they are raw margin instead of probability of positive class for binary task.
Nikita Titov's avatar
Nikita Titov committed
464
            group : array-like
465
466
467
                Group/query data.
                Only used in the learning-to-rank task.
                sum(group) = n_samples.
468
469
                For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
                where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
Nikita Titov's avatar
Nikita Titov committed
470
            grad : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
471
472
                The value of the first order derivative (gradient) of the loss
                with respect to the elements of y_pred for each sample point.
Nikita Titov's avatar
Nikita Titov committed
473
            hess : array-like of shape = [n_samples] or shape = [n_samples * n_classes] (for multi-class task)
474
475
                The value of the second order derivative (Hessian) of the loss
                with respect to the elements of y_pred for each sample point.
wxchan's avatar
wxchan committed
476

477
478
479
        For multi-class task, the y_pred is group by class_id first, then group by row_id.
        If you want to get i-th row y_pred in j-th class, the access way is y_pred[j * num_data + i]
        and you should group grad and hess in this way as well.
wxchan's avatar
wxchan committed
480
        """
wxchan's avatar
wxchan committed
481
        if not SKLEARN_INSTALLED:
482
483
            raise LightGBMError('scikit-learn is required for lightgbm.sklearn. '
                                'You must install scikit-learn and restart your session to use this module.')
wxchan's avatar
wxchan committed
484

485
        self.boosting_type = boosting_type
486
        self.objective = objective
wxchan's avatar
wxchan committed
487
488
489
490
        self.num_leaves = num_leaves
        self.max_depth = max_depth
        self.learning_rate = learning_rate
        self.n_estimators = n_estimators
wxchan's avatar
wxchan committed
491
        self.subsample_for_bin = subsample_for_bin
wxchan's avatar
wxchan committed
492
493
494
495
496
497
498
499
        self.min_split_gain = min_split_gain
        self.min_child_weight = min_child_weight
        self.min_child_samples = min_child_samples
        self.subsample = subsample
        self.subsample_freq = subsample_freq
        self.colsample_bytree = colsample_bytree
        self.reg_alpha = reg_alpha
        self.reg_lambda = reg_lambda
500
501
        self.random_state = random_state
        self.n_jobs = n_jobs
wxchan's avatar
wxchan committed
502
        self.silent = silent
503
        self.importance_type = importance_type
wxchan's avatar
wxchan committed
504
        self._Booster = None
505
506
507
508
        self._evals_result = None
        self._best_score = None
        self._best_iteration = None
        self._other_params = {}
509
        self._objective = objective
510
        self.class_weight = class_weight
511
512
        self._class_weight = None
        self._class_map = None
513
        self._n_features = None
514
        self._n_features_in = None
515
516
        self._classes = None
        self._n_classes = None
517
        self.set_params(**kwargs)
wxchan's avatar
wxchan committed
518

Nikita Titov's avatar
Nikita Titov committed
519
    def _more_tags(self):
520
521
522
523
524
525
526
527
528
529
        return {
            'allow_nan': True,
            'X_types': ['2darray', 'sparse', '1dlabels'],
            '_xfail_checks': {
                'check_no_attributes_set_in_init':
                'scikit-learn incorrectly asserts that private attributes '
                'cannot be set in __init__: '
                '(see https://github.com/microsoft/LightGBM/issues/2628)'
            }
        }
Nikita Titov's avatar
Nikita Titov committed
530

wxchan's avatar
wxchan committed
531
    def get_params(self, deep=True):
532
533
534
535
536
537
538
539
540
541
542
543
544
        """Get parameters for this estimator.

        Parameters
        ----------
        deep : bool, optional (default=True)
            If True, will return the parameters for this estimator and
            contained subobjects that are estimators.

        Returns
        -------
        params : dict
            Parameter names mapped to their values.
        """
545
        params = super().get_params(deep=deep)
546
        params.update(self._other_params)
wxchan's avatar
wxchan committed
547
548
549
        return params

    def set_params(self, **params):
550
551
552
553
554
555
556
557
558
559
560
561
        """Set the parameters of this estimator.

        Parameters
        ----------
        **params
            Parameter names with their new values.

        Returns
        -------
        self : object
            Returns self.
        """
wxchan's avatar
wxchan committed
562
563
        for key, value in params.items():
            setattr(self, key, value)
564
565
            if hasattr(self, f"_{key}"):
                setattr(self, f"_{key}", value)
566
            self._other_params[key] = value
wxchan's avatar
wxchan committed
567
        return self
wxchan's avatar
wxchan committed
568

Guolin Ke's avatar
Guolin Ke committed
569
    def fit(self, X, y,
570
            sample_weight=None, init_score=None, group=None,
571
            eval_set=None, eval_names=None, eval_sample_weight=None,
572
573
            eval_class_weight=None, eval_init_score=None, eval_group=None,
            eval_metric=None, early_stopping_rounds=None, verbose=True,
574
575
            feature_name='auto', categorical_feature='auto',
            callbacks=None, init_model=None):
576
        """Docstring is set after definition, using a template."""
577
578
579
580
581
582
583
584
585
586
        if self._objective is None:
            if isinstance(self, LGBMRegressor):
                self._objective = "regression"
            elif isinstance(self, LGBMClassifier):
                self._objective = "binary"
            elif isinstance(self, LGBMRanker):
                self._objective = "lambdarank"
            else:
                raise ValueError("Unknown LGBMModel type.")
        if callable(self._objective):
587
            self._fobj = _ObjectiveFunctionWrapper(self._objective)
588
589
        else:
            self._fobj = None
wxchan's avatar
wxchan committed
590
591
        evals_result = {}
        params = self.get_params()
wxchan's avatar
wxchan committed
592
        # user can set verbose with kwargs, it has higher priority
593
594
595
596
597
598
599
        if self.silent != "warn":
            _log_warning("'silent' argument is deprecated and will be removed in a future release of LightGBM. "
                         "Pass 'verbose' parameter via keyword arguments instead.")
            silent = self.silent
        else:
            silent = True
        if not any(verbose_alias in params for verbose_alias in _ConfigAliases.get("verbosity")) and silent:
600
            params['verbose'] = -1
wxchan's avatar
wxchan committed
601
        params.pop('silent', None)
602
        params.pop('importance_type', None)
wxchan's avatar
wxchan committed
603
        params.pop('n_estimators', None)
604
        params.pop('class_weight', None)
605
606
        if isinstance(params['random_state'], np.random.RandomState):
            params['random_state'] = params['random_state'].randint(np.iinfo(np.int32).max)
607
608
        for alias in _ConfigAliases.get('objective'):
            params.pop(alias, None)
609
        if self._n_classes is not None and self._n_classes > 2:
610
611
            for alias in _ConfigAliases.get('num_class'):
                params.pop(alias, None)
612
613
            params['num_class'] = self._n_classes
        if hasattr(self, '_eval_at'):
614
            eval_at = self._eval_at
615
            for alias in _ConfigAliases.get('eval_at'):
616
617
618
619
                if alias in params:
                    _log_warning(f"Found '{alias}' in params. Will use it instead of 'eval_at' argument")
                    eval_at = params.pop(alias)
            params['eval_at'] = eval_at
620
621
        params['objective'] = self._objective
        if self._fobj:
wxchan's avatar
wxchan committed
622
            params['objective'] = 'None'  # objective = nullptr for unknown objective
wxchan's avatar
wxchan committed
623

624
625
626
627
628
629
630
631
        # Do not modify original args in fit function
        # Refer to https://github.com/microsoft/LightGBM/pull/2619
        eval_metric_list = copy.deepcopy(eval_metric)
        if not isinstance(eval_metric_list, list):
            eval_metric_list = [eval_metric_list]

        # Separate built-in from callable evaluation metrics
        eval_metrics_callable = [_EvalFunctionWrapper(f) for f in eval_metric_list if callable(f)]
632
        eval_metrics_builtin = [m for m in eval_metric_list if isinstance(m, str)]
633
634

        # register default metric for consistency with callable eval_metric case
635
        original_metric = self._objective if isinstance(self._objective, str) else None
636
637
638
639
640
641
642
643
644
645
        if original_metric is None:
            # try to deduce from class instance
            if isinstance(self, LGBMRegressor):
                original_metric = "l2"
            elif isinstance(self, LGBMClassifier):
                original_metric = "multi_logloss" if self._n_classes > 2 else "binary_logloss"
            elif isinstance(self, LGBMRanker):
                original_metric = "ndcg"

        # overwrite default metric by explicitly set metric
646
        params = _choose_param_value("metric", params, original_metric)
647
648

        # concatenate metric from params (or default if not provided in params) and eval_metric
649
650
        params['metric'] = [params['metric']] if isinstance(params['metric'], (str, type(None))) else params['metric']
        params['metric'] = [e for e in eval_metrics_builtin if e not in params['metric']] + params['metric']
651
        params['metric'] = [metric for metric in params['metric'] if metric is not None]
wxchan's avatar
wxchan committed
652

653
        if not isinstance(X, (pd_DataFrame, dt_DataTable)):
654
            _X, _y = _LGBMCheckXY(X, y, accept_sparse=True, force_all_finite=False, ensure_min_samples=2)
655
656
            if sample_weight is not None:
                sample_weight = _LGBMCheckSampleWeight(sample_weight, _X)
657
658
        else:
            _X, _y = X, y
659

660
661
662
663
        if self._class_weight is None:
            self._class_weight = self.class_weight
        if self._class_weight is not None:
            class_sample_weight = _LGBMComputeSampleWeight(self._class_weight, y)
664
665
666
667
            if sample_weight is None or len(sample_weight) == 0:
                sample_weight = class_sample_weight
            else:
                sample_weight = np.multiply(sample_weight, class_sample_weight)
668

669
        self._n_features = _X.shape[1]
670
671
        # copy for consistency
        self._n_features_in = self._n_features
672

673
674
        def _construct_dataset(X, y, sample_weight, init_score, group, params,
                               categorical_feature='auto'):
675
            return Dataset(X, label=y, weight=sample_weight, group=group,
676
677
                           init_score=init_score, params=params,
                           categorical_feature=categorical_feature)
Guolin Ke's avatar
Guolin Ke committed
678

679
680
        train_set = _construct_dataset(_X, _y, sample_weight, init_score, group, params,
                                       categorical_feature=categorical_feature)
Guolin Ke's avatar
Guolin Ke committed
681
682
683

        valid_sets = []
        if eval_set is not None:
684

685
            def _get_meta_data(collection, name, i):
686
687
688
689
690
691
692
                if collection is None:
                    return None
                elif isinstance(collection, list):
                    return collection[i] if len(collection) > i else None
                elif isinstance(collection, dict):
                    return collection.get(i, None)
                else:
693
                    raise TypeError(f"{name} should be dict or list")
694

Guolin Ke's avatar
Guolin Ke committed
695
696
697
            if isinstance(eval_set, tuple):
                eval_set = [eval_set]
            for i, valid_data in enumerate(eval_set):
698
                # reduce cost for prediction training data
Guolin Ke's avatar
Guolin Ke committed
699
700
701
                if valid_data[0] is X and valid_data[1] is y:
                    valid_set = train_set
                else:
702
703
704
705
706
707
                    valid_weight = _get_meta_data(eval_sample_weight, 'eval_sample_weight', i)
                    valid_class_weight = _get_meta_data(eval_class_weight, 'eval_class_weight', i)
                    if valid_class_weight is not None:
                        if isinstance(valid_class_weight, dict) and self._class_map is not None:
                            valid_class_weight = {self._class_map[k]: v for k, v in valid_class_weight.items()}
                        valid_class_sample_weight = _LGBMComputeSampleWeight(valid_class_weight, valid_data[1])
708
709
710
711
                        if valid_weight is None or len(valid_weight) == 0:
                            valid_weight = valid_class_sample_weight
                        else:
                            valid_weight = np.multiply(valid_weight, valid_class_sample_weight)
712
713
                    valid_init_score = _get_meta_data(eval_init_score, 'eval_init_score', i)
                    valid_group = _get_meta_data(eval_group, 'eval_group', i)
714
715
                    valid_set = _construct_dataset(valid_data[0], valid_data[1],
                                                   valid_weight, valid_init_score, valid_group, params)
Guolin Ke's avatar
Guolin Ke committed
716
717
                valid_sets.append(valid_set)

718
719
720
        if isinstance(init_model, LGBMModel):
            init_model = init_model.booster_

Guolin Ke's avatar
Guolin Ke committed
721
        self._Booster = train(params, train_set,
722
                              self.n_estimators, valid_sets=valid_sets, valid_names=eval_names,
wxchan's avatar
wxchan committed
723
                              early_stopping_rounds=early_stopping_rounds,
724
                              evals_result=evals_result, fobj=self._fobj, feval=eval_metrics_callable,
Guolin Ke's avatar
Guolin Ke committed
725
                              verbose_eval=verbose, feature_name=feature_name,
726
                              callbacks=callbacks, init_model=init_model)
wxchan's avatar
wxchan committed
727
728

        if evals_result:
729
            self._evals_result = evals_result
wxchan's avatar
wxchan committed
730

731
        if early_stopping_rounds is not None and early_stopping_rounds > 0:
732
            self._best_iteration = self._Booster.best_iteration
733
734

        self._best_score = self._Booster.best_score
wxchan's avatar
wxchan committed
735

736
737
        self.fitted_ = True

wxchan's avatar
wxchan committed
738
        # free dataset
739
        self._Booster.free_dataset()
wxchan's avatar
wxchan committed
740
        del train_set, valid_sets
wxchan's avatar
wxchan committed
741
742
        return self

743
744
745
746
    fit.__doc__ = _lgbmmodel_doc_fit.format(
        X_shape="array-like or sparse matrix of shape = [n_samples, n_features]",
        y_shape="array-like of shape = [n_samples]",
        sample_weight_shape="array-like of shape = [n_samples] or None, optional (default=None)",
747
        init_score_shape="array-like of shape = [n_samples] or None, optional (default=None)",
748
        group_shape="array-like or None, optional (default=None)",
749
750
751
        eval_sample_weight_shape="list of array, or None, optional (default=None)",
        eval_init_score_shape="list of array, or None, optional (default=None)",
        eval_group_shape="list of array, or None, optional (default=None)"
752
753
    ) + "\n\n" + _lgbmmodel_doc_custom_eval_note

754
    def predict(self, X, raw_score=False, start_iteration=0, num_iteration=None,
755
                pred_leaf=False, pred_contrib=False, **kwargs):
756
        """Docstring is set after definition, using a template."""
757
        if self._n_features is None:
758
            raise LGBMNotFittedError("Estimator not fitted, call fit before exploiting the model.")
759
        if not isinstance(X, (pd_DataFrame, dt_DataTable)):
760
            X = _LGBMCheckArray(X, accept_sparse=True, force_all_finite=False)
761
762
763
        n_features = X.shape[1]
        if self._n_features != n_features:
            raise ValueError("Number of features of the model must "
764
765
                             f"match the input. Model n_features_ is {self._n_features} and "
                             f"input n_features is {n_features}")
766
        return self._Booster.predict(X, raw_score=raw_score, start_iteration=start_iteration, num_iteration=num_iteration,
767
                                     pred_leaf=pred_leaf, pred_contrib=pred_contrib, **kwargs)
wxchan's avatar
wxchan committed
768

769
770
771
772
773
774
775
776
777
    predict.__doc__ = _lgbmmodel_doc_predict.format(
        description="Return the predicted value for each sample.",
        X_shape="array-like or sparse matrix of shape = [n_samples, n_features]",
        output_name="predicted_result",
        predicted_result_shape="array-like of shape = [n_samples] or shape = [n_samples, n_classes]",
        X_leaves_shape="array-like of shape = [n_samples, n_trees] or shape = [n_samples, n_trees * n_classes]",
        X_SHAP_values_shape="array-like of shape = [n_samples, n_features + 1] or shape = [n_samples, (n_features + 1) * n_classes] or list with n_classes length of such objects"
    )

778
779
    @property
    def n_features_(self):
780
        """:obj:`int`: The number of features of fitted model."""
781
782
783
784
        if self._n_features is None:
            raise LGBMNotFittedError('No n_features found. Need to call fit beforehand.')
        return self._n_features

785
786
787
788
789
790
791
    @property
    def n_features_in_(self):
        """:obj:`int`: The number of features of fitted model."""
        if self._n_features_in is None:
            raise LGBMNotFittedError('No n_features_in found. Need to call fit beforehand.')
        return self._n_features_in

792
793
    @property
    def best_score_(self):
794
        """:obj:`dict` or :obj:`None`: The best score of fitted model."""
795
796
797
798
799
800
        if self._n_features is None:
            raise LGBMNotFittedError('No best_score found. Need to call fit beforehand.')
        return self._best_score

    @property
    def best_iteration_(self):
801
        """:obj:`int` or :obj:`None`: The best iteration of fitted model if ``early_stopping_rounds`` has been specified."""
802
803
804
805
806
807
        if self._n_features is None:
            raise LGBMNotFittedError('No best_iteration found. Need to call fit with early_stopping_rounds beforehand.')
        return self._best_iteration

    @property
    def objective_(self):
808
        """:obj:`str` or :obj:`callable`: The concrete objective used while fitting this model."""
809
810
811
812
        if self._n_features is None:
            raise LGBMNotFittedError('No objective found. Need to call fit beforehand.')
        return self._objective

813
814
    @property
    def booster_(self):
815
        """Booster: The underlying Booster of this model."""
816
        if self._Booster is None:
817
            raise LGBMNotFittedError('No booster found. Need to call fit beforehand.')
818
        return self._Booster
wxchan's avatar
wxchan committed
819

820
821
    @property
    def evals_result_(self):
822
        """:obj:`dict` or :obj:`None`: The evaluation results if ``early_stopping_rounds`` has been specified."""
823
824
825
        if self._n_features is None:
            raise LGBMNotFittedError('No results found. Need to call fit with eval_set beforehand.')
        return self._evals_result
826
827

    @property
828
    def feature_importances_(self):
829
        """:obj:`array` of shape = [n_features]: The feature importances (the higher, the more important).
830

Nikita Titov's avatar
Nikita Titov committed
831
832
833
834
        .. note::

            ``importance_type`` attribute is passed to the function
            to configure the type of importance values to be extracted.
835
        """
836
837
        if self._n_features is None:
            raise LGBMNotFittedError('No feature_importances found. Need to call fit beforehand.')
838
        return self._Booster.feature_importance(importance_type=self.importance_type)
wxchan's avatar
wxchan committed
839

840
841
    @property
    def feature_name_(self):
842
        """:obj:`array` of shape = [n_features]: The names of features."""
843
844
845
846
        if self._n_features is None:
            raise LGBMNotFittedError('No feature_name found. Need to call fit beforehand.')
        return self._Booster.feature_name()

wxchan's avatar
wxchan committed
847

848
class LGBMRegressor(_LGBMRegressorBase, LGBMModel):
849
    """LightGBM regressor."""
wxchan's avatar
wxchan committed
850

Guolin Ke's avatar
Guolin Ke committed
851
852
    def fit(self, X, y,
            sample_weight=None, init_score=None,
853
            eval_set=None, eval_names=None, eval_sample_weight=None,
854
            eval_init_score=None, eval_metric=None, early_stopping_rounds=None,
855
856
            verbose=True, feature_name='auto', categorical_feature='auto',
            callbacks=None, init_model=None):
857
        """Docstring is inherited from the LGBMModel."""
858
859
860
861
862
        super().fit(X, y, sample_weight=sample_weight, init_score=init_score,
                    eval_set=eval_set, eval_names=eval_names, eval_sample_weight=eval_sample_weight,
                    eval_init_score=eval_init_score, eval_metric=eval_metric,
                    early_stopping_rounds=early_stopping_rounds, verbose=verbose, feature_name=feature_name,
                    categorical_feature=categorical_feature, callbacks=callbacks, init_model=init_model)
Guolin Ke's avatar
Guolin Ke committed
863
864
        return self

865
    _base_doc = LGBMModel.fit.__doc__
866
867
    _base_doc = (_base_doc[:_base_doc.find('group :')]  # type: ignore
                 + _base_doc[_base_doc.find('eval_set :'):])  # type: ignore
868
869
870
871
    _base_doc = (_base_doc[:_base_doc.find('eval_class_weight :')]
                 + _base_doc[_base_doc.find('eval_init_score :'):])
    fit.__doc__ = (_base_doc[:_base_doc.find('eval_group :')]
                   + _base_doc[_base_doc.find('eval_metric :'):])
wxchan's avatar
wxchan committed
872

873

874
class LGBMClassifier(_LGBMClassifierBase, LGBMModel):
875
    """LightGBM classifier."""
wxchan's avatar
wxchan committed
876

Guolin Ke's avatar
Guolin Ke committed
877
878
    def fit(self, X, y,
            sample_weight=None, init_score=None,
879
            eval_set=None, eval_names=None, eval_sample_weight=None,
880
            eval_class_weight=None, eval_init_score=None, eval_metric=None,
wxchan's avatar
wxchan committed
881
            early_stopping_rounds=None, verbose=True,
882
883
            feature_name='auto', categorical_feature='auto',
            callbacks=None, init_model=None):
884
        """Docstring is inherited from the LGBMModel."""
885
        _LGBMAssertAllFinite(y)
886
887
        _LGBMCheckClassificationTargets(y)
        self._le = _LGBMLabelEncoder().fit(y)
888
        _y = self._le.transform(y)
889
        self._class_map = dict(zip(self._le.classes_, self._le.transform(self._le.classes_)))
890
891
        if isinstance(self.class_weight, dict):
            self._class_weight = {self._class_map[k]: v for k, v in self.class_weight.items()}
892

893
894
        self._classes = self._le.classes_
        self._n_classes = len(self._classes)
895

896
        if self._n_classes > 2:
wxchan's avatar
wxchan committed
897
            # Switch to using a multiclass objective in the underlying LGBM instance
898
            ova_aliases = {"multiclassova", "multiclass_ova", "ova", "ovr"}
899
            if self._objective not in ova_aliases and not callable(self._objective):
900
                self._objective = "multiclass"
901
902

        if not callable(eval_metric):
903
            if isinstance(eval_metric, (str, type(None))):
904
905
906
907
908
909
910
911
912
913
914
915
916
                eval_metric = [eval_metric]
            if self._n_classes > 2:
                for index, metric in enumerate(eval_metric):
                    if metric in {'logloss', 'binary_logloss'}:
                        eval_metric[index] = "multi_logloss"
                    elif metric in {'error', 'binary_error'}:
                        eval_metric[index] = "multi_error"
            else:
                for index, metric in enumerate(eval_metric):
                    if metric in {'logloss', 'multi_logloss'}:
                        eval_metric[index] = 'binary_logloss'
                    elif metric in {'error', 'multi_error'}:
                        eval_metric[index] = 'binary_error'
wxchan's avatar
wxchan committed
917

918
919
        # do not modify args, as it causes errors in model selection tools
        valid_sets = None
wxchan's avatar
wxchan committed
920
        if eval_set is not None:
921
922
            if isinstance(eval_set, tuple):
                eval_set = [eval_set]
923
            valid_sets = [None] * len(eval_set)
924
925
            for i, (valid_x, valid_y) in enumerate(eval_set):
                if valid_x is X and valid_y is y:
926
                    valid_sets[i] = (valid_x, _y)
927
                else:
928
                    valid_sets[i] = (valid_x, self._le.transform(valid_y))
929

930
931
932
933
934
935
        super().fit(X, _y, sample_weight=sample_weight, init_score=init_score, eval_set=valid_sets,
                    eval_names=eval_names, eval_sample_weight=eval_sample_weight,
                    eval_class_weight=eval_class_weight, eval_init_score=eval_init_score,
                    eval_metric=eval_metric, early_stopping_rounds=early_stopping_rounds,
                    verbose=verbose, feature_name=feature_name, categorical_feature=categorical_feature,
                    callbacks=callbacks, init_model=init_model)
wxchan's avatar
wxchan committed
936
937
        return self

938
    _base_doc = LGBMModel.fit.__doc__
939
940
    _base_doc = (_base_doc[:_base_doc.find('group :')]  # type: ignore
                 + _base_doc[_base_doc.find('eval_set :'):])  # type: ignore
941
942
    fit.__doc__ = (_base_doc[:_base_doc.find('eval_group :')]
                   + _base_doc[_base_doc.find('eval_metric :'):])
943

944
    def predict(self, X, raw_score=False, start_iteration=0, num_iteration=None,
945
                pred_leaf=False, pred_contrib=False, **kwargs):
946
        """Docstring is inherited from the LGBMModel."""
947
        result = self.predict_proba(X, raw_score, start_iteration, num_iteration,
948
                                    pred_leaf, pred_contrib, **kwargs)
949
        if callable(self._objective) or raw_score or pred_leaf or pred_contrib:
950
951
952
953
            return result
        else:
            class_index = np.argmax(result, axis=1)
            return self._le.inverse_transform(class_index)
wxchan's avatar
wxchan committed
954

955
956
    predict.__doc__ = LGBMModel.predict.__doc__

957
    def predict_proba(self, X, raw_score=False, start_iteration=0, num_iteration=None,
958
                      pred_leaf=False, pred_contrib=False, **kwargs):
959
        """Docstring is set after definition, using a template."""
960
        result = super().predict(X, raw_score, start_iteration, num_iteration, pred_leaf, pred_contrib, **kwargs)
961
        if callable(self._objective) and not (raw_score or pred_leaf or pred_contrib):
962
963
964
            _log_warning("Cannot compute class probabilities or labels "
                         "due to the usage of customized objective function.\n"
                         "Returning raw scores instead.")
965
966
            return result
        elif self._n_classes > 2 or raw_score or pred_leaf or pred_contrib:
967
            return result
wxchan's avatar
wxchan committed
968
        else:
969
            return np.vstack((1. - result, result)).transpose()
970

971
972
973
974
    predict_proba.__doc__ = _lgbmmodel_doc_predict.format(
        description="Return the predicted probability for each class for each sample.",
        X_shape="array-like or sparse matrix of shape = [n_samples, n_features]",
        output_name="predicted_probability",
975
        predicted_result_shape="array-like of shape = [n_samples] or shape = [n_samples, n_classes]",
976
977
978
979
        X_leaves_shape="array-like of shape = [n_samples, n_trees] or shape = [n_samples, n_trees * n_classes]",
        X_SHAP_values_shape="array-like of shape = [n_samples, n_features + 1] or shape = [n_samples, (n_features + 1) * n_classes] or list with n_classes length of such objects"
    )

980
981
    @property
    def classes_(self):
982
        """:obj:`array` of shape = [n_classes]: The class label array."""
983
984
985
        if self._classes is None:
            raise LGBMNotFittedError('No classes found. Need to call fit beforehand.')
        return self._classes
986
987
988

    @property
    def n_classes_(self):
989
        """:obj:`int`: The number of classes."""
990
991
992
        if self._n_classes is None:
            raise LGBMNotFittedError('No classes found. Need to call fit beforehand.')
        return self._n_classes
wxchan's avatar
wxchan committed
993

wxchan's avatar
wxchan committed
994

wxchan's avatar
wxchan committed
995
class LGBMRanker(LGBMModel):
996
997
998
999
1000
1001
1002
1003
    """LightGBM ranker.

    .. warning::

        scikit-learn doesn't support ranking applications yet,
        therefore this class is not really compatible with the sklearn ecosystem.
        Please use this class mainly for training and applying ranking models in common sklearnish way.
    """
wxchan's avatar
wxchan committed
1004

Guolin Ke's avatar
Guolin Ke committed
1005
    def fit(self, X, y,
1006
            sample_weight=None, init_score=None, group=None,
1007
            eval_set=None, eval_names=None, eval_sample_weight=None,
1008
            eval_init_score=None, eval_group=None, eval_metric=None,
1009
            eval_at=(1, 2, 3, 4, 5), early_stopping_rounds=None, verbose=True,
1010
1011
            feature_name='auto', categorical_feature='auto',
            callbacks=None, init_model=None):
1012
        """Docstring is inherited from the LGBMModel."""
1013
        # check group data
Guolin Ke's avatar
Guolin Ke committed
1014
        if group is None:
1015
            raise ValueError("Should set group for ranking task")
wxchan's avatar
wxchan committed
1016
1017

        if eval_set is not None:
Guolin Ke's avatar
Guolin Ke committed
1018
            if eval_group is None:
1019
                raise ValueError("Eval_group cannot be None when eval_set is not None")
Guolin Ke's avatar
Guolin Ke committed
1020
            elif len(eval_group) != len(eval_set):
1021
                raise ValueError("Length of eval_group should be equal to eval_set")
1022
            elif (isinstance(eval_group, dict)
1023
                  and any(i not in eval_group or eval_group[i] is None for i in range(len(eval_group)))
1024
1025
                  or isinstance(eval_group, list)
                  and any(group is None for group in eval_group)):
1026
1027
                raise ValueError("Should set group for all eval datasets for ranking task; "
                                 "if you use dict, the index should start from 0")
1028

1029
        self._eval_at = eval_at
1030
1031
1032
1033
1034
        super().fit(X, y, sample_weight=sample_weight, init_score=init_score, group=group,
                    eval_set=eval_set, eval_names=eval_names, eval_sample_weight=eval_sample_weight,
                    eval_init_score=eval_init_score, eval_group=eval_group, eval_metric=eval_metric,
                    early_stopping_rounds=early_stopping_rounds, verbose=verbose, feature_name=feature_name,
                    categorical_feature=categorical_feature, callbacks=callbacks, init_model=init_model)
wxchan's avatar
wxchan committed
1035
        return self
1036

1037
    _base_doc = LGBMModel.fit.__doc__
1038
1039
    fit.__doc__ = (_base_doc[:_base_doc.find('eval_class_weight :')]  # type: ignore
                   + _base_doc[_base_doc.find('eval_init_score :'):])  # type: ignore
1040
    _base_doc = fit.__doc__
1041
    _before_early_stop, _early_stop, _after_early_stop = _base_doc.partition('early_stopping_rounds :')
1042
1043
1044
    fit.__doc__ = f"""{_before_early_stop}eval_at : iterable of int, optional (default=(1, 2, 3, 4, 5))
        The evaluation positions of the specified metric.
    {_early_stop}{_after_early_stop}"""