[Feature] return probs in binary classification (#36)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "rektgbm"
-version = "0.1.4"
+version = "0.1.5"
 description = "No-brainer machine learning solution to achieve satisfactory performance"
 authors = ["RektPunk <[email protected]>"]
 license = "MIT"

rektgbm/__init__.py

@@ -3,4 +3,4 @@
 from rektgbm.gbm import RektGBM
 from rektgbm.optimizer import RektOptimizer
 
-__version__ = "0.1.4"
+__version__ = "0.1.5"

rektgbm/gbm.py

@@ -49,7 +49,7 @@ def fit(
         )
         if self._task_type in {TaskType.binary, TaskType.multiclass, TaskType.rank}:
             self.label_encoder = dataset.fit_transform_label()
-            self._label_encoder_used = True
+            self._is_label_encoder_used = True
 
         if valid_set is not None and self.__is_label_encoder_used:
             valid_set.transform_label(label_encoder=self.label_encoder)
@@ -66,13 +66,9 @@ def fit(
 
     def predict(self, dataset: RektDataset):
         preds = self.engine.predict(dataset=dataset)
-
-        if self._task_type in {TaskType.regression, TaskType.rank}:
+        if self._task_type in {TaskType.binary, TaskType.regression, TaskType.rank}:
             return preds
 
-        if self._task_type == TaskType.binary:
-            preds = np.around(preds).astype(int)
-
         if self._task_type == TaskType.multiclass:
             if self.method == MethodName.lightgbm:
                 preds = np.argmax(preds, axis=1).astype(int)
@@ -81,9 +77,8 @@ def predict(self, dataset: RektDataset):
 
         if self.__is_label_encoder_used:
             preds = self.label_encoder.inverse_transform(series=preds)
-
         return preds
 
     @property
     def __is_label_encoder_used(self) -> bool:
-        return getattr(self, "_label_encoder_used", False)
+        return getattr(self, "_is_label_encoder_used", False)

tests/test_gbm.py

@@ -76,7 +76,7 @@ def test_rektgbm_predict_binary(mock_dataset, mock_engine):
     gbm._is_fitted = True
 
     preds = gbm.predict(dataset=mock_dataset)
-    np.testing.assert_allclose(preds, [0, 1, 0, 1], rtol=1e-5)
+    np.testing.assert_allclose(preds, [0.1, 0.9, 0.2, 0.8], rtol=1e-5)
 
 
 def test_rektgbm_predict_multiclass(mock_dataset, mock_engine):