remove feature_configs for drop features

alibaba · Oct 20, 2023 · 0108c22 · 0108c22
1 parent e18caf0
commit 0108c22
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/easy_rec/python/tools/add_feature_info_to_config.py b/easy_rec/python/tools/add_feature_info_to_config.py
@@ -59,9 +59,9 @@ def main(argv):
       except common_io.exception.OutOfRangeException:
         reader.close()
         break
-
-  for feature_config in config_util.get_compatible_feature_configs(
-      pipeline_config):
+  
+  feature_configs = config_util.get_compatible_feature_configs(pipeline_config, drop_feature_names)
+  for feature_config in feature_configs:
     feature_name = feature_config.input_names[0]
     if feature_name in feature_info_map:
       logging.info('edited %s' % feature_name)

diff --git a/easy_rec/python/utils/config_util.py b/easy_rec/python/utils/config_util.py
@@ -390,14 +390,19 @@ def add_boundaries_to_config(pipeline_config, tables):
       logging.info('edited %s' % feature_name)
 
 
-def get_compatible_feature_configs(pipeline_config):
+def get_compatible_feature_configs(pipeline_config,drop_feature_names=None):
   if pipeline_config.feature_configs:
     feature_configs = pipeline_config.feature_configs
   else:
     feature_configs = pipeline_config.feature_config.features
+  if drop_feature_names:
+    tmp_feature_configs = feature_configs[:]
+    for fea_cfg in tmp_feature_configs:
+      fea_name = fea_cfg.input_names[0]
+      if fea_name in drop_feature_names:
+        feature_configs.remove(fea_cfg)
   return feature_configs
 
-
 def parse_time(time_data):
   """Parse time string to timestamp.