MAINT: rebase Pipeline with sklearn (#486)

Author: glemaitre

imblearn/pipeline.py

@@ -316,7 +316,7 @@ def fit_resample(self, X, y=None, **fit_params):
             return last_step.fit_resample(Xt, yt, **fit_params)
 
     @if_delegate_has_method(delegate='_final_estimator')
-    def predict(self, X):
+    def predict(self, X, **predict_params):
         """Apply transformers/samplers to the data, and predict with the final
         estimator
 
@@ -326,6 +326,14 @@ def predict(self, X):
             Data to predict on. Must fulfill input requirements of first step
             of the pipeline.
 
+        **predict_params : dict of string -> object
+            Parameters to the ``predict`` called at the end of all
+            transformations in the pipeline. Note that while this may be
+            used to return uncertainties from some models with return_std
+            or return_cov, uncertainties that are generated by the
+            transformations in the pipeline are not propagated to the
+            final estimator.
+
         Returns
         -------
         y_pred : array-like
@@ -339,7 +347,7 @@ def predict(self, X):
                 pass
             else:
                 Xt = transform.transform(Xt)
-        return self.steps[-1][-1].predict(Xt)
+        return self.steps[-1][-1].predict(Xt, **predict_params)
 
     @if_delegate_has_method(delegate='_final_estimator')
     def fit_predict(self, X, y=None, **fit_params):

imblearn/tests/test_pipeline.py

@@ -139,6 +139,16 @@ def fit(self, X, y):
         return self
 
 
+class DummyEstimatorParams(BaseEstimator):
+    """Mock classifier that takes params on predict"""
+    def fit(self, X, y):
+        return self
+
+    def predict(self, X, got_attribute=False):
+        self.got_attribute = got_attribute
+        return self
+
+
 class DummySampler(NoTrans):
     """Samplers which returns a balanced number of samples"""
 
@@ -1085,3 +1095,12 @@ def test_make_pipeline_memory():
         assert pipeline.memory is None
     finally:
         shutil.rmtree(cachedir)
+
+
+def test_predict_with_predict_params():
+    # tests that Pipeline passes predict_params to the final estimator
+    # when predict is invoked
+    pipe = Pipeline([('transf', Transf()), ('clf', DummyEstimatorParams())])
+    pipe.fit(None, None)
+    pipe.predict(X=None, got_attribute=True)
+    assert pipe.named_steps['clf'].got_attribute