Calibrated the Regressor Model

bugbug/models/regressor.py

@@ -18,6 +18,7 @@
 
 from bugbug import bugzilla, commit_features, db, feature_cleanup, repository, utils
 from bugbug.model import CommitModel
+from bugbug.model_calibration import IsotonicRegressionCalibrator
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -51,6 +52,7 @@ class RegressorModel(CommitModel):
 
     def __init__(
         self,
+        calibration: bool = True,
         lemmatization: bool = False,
         interpretable: bool = True,
         use_finder: bool = False,
@@ -127,9 +129,14 @@ def __init__(
                 ("union", ColumnTransformer(column_transformers)),
             ]
         )
-
-        self.clf = xgboost.XGBClassifier(n_jobs=utils.get_physical_cpu_count())
-        self.clf.set_params(predictor="cpu_predictor")
+        base_clf = xgboost.XGBClassifier(n_jobs=utils.get_physical_cpu_count())
+        base_clf.set_params(predictor="cpu_predictor")
+        if calibration:
+            self.clf = IsotonicRegressionCalibrator(base_clf)
+            # This is a temporary workaround for the error : "Model type not yet supported by TreeExplainer"
+            self.calculate_importance = False
+        else:
+            self.clf = base_clf
 
     def get_labels(self):
         classes = {}