Update ModelFamily xgboost/catboost values (#677)

* Update ModelFamily enum vals. * Changelog * Fix test
alteryx · Apr 20, 2020 · 70d6821 · 70d6821
1 parent b8642ca
commit 70d6821
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
@@ -33,6 +33,7 @@ Changelog
         * Remove s3fs as a dev dependency :pr:`664`
         * Changed requirements-parser to be a core dependency :pr:`673`
         * Replace `supported_problem_types` field on pipelines with `problem_type` attribute on base classes :pr:`678`
+        * Update `ModelFamily` values: don't list xgboost/catboost as classifiers now that we have regression pipelines for them :pr:`677`
     * Documentation Changes
         * Fixed some sphinx warnings :pr:`593`
         * Fixed docstring for AutoClassificationSearch with correct command :pr:`599`

diff --git a/evalml/model_family/model_family.py b/evalml/model_family/model_family.py
@@ -11,8 +11,8 @@ class ModelFamily(Enum):
 
     def __str__(self):
         model_family_dict = {ModelFamily.RANDOM_FOREST.name: "Random Forest",
-                             ModelFamily.XGBOOST.name: "XGBoost Classifier",
-                             ModelFamily.LINEAR_MODEL.name: "Linear Model",
-                             ModelFamily.CATBOOST.name: "CatBoost Classifier",
+                             ModelFamily.XGBOOST.name: "XGBoost",
+                             ModelFamily.LINEAR_MODEL.name: "Linear",
+                             ModelFamily.CATBOOST.name: "CatBoost",
                              ModelFamily.NONE.name: "None"}
         return model_family_dict[self.name]
diff --git a/evalml/tests/pipeline_tests/test_pipelines.py b/evalml/tests/pipeline_tests/test_pipelines.py
@@ -196,7 +196,7 @@ def test_describe(X_y, capsys, lr_pipeline):
     out, err = capsys.readouterr()
     assert "Logistic Regression Binary Pipeline" in out
     assert "Problem Type: Binary Classification" in out
-    assert "Model Family: Linear Model" in out
+    assert "Model Family: Linear" in out
 
     for component in lrp.component_graph:
         if component.hyperparameter_ranges: