Merge pull request #112 from reuben/streaming_api

Expose streaming API

Author: Sean Naren

ctcdecode/src/ctc_beam_search_decoder.cpp

@@ -14,39 +14,33 @@
 
 using FSTMATCH = fst::SortedMatcher<fst::StdVectorFst>;
 
-std::vector<std::pair<double, Output>> ctc_beam_search_decoder(
-    const std::vector<std::vector<double>> &probs_seq,
-    const std::vector<std::string> &vocabulary,
-    size_t beam_size,
-    double cutoff_prob,
-    size_t cutoff_top_n,
-    size_t blank_id,
-    int log_input,
-    Scorer *ext_scorer) {
-  // dimension check
-  size_t num_time_steps = probs_seq.size();
-  for (size_t i = 0; i < num_time_steps; ++i) {
-    VALID_CHECK_EQ(probs_seq[i].size(),
-                   vocabulary.size(),
-                   "The shape of probs_seq does not match with "
-                   "the shape of the vocabulary");
-  }
-
-  // assign blank id
-  // size_t blank_id = vocabulary.size();
-
+DecoderState::DecoderState(const std::vector<std::string> &vocabulary,
+                           size_t beam_size,
+                           double cutoff_prob,
+                           size_t cutoff_top_n,
+                           size_t blank_id,
+                           int log_input,
+                           Scorer *ext_scorer)
+  : abs_time_step(0)
+  , beam_size(beam_size)
+  , cutoff_prob(cutoff_prob)
+  , cutoff_top_n(cutoff_top_n)
+  , blank_id(blank_id)
+  , log_input(log_input)
+  , vocabulary(vocabulary)
+  , ext_scorer(ext_scorer)
+{
   // assign space id
   auto it = std::find(vocabulary.begin(), vocabulary.end(), " ");
-  int space_id = it - vocabulary.begin();
   // if no space in vocabulary
-  if ((size_t)space_id >= vocabulary.size()) {
+  if (it == vocabulary.end()) {
     space_id = -2;
+  } else {
+    space_id = std::distance(vocabulary.begin(), it);
   }
 
   // init prefixes' root
-  PathTrie root;
   root.score = root.log_prob_b_prev = 0.0;
-  std::vector<PathTrie *> prefixes;
   prefixes.push_back(&root);
 
   if (ext_scorer != nullptr && !ext_scorer->is_character_based()) {
@@ -56,9 +50,22 @@ std::vector<std::pair<double, Output>> ctc_beam_search_decoder(
     auto matcher = std::make_shared<FSTMATCH>(*dict_ptr, fst::MATCH_INPUT);
     root.set_matcher(matcher);
   }
+}
+
+void
+DecoderState::next(const std::vector<std::vector<double>> &probs_seq)
+{
+  // dimension check
+  size_t num_time_steps = probs_seq.size();
+  for (size_t i = 0; i < num_time_steps; ++i) {
+    VALID_CHECK_EQ(probs_seq[i].size(),
+                   vocabulary.size(),
+                   "The shape of probs_seq does not match with "
+                   "the shape of the vocabulary");
+  }
 
   // prefix search over time
-  for (size_t time_step = 0; time_step < num_time_steps; ++time_step) {
+  for (size_t time_step = 0; time_step < num_time_steps; ++time_step, ++abs_time_step) {
     auto &prob = probs_seq[time_step];
 
     float min_cutoff = -NUM_FLT_INF;
@@ -97,7 +104,7 @@ std::vector<std::pair<double, Output>> ctc_beam_search_decoder(
               prefix->log_prob_nb_cur, log_prob_c + prefix->log_prob_nb_prev);
         }
         // get new prefix
-        auto prefix_new = prefix->get_path_trie(c, time_step, log_prob_c);
+        auto prefix_new = prefix->get_path_trie(c, abs_time_step, log_prob_c);
 
         if (prefix_new != nullptr) {
           float log_p = -NUM_FLT_INF;
@@ -147,45 +154,75 @@ std::vector<std::pair<double, Output>> ctc_beam_search_decoder(
       for (size_t i = beam_size; i < prefixes.size(); ++i) {
         prefixes[i]->remove();
       }
+
+      prefixes.resize(beam_size);
     }
   }  // end of loop over time
+}
+
+std::vector<std::pair<double, Output>>
+DecoderState::decode() const
+{
+  std::vector<PathTrie*> prefixes_copy = prefixes;
+  std::unordered_map<const PathTrie*, float> scores;
+  for (PathTrie* prefix : prefixes_copy) {
+    scores[prefix] = prefix->score;
+  }
 
   // score the last word of each prefix that doesn't end with space
   if (ext_scorer != nullptr && !ext_scorer->is_character_based()) {
-    for (size_t i = 0; i < beam_size && i < prefixes.size(); ++i) {
-      auto prefix = prefixes[i];
+    for (size_t i = 0; i < beam_size && i < prefixes_copy.size(); ++i) {
+      auto prefix = prefixes_copy[i];
       if (!prefix->is_empty() && prefix->character != space_id) {
         float score = 0.0;
         std::vector<std::string> ngram = ext_scorer->make_ngram(prefix);
         score = ext_scorer->get_log_cond_prob(ngram) * ext_scorer->alpha;
         score += ext_scorer->beta;
-        prefix->score += score;
+        scores[prefix] += score;
       }
     }
   }
 
-  size_t num_prefixes = std::min(prefixes.size(), beam_size);
-  std::sort(prefixes.begin(), prefixes.begin() + num_prefixes, prefix_compare);
+  using namespace std::placeholders;
+  size_t num_prefixes = std::min(prefixes_copy.size(), beam_size);
+  std::sort(prefixes_copy.begin(), prefixes_copy.begin() + num_prefixes,
+            std::bind(prefix_compare_external_scores, _1, _2, scores));
 
   // compute aproximate ctc score as the return score, without affecting the
   // return order of decoding result. To delete when decoder gets stable.
-  for (size_t i = 0; i < beam_size && i < prefixes.size(); ++i) {
-    double approx_ctc = prefixes[i]->score;
+  for (size_t i = 0; i < beam_size && i < prefixes_copy.size(); ++i) {
+    double approx_ctc = scores[prefixes_copy[i]];
     if (ext_scorer != nullptr) {
       std::vector<int> output;
       std::vector<int> timesteps;
-      prefixes[i]->get_path_vec(output, timesteps);
+      prefixes_copy[i]->get_path_vec(output, timesteps);
       auto prefix_length = output.size();
       auto words = ext_scorer->split_labels(output);
       // remove word insert
       approx_ctc = approx_ctc - prefix_length * ext_scorer->beta;
       // remove language model weight:
       approx_ctc -= (ext_scorer->get_sent_log_prob(words)) * ext_scorer->alpha;
     }
-    prefixes[i]->approx_ctc = approx_ctc;
+    prefixes_copy[i]->approx_ctc = approx_ctc;
   }
 
-  return get_beam_search_result(prefixes, beam_size);
+  return get_beam_search_result(prefixes_copy, beam_size);
+}
+
+std::vector<std::pair<double, Output>> ctc_beam_search_decoder(
+    const std::vector<std::vector<double>> &probs_seq,
+    const std::vector<std::string> &vocabulary,
+    size_t beam_size,
+    double cutoff_prob,
+    size_t cutoff_top_n,
+    size_t blank_id,
+    int log_input,
+    Scorer *ext_scorer)
+{
+  DecoderState state(vocabulary, beam_size, cutoff_prob, cutoff_top_n, blank_id,
+                     log_input, ext_scorer);
+  state.next(probs_seq);
+  return state.decode();
 }
 
 
@@ -199,7 +236,8 @@ ctc_beam_search_decoder_batch(
     size_t cutoff_top_n,
     size_t blank_id,
     int log_input,
-    Scorer *ext_scorer) {
+    Scorer *ext_scorer)
+{
   VALID_CHECK_GT(num_processes, 0, "num_processes must be nonnegative!");
   // thread pool
   ThreadPool pool(num_processes);

ctcdecode/src/ctc_beam_search_decoder.h

@@ -64,4 +64,57 @@ ctc_beam_search_decoder_batch(
     int log_input = 0,
     Scorer *ext_scorer = nullptr);
 
+class DecoderState
+{
+  int abs_time_step;
+  int space_id;
+  size_t beam_size;
+  double cutoff_prob;
+  size_t cutoff_top_n;
+  size_t blank_id;
+  int log_input;
+  std::vector<std::string> vocabulary;
+  Scorer *ext_scorer;
+
+  std::vector<PathTrie*> prefixes;
+  PathTrie root;
+
+public:
+  /* Initialize CTC beam search decoder for streaming
+   *
+   * Parameters:
+   *     vocabulary: A vector of vocabulary.
+   *     beam_size: The width of beam search.
+   *     cutoff_prob: Cutoff probability for pruning.
+   *     cutoff_top_n: Cutoff number for pruning.
+   *     ext_scorer: External scorer to evaluate a prefix, which consists of
+   *                 n-gram language model scoring and word insertion term.
+   *                 Default null, decoding the input sample without scorer.
+  */
+  DecoderState(const std::vector<std::string> &vocabulary,
+               size_t beam_size,
+               double cutoff_prob,
+               size_t cutoff_top_n,
+               size_t blank_id,
+               int log_input,
+               Scorer *ext_scorer);
+  ~DecoderState() = default;
+
+  /* Process logits in decoder stream
+   *
+   * Parameters:
+   *     probs: 2-D vector where each element is a vector of probabilities
+   *               over alphabet of one time step.
+  */
+  void next(const std::vector<std::vector<double>> &probs_seq);
+
+  /* Get current transcription from the decoder stream state
+   *
+   * Return:
+   *     A vector where each element is a pair of score and decoding result,
+   *     in descending order.
+  */
+  std::vector<std::pair<double, Output>> decode() const;
+};
+
 #endif  // CTC_BEAM_SEARCH_DECODER_H_

ctcdecode/src/decoder_utils.cpp

@@ -129,6 +129,19 @@ bool prefix_compare(const PathTrie *x, const PathTrie *y) {
   }
 }
 
+bool prefix_compare_external_scores(const PathTrie *x, const PathTrie *y,
+                                    const std::unordered_map<const PathTrie*, float>& scores) {
+  if (scores.at(x) == scores.at(y)) {
+    if (x->character == y->character) {
+      return false;
+    } else {
+      return (x->character < y->character);
+    }
+  } else {
+    return scores.at(x) > scores.at(y);
+  }
+}
+
 void add_word_to_fst(const std::vector<int> &word,
                      fst::StdVectorFst *dictionary) {
   if (dictionary->NumStates() == 0) {

ctcdecode/src/decoder_utils.h

@@ -1,7 +1,10 @@
 #ifndef DECODER_UTILS_H_
 #define DECODER_UTILS_H_
 
+#include <unordered_map>
 #include <utility>
+#include <vector>
+
 #include "fst/log.h"
 #include "path_trie.h"
 #include "output.h"
@@ -62,9 +65,12 @@ std::vector<std::pair<double, Output>> get_beam_search_result(
     const std::vector<PathTrie *> &prefixes,
     size_t beam_size);
 
-// Functor for prefix comparsion
+// Functor for prefix comparison
 bool prefix_compare(const PathTrie *x, const PathTrie *y);
 
+bool prefix_compare_external_scores(const PathTrie *x, const PathTrie *y,
+                                    const std::unordered_map<const PathTrie*, float>& scores);
+
 /* Get length of utf8 encoding string
  * See: http://stackoverflow.com/a/4063229
  */