From 88c4f4a25df3c171dc0d07efc400f73e6a09e165 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期二, 07 二月 2023 21:43:30 +0800
Subject: [PATCH] export model

---
 funasr/export/models/e2e_asr_paraformer.py |   11 ++++++++++-
 1 files changed, 10 insertions(+), 1 deletions(-)

diff --git a/funasr/export/models/e2e_asr_paraformer.py b/funasr/export/models/e2e_asr_paraformer.py
index dd87213..8388f4f 100644
--- a/funasr/export/models/e2e_asr_paraformer.py
+++ b/funasr/export/models/e2e_asr_paraformer.py
@@ -63,8 +63,9 @@
 
         decoder_out, _ = self.decoder(enc, enc_len, pre_acoustic_embeds, pre_token_length)
         decoder_out = torch.log_softmax(decoder_out, dim=-1)
+        sample_ids = decoder_out.argmax(dim=-1)
 
-        return decoder_out, pre_token_length
+        return decoder_out, sample_ids
     
     # def get_output_size(self):
     #     return self.model.encoders[0].size
@@ -74,6 +75,14 @@
         speech_lengths = torch.tensor([6, 30], dtype=torch.int32)
         return (speech, speech_lengths)
 
+    def get_dummy_inputs_txt(self, txt_file: str = "/mnt/workspace/data_fbank/0207/12345.wav.fea.txt"):
+        import numpy as np
+        fbank = np.loadtxt(txt_file)
+        fbank_lengths = np.array([fbank.shape[0], ], dtype=np.int32)
+        speech = torch.from_numpy(fbank[None, :, :].astype(np.float32))
+        speech_lengths = torch.from_numpy(fbank_lengths.astype(np.int32))
+        return (speech, speech_lengths)
+
     def get_input_names(self):
         return ['speech', 'speech_lengths']
 

--
Gitblit v1.9.1