@@ -136,7 +136,7 @@ def test_classification(self):
136136 self .assertTrue (nb_model .predict (features [2 ]) <= 0 )
137137 self .assertTrue (nb_model .predict (features [3 ]) > 0 )
138138
139- categoricalFeaturesInfo = {0 : 3 } # feature 0 has 3 categories
139+ categoricalFeaturesInfo = {0 : 3 } # feature 0 has 3 categories
140140 dt_model = \
141141 DecisionTree .trainClassifier (rdd , numClasses = 2 ,
142142 categoricalFeaturesInfo = categoricalFeaturesInfo )
@@ -176,9 +176,10 @@ def test_regression(self):
176176 self .assertTrue (rr_model .predict (features [2 ]) <= 0 )
177177 self .assertTrue (rr_model .predict (features [3 ]) > 0 )
178178
179- categoricalFeaturesInfo = {0 : 2 } # feature 0 has 2 categories
179+ categoricalFeaturesInfo = {0 : 2 } # feature 0 has 2 categories
180180 dt_model = \
181- DecisionTree .trainRegressor (rdd , categoricalFeaturesInfo = categoricalFeaturesInfo )
181+ DecisionTree .trainRegressor (
182+ rdd , categoricalFeaturesInfo = categoricalFeaturesInfo )
182183 self .assertTrue (dt_model .predict (features [0 ]) <= 0 )
183184 self .assertTrue (dt_model .predict (features [1 ]) > 0 )
184185 self .assertTrue (dt_model .predict (features [2 ]) <= 0 )
@@ -290,7 +291,7 @@ def test_classification(self):
290291 self .assertTrue (nb_model .predict (features [2 ]) <= 0 )
291292 self .assertTrue (nb_model .predict (features [3 ]) > 0 )
292293
293- categoricalFeaturesInfo = {0 : 3 } # feature 0 has 3 categories
294+ categoricalFeaturesInfo = {0 : 3 } # feature 0 has 3 categories
294295 dt_model = DecisionTree .trainClassifier (rdd , numClasses = 2 ,
295296 categoricalFeaturesInfo = categoricalFeaturesInfo )
296297 self .assertTrue (dt_model .predict (features [0 ]) <= 0 )
@@ -329,8 +330,9 @@ def test_regression(self):
329330 self .assertTrue (rr_model .predict (features [2 ]) <= 0 )
330331 self .assertTrue (rr_model .predict (features [3 ]) > 0 )
331332
332- categoricalFeaturesInfo = {0 : 2 } # feature 0 has 2 categories
333- dt_model = DecisionTree .trainRegressor (rdd , categoricalFeaturesInfo = categoricalFeaturesInfo )
333+ categoricalFeaturesInfo = {0 : 2 } # feature 0 has 2 categories
334+ dt_model = DecisionTree .trainRegressor (
335+ rdd , categoricalFeaturesInfo = categoricalFeaturesInfo )
334336 self .assertTrue (dt_model .predict (features [0 ]) <= 0 )
335337 self .assertTrue (dt_model .predict (features [1 ]) > 0 )
336338 self .assertTrue (dt_model .predict (features [2 ]) <= 0 )
0 commit comments