Resolve "integration-sota"

code/bolsonaro/trainer.py

@@ -2,6 +2,7 @@ from bolsonaro.models.model_raw_results import ModelRawResults
 from bolsonaro.models.omp_forest_regressor import OmpForestRegressor
 from bolsonaro.models.omp_forest_classifier import OmpForestBinaryClassifier, OmpForestMulticlassClassifier
 from bolsonaro.models.similarity_forest_regressor import SimilarityForestRegressor
+from bolsonaro.models.kmeans_forest_regressor import KMeansForestRegressor
 from bolsonaro.error_handling.logger_factory import LoggerFactory
 from bolsonaro.data.task import Task
 from . import LOG_PATH
@@ -96,7 +97,7 @@ class Trainer(object):
         self._end_time = time.time()
 
     def __score_func(self, model, X, y_true, weights=True):
-        if type(model) in [OmpForestRegressor, RandomForestRegressor, SimilarityForestRegressor]:
+        if type(model) in [OmpForestRegressor, RandomForestRegressor]:
             if weights:
                 y_pred = model.predict(X)
             else:
@@ -109,12 +110,14 @@ class Trainer(object):
                 y_pred = model.predict_no_weights(X)
             if type(model) is OmpForestBinaryClassifier:
                 y_pred = np.sign(y_pred)
-                y_pred = np.where(y_pred==0, 1, y_pred)
+                y_pred = np.where(y_pred == 0, 1, y_pred)
             result = self._classification_score_metric(y_true, y_pred)
+        elif type(model) in [SimilarityForestRegressor, KMeansForestRegressor]:
+            result = model.score(X, y_true)
         return result
 
     def __score_func_base(self, model, X, y_true):
-        if type(model) == OmpForestRegressor:
+        if type(model) in [OmpForestRegressor, SimilarityForestRegressor, KMeansForestRegressor]:
             y_pred = model.predict_base_estimator(X)
             result = self._base_regression_score_metric(y_true, y_pred)
         elif type(model) in [OmpForestBinaryClassifier, OmpForestMulticlassClassifier]:
@@ -123,7 +126,7 @@ class Trainer(object):
         elif type(model) == RandomForestClassifier:
             y_pred = model.predict(X)
             result = self._base_classification_score_metric(y_true, y_pred)
-        elif type(model) in [RandomForestRegressor, SimilarityForestRegressor]:
+        elif type(model) is RandomForestRegressor:
             y_pred = model.predict(X)
             result = self._base_regression_score_metric(y_true, y_pred)
         return result