add feature importance to fit step

tryolabs · Mar 5, 2024 · c5210f1 · c5210f1
1 parent aa670ec
commit c5210f1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/pipeline_lib/core/data_container.py b/pipeline_lib/core/data_container.py
@@ -34,6 +34,7 @@ class DataContainer:
     EXPLAINER = "explainer"
     TUNING_PARAMS = "tuning_params"
     TARGET = "target"
+    IMPORTANCE = "importance"
 
     def __init__(self, initial_data: Optional[dict] = None):
         """

diff --git a/pipeline_lib/implementation/tabular/xgboost/fit_model.py b/pipeline_lib/implementation/tabular/xgboost/fit_model.py
@@ -75,6 +75,10 @@ def execute(self, data: DataContainer) -> DataContainer:
         # Save the model to the data container
         data[DataContainer.MODEL] = model
 
+        importance = model.get_booster().get_score(importance_type="gain")
+        importance = dict(sorted(importance.items(), key=lambda item: item[1], reverse=True))
+        data[DataContainer.IMPORTANCE] = importance
+
         end_time = time.time()
         elapsed_time = end_time - start_time
         minutes = int(elapsed_time // 60)