Merge pull request #517 from jarivk/dev/separate_reset_flags

Provide separate flags to reset offsets and words
NVIDIA · Nov 19, 2019 · 61204b2 · 61204b2
2 parents a81babd + 5ccc608
commit 61204b2
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;