From 8c3ea93fc8a9be5dcac2bd0ad0d8e34cc13f6dd3 Mon Sep 17 00:00:00 2001
From: Fangjun Kuang <csukuangfj@gmail.com>
Date: Mon, 27 Mar 2023 11:39:29 +0800
Subject: [PATCH] Save meta data to exported ONNX models (#968)

---
 .../ASR/pruned_transducer_stateless3/export-onnx.py    | 10 ++++++++++
 .../ASR/pruned_transducer_stateless5/export-onnx.py    | 10 ++++++++++
 .../ASR/pruned_transducer_stateless7/export-onnx.py    | 10 ++++++++++
 3 files changed, 30 insertions(+)

diff --git a/egs/librispeech/ASR/pruned_transducer_stateless3/export-onnx.py b/egs/librispeech/ASR/pruned_transducer_stateless3/export-onnx.py
index ca8be307c..36e57e946 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless3/export-onnx.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless3/export-onnx.py
@@ -273,6 +273,16 @@ def export_encoder_model_onnx(
         },
     )
 
+    meta_data = {
+        "model_type": "conformer",
+        "version": "1",
+        "model_author": "k2-fsa",
+        "comment": "stateless3",
+    }
+    logging.info(f"meta_data: {meta_data}")
+
+    add_meta_data(filename=encoder_filename, meta_data=meta_data)
+
 
 def export_decoder_model_onnx(
     decoder_model: OnnxDecoder,
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless5/export-onnx.py b/egs/librispeech/ASR/pruned_transducer_stateless5/export-onnx.py
index 743fe8a92..3d94760dc 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless5/export-onnx.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless5/export-onnx.py
@@ -296,6 +296,16 @@ def export_encoder_model_onnx(
         },
     )
 
+    meta_data = {
+        "model_type": "conformer",
+        "version": "1",
+        "model_author": "k2-fsa",
+        "comment": "stateless5",
+    }
+    logging.info(f"meta_data: {meta_data}")
+
+    add_meta_data(filename=encoder_filename, meta_data=meta_data)
+
 
 def export_decoder_model_onnx(
     decoder_model: OnnxDecoder,
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/export-onnx.py b/egs/librispeech/ASR/pruned_transducer_stateless7/export-onnx.py
index f76915a74..93eb8df3d 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/export-onnx.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/export-onnx.py
@@ -291,6 +291,16 @@ def export_encoder_model_onnx(
         },
     )
 
+    meta_data = {
+        "model_type": "zipformer",
+        "version": "1",
+        "model_author": "k2-fsa",
+        "comment": "stateless7",
+    }
+    logging.info(f"meta_data: {meta_data}")
+
+    add_meta_data(filename=encoder_filename, meta_data=meta_data)
+
 
 def export_decoder_model_onnx(
     decoder_model: OnnxDecoder,