using a function

Author: chkoar

imblearn/under_sampling/instance_hardness_threshold.py

@@ -3,20 +3,28 @@
 from __future__ import division, print_function
 
 import warnings
-
 from collections import Counter
 
 import numpy as np
-
+from six import string_types
 from sklearn.base import ClassifierMixin
 from sklearn.ensemble import RandomForestClassifier
-from sklearn.cross_validation import StratifiedKFold
-
-from six import string_types
 
 from ..base import BaseBinarySampler
 
 
+def _get_cv_splits(X, y, cv, random_state):
+    try:
+        from sklearn.model_selection import StratifiedKFold
+        cv_iterator = StratifiedKFold(
+            n_splits=cv, shuffle=False, random_state=random_state).split(X, y)
+    except:
+        from sklearn.cross_validation import StratifiedKFold
+        cv_iterator = StratifiedKFold(
+            y, n_folds=cv, shuffle=False, random_state=random_state)
+    return cv_iterator
+
+
 class InstanceHardnessThreshold(BaseBinarySampler):
     """Class to perform under-sampling based on the instance hardness
     threshold.
@@ -225,8 +233,7 @@ def _sample(self, X, y):
         """
 
         # Create the different folds
-        skf = StratifiedKFold(y, n_folds=self.cv, shuffle=False,
-                              random_state=self.random_state)
+        skf = _get_cv_splits(X, y, self.cv, self.random_state)
 
         probabilities = np.zeros(y.shape[0], dtype=float)