Cleanup and additional fixes

Signed-off-by: Anand Joseph <anajoseph@nvidia.com>

Author: anand-nv

nemo2riva/cookbook.py

@@ -48,8 +48,8 @@ def export_model(model, cfg, args, artifacts, metadata):
         format_meta = {"has_pytorch_checkpoint": True, "runtime": "PyTorch"}
     elif cfg.export_format == "NEMO":
         format_meta = {"has_pytorch_checkpoint": True, "runtime": "Python"}
-    elif cfg.export_format == "TRT-LLM":
-        format_meta = {"has_pytorch_checkpoint": False, "runtime": "TRT-LLM"}
+    elif cfg.export_format == "STATE":
+        format_meta = {"has_pytorch_checkpoint": False, "runtime": "Python"}
     # TODO: use submodel sections
     metadata.update(format_meta)
     runtime = format_meta["runtime"]
@@ -142,14 +142,13 @@ def export_model(model, cfg, args, artifacts, metadata):
 
         elif cfg.export_format == "NEMO":
             model.save_to(export_file)
-        elif cfg.export_format == "TRT-LLM":
+        elif cfg.export_format == "STATE":
             if not isinstance(model, Exportable):
                 logging.error("Your NeMo model class ({}) is not Exportable.".format(metadata['obj_cls']))
                 sys.exit(1)
-            model.transf_decoder.freeze()
-            decoder_params = model.transf_decoder.state_dict()
-            decoder_params.update(model.log_softmax.state_dict())
-            torch.save(decoder_params, export_file)
+            model.freeze()
+            model_params = model.state_dict()
+            torch.save(model_params, export_file)
 
 
 

nemo2riva/patches/aed_canary.py

@@ -72,7 +72,7 @@ def make_vocabulary_file(model, artifacts, **kwargs):
             lang = model.tokenizer.ids_to_lang([t_id])
             tokenizer_vocab['tokens'][lang][t_id] = model.tokenizer.ids_to_tokens([t_id])[0]
 
-        artifacts['vocab.json']={}
-        artifacts['vocab.json']['content'] = json.dumps(tokenizer_vocab).encode('utf-8')
+        #artifacts['vocab.json']={}
+        #artifacts['vocab.json']['content'] = json.dumps(tokenizer_vocab).encode('utf-8')
         return tokenizer_vocab
 

nemo2riva/schema.py

@@ -17,7 +17,7 @@
 
 schema_dict = None
 
-supported_formats = ["ONNX", "CKPT", "TS", "NEMO", "PYTORCH", "TRT-LLM"]
+supported_formats = ["ONNX", "CKPT", "TS", "NEMO", "PYTORCH", "STATE"]
 
 
 @dataclass
@@ -50,12 +50,13 @@ def get_export_config(export_obj, args):
         conf.export_file = list(export_obj)[0]
         attribs = export_obj[conf.export_file]
         conf.export_subnet = attribs.get('export_subnet', None)
+
         conf.is_onnx=attribs.get('onnx', False)
-        conf.trt_llm = False
-        conf.weights_only = False
+
+
 
         if not conf.is_onnx:
-            conf.trt_llm = attribs.get('trtllm', False)
+            conf.states_only = attribs.get('states_only', False)
             conf.is_torch = attribs.get('torch', False)
 
         if conf.export_file.endswith('.onnx'):
@@ -64,10 +65,11 @@ def get_export_config(export_obj, args):
             conf.export_format = "TS"
         elif conf.export_file.endswith('.nemo'):
             conf.export_format = "NEMO"
-        elif conf.trt_llm:
-            conf.export_format = "TRT-LLM"
-        elif conf.export_file.endswith('.pt'):
-            conf.export_format = "PYTORCH"
+        elif conf.is_torch:
+            if conf.states_only:
+                conf.export_format = "STATE"
+            else:
+                conf.export_format = "PYTORCH"
         else:
             conf.export_format = "CKPT"
         conf.autocast = attribs.get('autocast', False)

nemo2riva/validation_schemas/asr-scr-exported-aedmodel.yaml

@@ -14,7 +14,13 @@ artifact_properties:
       export_subnet: encoder
       onnx: True
   - decoder.pt:
-      trtllm: True
+      export_subnet: transf_decoder
+      states_only: True
+      torch: True
+      onnx: False
+  - log_softmax.pt:
+      export_subnet: log_softmax
+      states_only: True
       torch: True
       onnx: False
 

requirements.txt

@@ -2,10 +2,11 @@
 # SPDX-License-Identifier: MIT
 
 nemo_toolkit>=1.6.0
+torch>=2.4.0
 nvidia-eff>=0.6.4
 nvidia-eff-tao-encryption>=0.1.8
 nvidia-pyindex==1.0.6
-onnx==1.14.1
+onnx==1.16.1
 onnxruntime==1.16.3
 onnxruntime-gpu==1.16.3
 onnx-graphsurgeon==0.3.27