Provide separate flags to reset offsets and words

Separate reset flags are required in cases where we want to get continuous time stamps across reset but don't want past words. Signed-off-by: Viraj Karandikar <[email protected]>
NVIDIA · Nov 19, 2019 · 5ccc608 · 5ccc608
1 parent a81babd
commit 5ccc608
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/decoders/ctc_beam_search_decoder.cpp b/decoders/ctc_beam_search_decoder.cpp
@@ -259,7 +259,7 @@ BeamDecoder::~BeamDecoder()
 }
 
 
-void BeamDecoder::reset(bool partial /* default = false */)
+void BeamDecoder::reset(bool keep_offset /*default = false*/, bool keep_words /*default = false*/)
 {
   // init prefixes' root
   if (root != nullptr) {
@@ -279,15 +279,20 @@ void BeamDecoder::reset(bool partial /* default = false */)
     root->set_matcher(matcher);
   }
 
-  if (partial) {
+  if (keep_offset) {
+    prev_time_offset += last_decoded_timestep + time_offset;
+  } else {
+    prev_time_offset = 0;
+  }
+
+  if (keep_words) {
     prev_wordlist.insert(
         std::end(prev_wordlist), std::begin(wordlist),
         std::end(wordlist));
-    prev_time_offset += last_decoded_timestep + time_offset;
   } else {
     prev_wordlist.clear();
-    prev_time_offset = 0;
   }
+
   wordlist.clear();
   time_offset = 0;
   last_decoded_timestep = 0;

diff --git a/decoders/ctc_beam_search_decoder.h b/decoders/ctc_beam_search_decoder.h
@@ -51,7 +51,7 @@ class BeamDecoder {
   void set_start_offset(int offset) { time_offset = offset; }
 
   // reset state
-  void reset(bool partial = false);
+  void reset(bool keep_offset = false, bool keep_words = false);
 
 private:
   Scorer *ext_scorer;