Repository: mlflow/mlflow

Issue:
Fix shap plot feature name de-duplication
Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

<!-- 🚨 We recommend pull requests be filed from a non-master branch on a repository fork (e.g. <username>:fix-xxx). 🚨 -->

## Related Issues/PRs

<!--
Please reference any related feature requests, issues, or PRs here. For example, `#123`. To automatically close the referenced items when this PR is merged, please use a closing keyword (close, fix, or resolve). For example, `Close #123`. See https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue for more information.
-->
#xxx

## What changes are proposed in this pull request?

Fix shap plot feature name de-duplication.

Before:
It always append suffix `(f_{feature_index})` on truncated names.

After:
Only append suffix `({dup_index}` on duplicated names. 

## How is this patch tested?
Unit tests.

## Does this PR change the documentation?

- [x] No. You can skip the rest of this section.
- [ ] Yes. Make sure the changed pages / sections render correctly by following the steps below.

1. Click the `Details` link on the `Preview docs` check.
2. Find the changed pages / sections and make sure they render correctly.

## Release Notes

### Is this a user-facing change?

- [ ] No. You can skip the rest of this section.
- [x] Yes. Give a description of this change to be included in the release notes for MLflow users.

Before:
It always append suffix (f_{feature_index}) on truncated names.

After:
Only append suffix ({dup_index} on duplicated names.

### What component(s), interfaces, languages, and integrations does this PR affect?
Components
- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [x] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface
- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Language
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

<!--
Insert an empty named anchor here to allow jumping to this section with a fragment URL
(e.g. https://github.com/mlflow/mlflow/pull/123#user-content-release-note-category).
Note that GitHub prefixes anchor names in markdown with "user-content-".
-->
<a name="release-note-category"></a>
### How should the PR be classified in the release notes? Choose one:

- [ ] `rn/breaking-change` - The PR will be mentioned in the "Breaking Changes" section
- [x] `rn/none` - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
- [ ] `rn/feature` - A new user-facing feature worth mentioning in the release notes
- [ ] `rn/bug-fix` - A user-facing bug fix worth mentioning in the release notes
- [ ] `rn/documentation` - A user-facing documentation change worth mentioning in the release notes



Discussion:
Discussion:
I am rethinking whether the de-dup logic in this PR is good. e.g., let's see the following example (applying previous de-dup logic) :
![shap_beeswarm_plot_on_data_diabetes_dataset](https://user-images.githubusercontent.com/19235986/188540555-2ec21869-0c2b-4ba3-920a-7cde6b5fe08d.png)

Note the feature display order is changed on the plot. If using the de-dup logic in this PR, seemingly we cannot distinguish the original feature well. But if we append suffix with the exact "feature_index", it is easier for user to distinguish them.
> But if we append suffix with the exact "feature_index", it is easier for user to distinguish them.

If users understand what `f_{feature_index}` means.


btw do we really need to truncate the feature names?
> btw do we really need to truncate the feature names?

If we don't truncate feature names, too long feature name might cause the plot layout broken.
> If we don't truncate feature names, too long feature name might cause the plot layout broken.

Can we create a function that creates the plot so that we can control the layout?
> > If we don't truncate feature names, too long feature name might cause the plot layout broken.
> 
> Can we create a function that creates the plot so that we can control the layout?

Then we might need to copy and modify shap plot generation code, which makes maintenance harder.
And for long feature names, without truncation we cannot show it in plot well, we have to either truncate it or make font size smaller.


Resolving patch:
diff --git a/mlflow/models/evaluation/default_evaluator.py b/mlflow/models/evaluation/default_evaluator.py
--- a/mlflow/models/evaluation/default_evaluator.py
+++ b/mlflow/models/evaluation/default_evaluator.py
@@ -9,7 +9,6 @@
 from mlflow.entities.metric import Metric
 from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
 from mlflow.utils.file_utils import TempDir
-from mlflow.utils.string_utils import truncate_str_from_middle
 from mlflow.models.utils import plot_lines
 from mlflow.models.evaluation.artifacts import (
     ImageEvaluationArtifact,
@@ -474,7 +473,7 @@ def _log_image_artifact(
         try:
             pyplot.clf()
             do_plot()
-            pyplot.savefig(artifact_file_local_path)
+            pyplot.savefig(artifact_file_local_path, bbox_inches="tight")
         finally:
             pyplot.close(pyplot.gcf())
 
@@ -560,19 +559,7 @@ def _log_model_explainability(self):
             "explainability_nsamples", _DEFAULT_SAMPLE_ROWS_FOR_SHAP
         )
 
-        truncated_feature_names = [truncate_str_from_middle(f, 20) for f in self.feature_names]
-        for i, truncated_name in enumerate(truncated_feature_names):
-            if truncated_name != self.feature_names[i]:
-                # For duplicated truncated name, attach "(f_{feature_index})" at the end
-                truncated_feature_names[i] = f"{truncated_name}(f_{i + 1})"
-
-        truncated_feature_name_map = dict(zip(self.feature_names, truncated_feature_names))
-
-        # For some shap explainer, the plot will use the DataFrame column names instead of
-        # using feature_names argument value. So rename the dataframe column names.
-        X_df = self.X.copy_to_avoid_mutation().rename(
-            columns=truncated_feature_name_map, copy=False
-        )
+        X_df = self.X.copy_to_avoid_mutation()
 
         sampled_X = shap.sample(X_df, sample_rows, random_state=0)
 
@@ -614,7 +601,7 @@ def _log_model_explainability(self):
                     explainer = shap.Explainer(
                         shap_predict_fn,
                         sampled_X,
-                        feature_names=truncated_feature_names,
+                        feature_names=self.feature_names,
                         algorithm=algorithm,
                     )
             else:
@@ -627,19 +614,19 @@ def _log_model_explainability(self):
                     # for raw model, this case shap plot doesn't support it well, so exclude the
                     # multinomial_classifier case here.
                     explainer = shap.Explainer(
-                        self.raw_model, sampled_X, feature_names=truncated_feature_names
+                        self.raw_model, sampled_X, feature_names=self.feature_names
                     )
                 else:
                     # fallback to default explainer
                     explainer = shap.Explainer(
-                        shap_predict_fn, sampled_X, feature_names=truncated_feature_names
+                        shap_predict_fn, sampled_X, feature_names=self.feature_names
                     )
 
             _logger.info(f"Shap explainer {explainer.__class__.__name__} is used.")
 
             if algorithm == "kernel":
                 shap_values = shap.Explanation(
-                    explainer.shap_values(sampled_X), feature_names=truncated_feature_names
+                    explainer.shap_values(sampled_X), feature_names=self.feature_names
                 )
             else:
                 shap_values = explainer(sampled_X)
@@ -669,9 +656,18 @@ def _log_model_explainability(self):
             )
             _logger.debug("", exc_info=True)
 
+        def _adjust_color_bar():
+            pyplot.gcf().axes[-1].set_aspect("auto")
+            pyplot.gcf().axes[-1].set_box_aspect(50)
+
+        def _adjust_axis_tick():
+            pyplot.xticks(fontsize=10)
+            pyplot.yticks(fontsize=10)
+
         def plot_beeswarm():
-            pyplot.subplots_adjust(bottom=0.2, left=0.4)
-            shap.plots.beeswarm(shap_values, show=False)
+            shap.plots.beeswarm(shap_values, show=False, color_bar=True)
+            _adjust_color_bar()
+            _adjust_axis_tick()
 
         self._log_image_artifact(
             plot_beeswarm,
@@ -679,8 +675,9 @@ def plot_beeswarm():
         )
 
         def plot_summary():
-            pyplot.subplots_adjust(bottom=0.2, left=0.4)
-            shap.summary_plot(shap_values, show=False)
+            shap.summary_plot(shap_values, show=False, color_bar=True)
+            _adjust_color_bar()
+            _adjust_axis_tick()
 
         self._log_image_artifact(
             plot_summary,
@@ -688,8 +685,8 @@ def plot_summary():
         )
 
         def plot_feature_importance():
-            pyplot.subplots_adjust(bottom=0.2, left=0.4)
             shap.plots.bar(shap_values, show=False)
+            _adjust_axis_tick()
 
         self._log_image_artifact(
             plot_feature_importance,


Test patch:
diff --git a/tests/models/test_default_evaluator.py b/tests/models/test_default_evaluator.py
--- a/tests/models/test_default_evaluator.py
+++ b/tests/models/test_default_evaluator.py
@@ -53,7 +53,6 @@
     iris_dataset,
     iris_pandas_df_dataset,
     iris_pandas_df_num_cols_dataset,
-    iris_pandas_df_longname_cols_dataset,
     binary_logistic_regressor_model_uri,
     breast_cancer_dataset,
     spark_linear_regressor_model_uri,
@@ -676,14 +675,6 @@ def test_default_explainer_pandas_df_num_cols(
     )
 
 
-def test_default_explainer_pandas_df_longname_cols(
-    multiclass_logistic_regressor_model_uri, iris_pandas_df_longname_cols_dataset
-):
-    _evaluate_explainer_with_exceptions(
-        multiclass_logistic_regressor_model_uri, iris_pandas_df_longname_cols_dataset
-    )
-
-
 def test_svm_classifier_evaluation_disable_logging_metrics_and_artifacts(
     svm_model_uri, breast_cancer_dataset
 ):
@@ -1516,29 +1507,6 @@ def test_autologging_is_disabled_during_evaluate(model):
         mlflow.sklearn.autolog(disable=True)
 
 
-def test_truncation_works_for_long_feature_names(linear_regressor_model_uri, diabetes_dataset):
-    evaluate(
-        linear_regressor_model_uri,
-        diabetes_dataset._constructor_args["data"],
-        model_type="regressor",
-        targets=diabetes_dataset._constructor_args["targets"],
-        dataset_name=diabetes_dataset.name,
-        feature_names=[
-            "f1",
-            "f2",
-            "f3longnamelongnamelongname",
-            "f4",
-            "f5",
-            "f6",
-            "f7longlonglonglong",
-            "f8",
-            "f9",
-            "f10",
-        ],
-        evaluators="default",
-    )
-
-
 def test_evaluation_works_with_model_pipelines_that_modify_input_data():
     iris = load_iris()
     X = pd.DataFrame(iris.data, columns=["0", "1", "2", "3"])
diff --git a/tests/models/test_evaluation.py b/tests/models/test_evaluation.py
--- a/tests/models/test_evaluation.py
+++ b/tests/models/test_evaluation.py
@@ -342,29 +342,6 @@ def iris_pandas_df_num_cols_dataset():
     return ds
 
 
-@pytest.fixture
-def iris_pandas_df_longname_cols_dataset():
-    X, y = get_iris()
-    eval_X, eval_y = X[0::3], y[0::3]
-    data = pd.DataFrame(
-        {
-            "f1": eval_X[:, 0],
-            "f2": eval_X[:, 1],
-            "f3longnamelongnamelongname": eval_X[:, 2],
-            "f4longnamelongnamelongname": eval_X[:, 3],
-            "y": eval_y,
-        }
-    )
-    constructor_args = {
-        "data": data,
-        "targets": "y",
-        "name": "iris_pandas_df_longname_cols_dataset",
-    }
-    ds = EvaluationDataset(**constructor_args)
-    ds._constructor_args = constructor_args
-    return ds
-
-
 @pytest.fixture
 def baseline_model_uri(request):
     if request.param == "linear_regressor_model_uri":