Fix binary variables being visualized as 0 and 1 (#3691)

ludwig-ai · Oct 6, 2023 · e4fa42c · e4fa42c
1 parent d74b311
commit e4fa42c
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/ludwig/visualize.py b/ludwig/visualize.py
@@ -3732,13 +3732,15 @@ def confusion_matrix(
                 model_name_name = (
                     model_names_list[i] if (model_names_list is not None and i < len(model_names_list)) else ""
                 )
-
                 if (
                     metadata is not None
                     and output_feature_name in metadata
-                    and "idx2str" in metadata[output_feature_name]
+                    and ("idx2str" in metadata[output_feature_name] or "bool2str" in metadata[output_feature_name])
                 ):
-                    labels = metadata[output_feature_name]["idx2str"]
+                    if "bool2str" in metadata[output_feature_name]:  # Handles the binary output case
+                        labels = metadata[output_feature_name]["bool2str"]
+                    else:
+                        labels = metadata[output_feature_name]["idx2str"]
                 else:
                     labels = list(range(len(_confusion_matrix)))