Merge branch 'feature/fix_regression_models' into dev

Author: maks-sh

sklift/models/models.py

@@ -355,7 +355,7 @@ def fit(self, X, y, treatment, estimator_trmnt_fit_params=None, estimator_ctrl_f
             if self._type_of_target == 'binary':
                 ddr_treatment = self.estimator_trmnt.predict_proba(X_ctrl)[:, 1]
             else:
-                ddr_treatment = self.estimator_trmnt.predict(X_ctrl)[:, 1]
+                ddr_treatment = self.estimator_trmnt.predict(X_ctrl)
 
             if isinstance(X_ctrl, np.ndarray):
                 X_ctrl_mod = np.column_stack((X_ctrl, ddr_treatment))
@@ -393,21 +393,29 @@ def predict(self, X):
                 X_mod = X.assign(ddr_control=self.ctrl_preds_)
             else:
                 raise TypeError("Expected numpy.ndarray or pandas.DataFrame, got %s" % type(X))
-            self.trmnt_preds_ = self.estimator_trmnt.predict_proba(X_mod)[:, 1]
+
+            if self._type_of_target == 'binary':
+                self.trmnt_preds_ = self.estimator_trmnt.predict_proba(X_mod)[:, 1]
+            else:
+                self.trmnt_preds_ = self.estimator_trmnt.predict(X_mod)
 
         elif self.method == 'ddr_treatment':
             if self._type_of_target == 'binary':
                 self.trmnt_preds_ = self.estimator_trmnt.predict_proba(X)[:, 1]
             else:
-                self.trmnt_preds_ = self.estimator_trmnt.predict_proba(X)[:, 1]
+                self.trmnt_preds_ = self.estimator_trmnt.predict(X)
 
             if isinstance(X, np.ndarray):
                 X_mod = np.column_stack((X, self.trmnt_preds_))
             elif isinstance(X, pd.DataFrame):
                 X_mod = X.assign(ddr_treatment=self.trmnt_preds_)
             else:
                 raise TypeError("Expected numpy.ndarray or pandas.DataFrame, got %s" % type(X))
-            self.ctrl_preds_ = self.estimator_ctrl.predict_proba(X_mod)[:, 1]
+
+            if self._type_of_target == 'binary':
+                self.ctrl_preds_ = self.estimator_ctrl.predict_proba(X_mod)[:, 1]
+            else:
+                self.ctrl_preds_ = self.estimator_ctrl.predict(X_mod)
 
         else:
             if self._type_of_target == 'binary':