From 251aec806fc319ada281e00aa1bf2757601c5efd Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Wed, 5 Jul 2023 17:06:33 -0700 Subject: [PATCH 1/2] Changed from **kwargs Signed-off-by: Boris Fomitchev --- nemo/collections/asr/models/asr_model.py | 8 ++++---- nemo/core/classes/exportable.py | 2 +- scripts/export.py | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/nemo/collections/asr/models/asr_model.py b/nemo/collections/asr/models/asr_model.py index 1e38b8b93062..20be6cc16203 100644 --- a/nemo/collections/asr/models/asr_model.py +++ b/nemo/collections/asr/models/asr_model.py @@ -240,10 +240,10 @@ def disabled_deployment_input_names(self): def disabled_deployment_output_names(self): return self.encoder.disabled_deployment_output_names - def set_export_config(self, **kwargs): - if 'cache_support' in kwargs: - enable = bool(kwargs['cache_support']) + def set_export_config(self, args): + if 'cache_support' in args: + enable = bool(args['cache_support']) self.encoder.export_cache_support = enable logging.info(f"Caching support enabled: {enable}") self.encoder.setup_streaming_params() - super().set_export_config(**kwargs) + super().set_export_config(args) diff --git a/nemo/core/classes/exportable.py b/nemo/core/classes/exportable.py index 3f54fb9a0f02..0a77ccfee41a 100644 --- a/nemo/core/classes/exportable.py +++ b/nemo/core/classes/exportable.py @@ -306,7 +306,7 @@ def list_export_subnets(self): def get_export_config(self): return getattr(self, 'export_config', {}) - def set_export_config(self, **kwargs): + def set_export_config(self, kwargs): ex_config = self.get_export_config() ex_config.update(kwargs) self.export_config = ex_config diff --git a/scripts/export.py b/scripts/export.py index 7c68fa4cd41c..e22023786448 100644 --- a/scripts/export.py +++ b/scripts/export.py @@ -140,11 +140,11 @@ def nemo_export(argv): max_dim = args.max_dim if args.cache_support: - model.set_export_config(cache_support=True) + model.set_export_config({cache_support: True}) if args.config: kv = dict(map(lambda s: s.split('='), args.config)) - model.set_export_config(**kv) + model.set_export_config(kv) autocast = nullcontext if args.autocast: From 6392c43cee2a46739227b8753437f2614723958a Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Wed, 5 Jul 2023 17:30:46 -0700 Subject: [PATCH 2/2] Docstring Signed-off-by: Boris Fomitchev --- nemo/core/classes/exportable.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/nemo/core/classes/exportable.py b/nemo/core/classes/exportable.py index 0a77ccfee41a..e6f131fa0617 100644 --- a/nemo/core/classes/exportable.py +++ b/nemo/core/classes/exportable.py @@ -304,9 +304,15 @@ def list_export_subnets(self): return ['self'] def get_export_config(self): + """ + Returns export_config dictionary + """ return getattr(self, 'export_config', {}) def set_export_config(self, kwargs): + """ + Sets/updates export_config dictionary + """ ex_config = self.get_export_config() ex_config.update(kwargs) self.export_config = ex_config