only get feature_importance if select feature by top n. (#3235)

* only get feature_importance if select feature by top n. * remove extra import --------- Co-authored-by: Richard Li <[email protected]>
Azure · Aug 19, 2024 · a0567e3 · a0567e3
1 parent 7d13863
commit a0567e3
Showing 1 changed file with 7 additions and 11 deletions.
diff --git a/assets/model_monitoring/components/src/model_monitor_feature_selector/run.py b/assets/model_monitoring/components/src/model_monitor_feature_selector/run.py
@@ -13,7 +13,7 @@
 from shared_utilities.io_utils import (
     try_read_mltable_in_spark, try_read_mltable_in_spark_with_error, save_spark_df_as_mltable
 )
-from shared_utilities.momo_exceptions import DataNotFoundError, InvalidInputError
+from shared_utilities.momo_exceptions import InvalidInputError
 
 
 def run():
@@ -42,16 +42,12 @@ def run():
         input_df1 = input_df2
 
     feature_importance = None
-    try:
-        feature_importance = try_read_mltable_in_spark_with_error(args.feature_importance, "feature_importance")
-    except DataNotFoundError as e:
-        if args.filter_type == FeatureSelectorType.TOP_N_BY_ATTRIBUTION.name:
-            raise e
-    except Exception:
-        if args.filter_type == FeatureSelectorType.TOP_N_BY_ATTRIBUTION.name:
-            raise Exception(
-                "Error encountered when retrieving top N features. Please ensure target_column is defined."
-            )
+    if args.filter_type == FeatureSelectorType.TOP_N_BY_ATTRIBUTION.name:
+        try:
+            feature_importance = try_read_mltable_in_spark_with_error(args.feature_importance, "feature_importance")
+        except Exception as e:
+            raise Exception(f"Error encountered when retrieving top N features. {e}. "
+                            "The feature importance step may not have completed successfully.")
 
     feature_selector = FeatureSelectorFactory().produce(
         feature_selector_type=args.filter_type,