File size: 557 Bytes
5f58699
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
"""Probability calibration helpers."""

from __future__ import annotations

from typing import Any

import numpy as np
from sklearn.calibration import CalibratedClassifierCV


def fit_calibrator(
    estimator: Any,
    X: np.ndarray,
    y: np.ndarray,
    *,
    method: str = "isotonic",
    cv: int | str | None = "prefit",
) -> CalibratedClassifierCV:
    """Fit a ``CalibratedClassifierCV`` on top of a pre-trained estimator."""

    calibrator = CalibratedClassifierCV(estimator, method=method, cv=cv)
    calibrator.fit(X, y)
    return calibrator