diff --git a/easy_rec/python/tools/add_feature_info_to_config.py b/easy_rec/python/tools/add_feature_info_to_config.py index e82a04516..1c2d9b606 100644 --- 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 index b63a02f71..8443226c7 100644 --- 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.