Fix plotting utils

Author: kalama-ai

octopus/diagnostics/_plots.py

@@ -7,14 +7,16 @@
 from plotly.subplots import make_subplots
 from sklearn.metrics import confusion_matrix
 
+from octopus.types import FIResultLabel
+
 
 def plot_feature_importance_chart(
     df: pd.DataFrame,
     *,
     outersplit_id: int | str = 0,
     task_id: int | str = 0,
     training_id: str = "",
-    fi_method: str = "",
+    fi_method: str | FIResultLabel = "",
 ) -> go.Figure:
     """Create a feature importance bar chart.
 

octopus/diagnostics/core.py

@@ -30,7 +30,7 @@
     plot_optuna_trials_chart,
     plot_predictions_vs_truth_chart,
 )
-from octopus.types import MLType
+from octopus.types import FIResultLabel, MLType
 
 
 def _has_ipywidgets() -> bool:
@@ -154,7 +154,7 @@ def plot_feature_importance(
         outersplit_id: int | None = None,
         task_id: int | None = None,
         training_id: str | None = None,
-        fi_method: str | None = None,
+        fi_method: str | FIResultLabel | None = None,
     ) -> None:
         """Plot feature importance bar chart.