66#ifndef SHERPA_ONNX_CSRC_ONLINE_RECOGNIZER_TRANSDUCER_NEMO_IMPL_H_
77#define SHERPA_ONNX_CSRC_ONLINE_RECOGNIZER_TRANSDUCER_NEMO_IMPL_H_
88
9+ #include < algorithm>
910#include < fstream>
1011#include < ios>
1112#include < memory>
3233namespace sherpa_onnx {
3334
3435// defined in ./online-recognizer-transducer-impl.h
35- // static may or may not be here? TODDOs
36- static OnlineRecognizerResult Convert (const OnlineTransducerDecoderResult &src,
37- const SymbolTable &sym_table,
38- float frame_shift_ms,
39- int32_t subsampling_factor,
40- int32_t segment,
41- int32_t frames_since_start);
36+ OnlineRecognizerResult Convert (const OnlineTransducerDecoderResult &src,
37+ const SymbolTable &sym_table,
38+ float frame_shift_ms, int32_t subsampling_factor,
39+ int32_t segment, int32_t frames_since_start);
4240
4341class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl {
44- public:
42+ public:
4543 explicit OnlineRecognizerTransducerNeMoImpl (
4644 const OnlineRecognizerConfig &config)
4745 : config_(config),
4846 symbol_table_(config.model_config.tokens),
4947 endpoint_(config_.endpoint_config),
50- model_(std::make_unique<OnlineTransducerNeMoModel>(
51- config.model_config)) {
48+ model_(
49+ std::make_unique<OnlineTransducerNeMoModel>( config.model_config)) {
5250 if (config.decoding_method == " greedy_search"
5351 decoder_ = std::make_unique<OnlineTransducerGreedySearchNeMoDecoder>(
5452 model_.get (), config_.blank_penalty );
@@ -73,7 +71,7 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl {
7371 model_.get (), config_.blank_penalty );
7472 } else {
7573 SHERPA_ONNX_LOGE (" Unsupported decoding method: %s"
76- config.decoding_method .c_str ());
74+ config.decoding_method .c_str ());
7775 exit (-1 );
7876 }
7977
@@ -83,7 +81,6 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl {
8381
8482 std::unique_ptr<OnlineStream> CreateStream () const override {
8583 auto stream = std::make_unique<OnlineStream>(config_.feat_config );
86- stream->SetStates (model_->GetInitStates ());
8784 InitOnlineStream (stream.get ());
8885 return stream;
8986 }
@@ -94,14 +91,12 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl {
9491 }
9592
9693 OnlineRecognizerResult GetResult (OnlineStream *s) const override {
97- OnlineTransducerDecoderResult decoder_result = s->GetResult ();
98- decoder_->StripLeadingBlanks (&decoder_result);
99-
10094 // TODO(fangjun): Remember to change these constants if needed
10195 int32_t frame_shift_ms = 10 ;
102- int32_t subsampling_factor = 8 ;
103- return Convert (decoder_result, symbol_table_, frame_shift_ms, subsampling_factor,
104- s->GetCurrentSegment (), s->GetNumFramesSinceStart ());
96+ int32_t subsampling_factor = model_->SubsamplingFactor ();
97+ return Convert (s->GetResult (), symbol_table_, frame_shift_ms,
98+ subsampling_factor, s->GetCurrentSegment (),
99+ s->GetNumFramesSinceStart ());
105100 }
106101
107102 bool IsEndpoint (OnlineStream *s) const override {
@@ -114,8 +109,8 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl {
114109 // frame shift is 10 milliseconds
115110 float frame_shift_in_seconds = 0.01 ;
116111
117- // subsampling factor is 8
118- int32_t trailing_silence_frames = s->GetResult ().num_trailing_blanks * 8 ;
112+ int32_t trailing_silence_frames =
113+ s->GetResult ().num_trailing_blanks * model_-> SubsamplingFactor () ;
119114
120115 return endpoint_.IsEndpoint (num_processed_frames, trailing_silence_frames,
121116 frame_shift_in_seconds);
@@ -126,19 +121,16 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl {
126121 // segment is incremented only when the last
127122 // result is not empty
128123 const auto &r = s->GetResult ();
129- if (!r.tokens .empty () && r. tokens . back () != 0 ) {
124+ if (!r.tokens .empty ()) {
130125 s->GetCurrentSegment () += 1 ;
131126 }
132127 }
133128
134- // we keep the decoder_out
135- decoder_-> UpdateDecoderOut (&s-> GetResult ());
136- Ort::Value decoder_out = std::move ( s->GetResult (). decoder_out );
129+ s-> SetResult ({});
130+
131+ s->SetStates (model_-> GetEncoderInitStates () );
137132
138- auto r = decoder_->GetEmptyResult ();
139-
140- s->SetResult (r);
141- s->GetResult ().decoder_out = std::move (decoder_out);
133+ s->SetNeMoDecoderStates (model_->GetDecoderInitStates ());
142134
143135 // Note: We only update counters. The underlying audio samples
144136 // are not discarded.
@@ -151,10 +143,9 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl {
151143
152144 int32_t feature_dim = ss[0 ]->FeatureDim ();
153145
154- std::vector<OnlineTransducerDecoderResult> result (n);
155146 std::vector<float > features_vec (n * chunk_size * feature_dim);
156147 std::vector<std::vector<Ort::Value>> encoder_states (n);
157-
148+
158149 for (int32_t i = 0 ; i != n; ++i) {
159150 const auto num_processed_frames = ss[i]->GetNumProcessedFrames ();
160151 std::vector<float > features =
@@ -166,9 +157,7 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl {
166157 std::copy (features.begin (), features.end (),
167158 features_vec.data () + i * chunk_size * feature_dim);
168159
169- result[i] = std::move (ss[i]->GetResult ());
170160 encoder_states[i] = std::move (ss[i]->GetStates ());
171-
172161 }
173162
174163 auto memory_info =
@@ -180,42 +169,35 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl {
180169 features_vec.size (), x_shape.data (),
181170 x_shape.size ());
182171
183- // Batch size is 1
184- auto states = std::move (encoder_states[0 ]);
185- int32_t num_states = states.size (); // num_states = 3
172+ auto states = model_->StackStates (std::move (encoder_states));
173+ int32_t num_states = states.size (); // num_states = 3
186174 auto t = model_->RunEncoder (std::move (x), std::move (states));
187175 // t[0] encoder_out, float tensor, (batch_size, dim, T)
188176 // t[1] next states
189-
177+
190178 std::vector<Ort::Value> out_states;
191179 out_states.reserve (num_states);
192-
180+
193181 for (int32_t k = 1 ; k != num_states + 1 ; ++k) {
194182 out_states.push_back (std::move (t[k]));
195183 }
196184
185+ auto unstacked_states = model_->UnStackStates (std::move (out_states));
186+ for (int32_t i = 0 ; i != n; ++i) {
187+ ss[i]->SetStates (std::move (unstacked_states[i]));
188+ }
189+
197190 Ort::Value encoder_out = Transpose12 (model_->Allocator (), &t[0 ]);
198-
199- // defined in online-transducer-greedy-search-nemo-decoder.h
200- // get intial states of decoder.
201- std::vector<Ort::Value> &decoder_states = ss[0 ]->GetNeMoDecoderStates ();
202-
203- // Subsequent decoder states (for each chunks) are updated inside the Decode method.
204- // This returns the decoder state from the LAST chunk. We probably dont need it. So we can discard it.
205- decoder_states = decoder_->Decode (std::move (encoder_out),
206- std::move (decoder_states),
207- &result, ss, n);
208-
209- ss[0 ]->SetResult (result[0 ]);
210-
211- ss[0 ]->SetStates (std::move (out_states));
191+
192+ decoder_->Decode (std::move (encoder_out), ss, n);
212193 }
213194
214195 void InitOnlineStream (OnlineStream *stream) const {
215- auto r = decoder_->GetEmptyResult ();
196+ // set encoder states
197+ stream->SetStates (model_->GetEncoderInitStates ());
216198
217- stream-> SetResult (r);
218- stream->SetNeMoDecoderStates (model_->GetDecoderInitStates (1 ));
199+ // set decoder states
200+ stream->SetNeMoDecoderStates (model_->GetDecoderInitStates ());
219201 }
220202
221203 private:
@@ -250,7 +232,6 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl {
250232 symbol_table_.NumSymbols (), vocab_size);
251233 exit (-1 );
252234 }
253-
254235 }
255236
256237 private:
@@ -259,9 +240,8 @@ class OnlineRecognizerTransducerNeMoImpl : public OnlineRecognizerImpl {
259240 std::unique_ptr<OnlineTransducerNeMoModel> model_;
260241 std::unique_ptr<OnlineTransducerGreedySearchNeMoDecoder> decoder_;
261242 Endpoint endpoint_;
262-
263243};
264244
265245} // namespace sherpa_onnx
266246
267- #endif // SHERPA_ONNX_CSRC_ONLINE_RECOGNIZER_TRANSDUCER_NEMO_IMPL_H_
247+ #endif // SHERPA_ONNX_CSRC_ONLINE_RECOGNIZER_TRANSDUCER_NEMO_IMPL_H_
0 commit comments