Merge pull request #2183 from coqui-ai/decoder-wav2vec2-batch

Expose batch version of decoder API for wav2vec2 AM

Author: reuben

.github/workflows/build-and-test.yml

@@ -1567,6 +1567,9 @@ jobs:
         run: |
           sudo rm -r /Library/Developer/CommandLineTools
       - run: ./ci_scripts/host-build.sh ${{ matrix.arch }}
+      - name: Setup tmate session
+        uses: mxschmitt/action-tmate@v3
+        if: failure()
       - run: ./ci_scripts/package.sh
       - uses: actions/upload-artifact@v2
         with:

native_client/ctcdecode/__init__.py

@@ -311,6 +311,78 @@ def ctc_beam_search_decoder_batch(
     return batch_beam_results
 
 
+def ctc_beam_search_decoder_for_wav2vec2am_batch(
+    probs_seq,
+    seq_lengths,
+    alphabet,
+    beam_size,
+    num_threads,
+    cutoff_prob=1.0,
+    cutoff_top_n=40,
+    blank_id=-1,
+    ignored_symbols=frozenset(),
+    scorer=None,
+    hot_words=dict(),
+    num_results=1,
+):
+    """Wrapper for the batched CTC beam search decoder for wav2vec2 AM.
+
+    :param probs_seq: 3-D list with each element as an instance of 2-D list
+                      of probabilities used by ctc_beam_search_decoder().
+    :type probs_seq: 3-D list
+    :param alphabet: alphabet list.
+    :alphabet: Alphabet
+    :param beam_size: Width for beam search.
+    :type beam_size: int
+    :param num_threads: Number of threads to use for processing batch.
+    :type num_threads: int
+    :param cutoff_prob: Cutoff probability in alphabet pruning,
+                        default 1.0, no pruning.
+    :type cutoff_prob: float
+    :param cutoff_top_n: Cutoff number in pruning, only top cutoff_top_n
+                         characters with highest probs in alphabet will be
+                         used in beam search, default 40.
+    :type cutoff_top_n: int
+    :param scorer: External scorer for partially decoded sentence, e.g. word
+                   count or language model.
+    :type scorer: Scorer
+    :param hot_words: Map of words (keys) to their assigned boosts (values)
+    :type hot_words: dict[string, float]
+    :param num_results: Number of beams to return.
+    :type num_results: int
+    :return: List of tuples of confidence and sentence as decoding
+             results, in descending order of the confidence.
+    :rtype: list
+    """
+    batch_beam_results = swigwrapper.ctc_beam_search_decoder_for_wav2vec2am_batch(
+        probs_seq,
+        seq_lengths,
+        alphabet,
+        beam_size,
+        num_threads,
+        cutoff_prob,
+        cutoff_top_n,
+        blank_id,
+        ignored_symbols,
+        scorer,
+        hot_words,
+        num_results,
+    )
+    batch_beam_results = [
+        [
+            DecodeResult(
+                res.confidence,
+                alphabet.Decode(res.tokens),
+                [int(t) for t in res.tokens],
+                [int(t) for t in res.timesteps],
+            )
+            for res in beam_results
+        ]
+        for beam_results in batch_beam_results
+    ]
+    return batch_beam_results
+
+
 class FlashlightDecoderState(swigwrapper.FlashlightDecoderState):
     """
     This class contains constants used to specify the desired behavior for the

native_client/ctcdecode/ctc_beam_search_decoder.cpp

@@ -652,7 +652,7 @@ ctc_beam_search_decoder_batch(
 }
 
 std::vector<std::vector<Output>>
-ctc_beam_search_decoder_batch_for_wav2vec2am(
+ctc_beam_search_decoder_for_wav2vec2am_batch(
     const double *probs,
     int batch_size,
     int time_dim,
@@ -661,7 +661,7 @@ ctc_beam_search_decoder_batch_for_wav2vec2am(
     int seq_lengths_size,
     const Alphabet &alphabet,
     size_t beam_size,
-    size_t num_processes,
+    size_t num_threads,
     double cutoff_prob,
     size_t cutoff_top_n,
     int blank_id,
@@ -670,10 +670,10 @@ ctc_beam_search_decoder_batch_for_wav2vec2am(
     std::unordered_map<std::string, float> hot_words,
     size_t num_results)
 {
-  VALID_CHECK_GT(num_processes, 0, "num_processes must be nonnegative!");
+  VALID_CHECK_GT(num_threads, 0, "num_threads must be nonnegative!");
   VALID_CHECK_EQ(batch_size, seq_lengths_size, "must have one sequence length per batch element");
   // thread pool
-  ThreadPool pool(num_processes);
+  ThreadPool pool(num_threads);
 
   // enqueue the tasks of decoding
   std::vector<std::future<std::vector<Output>>> res;

native_client/ctcdecode/ctc_beam_search_decoder.h

@@ -333,7 +333,7 @@ ctc_beam_search_decoder_batch(
  *     result for one audio sample.
 */
 std::vector<std::vector<Output>>
-ctc_beam_search_decoder_batch_for_wav2vec2am(
+ctc_beam_search_decoder_for_wav2vec2am_batch(
     const double* probs,
     int batch_size,
     int time_dim,

native_client/definitions.mk

@@ -38,11 +38,7 @@ CFLAGS                  := -mmacosx-version-min=10.10 -target x86_64-apple-macos
 LDFLAGS                 := -mmacosx-version-min=10.10 -target x86_64-apple-macos10.10
 
 SOX_CFLAGS              := $(shell pkg-config --cflags sox)
-LIBSOX_PATH             := $(shell echo `pkg-config --libs-only-L sox | sed -e 's/^-L//'`/lib`pkg-config --libs-only-l sox | sed -e 's/^-l//'`.dylib)
-LIBOPUSFILE_PATH        := $(shell echo `pkg-config --libs-only-L opusfile | sed -e 's/^-L//'`/lib`pkg-config --libs-only-l opusfile | sed -e 's/^-l//'`.dylib)
-LIBSOX_STATIC_DEPS      := $(shell echo `otool -L $(LIBSOX_PATH) | tail -n +2 | cut -d' ' -f1 | grep /opt/ | sed -E "s/\.[[:digit:]]+\.dylib/\.a/" | tr '\n' ' '`)
-LIBOPUSFILE_STATIC_DEPS := $(shell echo `otool -L $(LIBOPUSFILE_PATH) | tail -n +2 | cut -d' ' -f1 | grep /opt/ | sed -E "s/\.[[:digit:]]+\.dylib/\.a/" | tr '\n' ' '`)
-SOX_LDFLAGS             := $(LIBSOX_STATIC_DEPS) $(LIBOPUSFILE_STATIC_DEPS) -framework CoreAudio -lz
+SOX_LDFLAGS             := $(shell pkg-config --libs sox) -framework CoreAudio -lz
 else
 SOX_LDFLAGS     := `pkg-config --libs sox`
 endif # OS others
@@ -128,11 +124,7 @@ CFLAGS                  := -mmacosx-version-min=11.0 -target arm64-apple-macos11
 LDFLAGS                 := -mmacosx-version-min=11.0 -target arm64-apple-macos11
 
 SOX_CFLAGS              := $(shell arm-pkg-config --cflags sox)
-LIBSOX_PATH             := $(shell echo `arm-pkg-config --libs-only-L sox | sed -e 's/^-L//'`/lib`arm-pkg-config --libs-only-l sox | sed -e 's/^-l//'`.dylib)
-LIBOPUSFILE_PATH        := $(shell echo `arm-pkg-config --libs-only-L opusfile | sed -e 's/^-L//'`/lib`arm-pkg-config --libs-only-l opusfile | sed -e 's/^-l//'`.dylib)
-LIBSOX_STATIC_DEPS      := $(shell echo `otool -L $(LIBSOX_PATH) | tail -n +2 | cut -d' ' -f1 | grep /opt/ | sed -E "s/\.[[:digit:]]+\.dylib/\.a/" | tr '\n' ' '`)
-LIBOPUSFILE_STATIC_DEPS := $(shell echo `otool -L $(LIBOPUSFILE_PATH) | tail -n +2 | cut -d' ' -f1 | grep /opt/ | sed -E "s/\.[[:digit:]]+\.dylib/\.a/" | tr '\n' ' '`)
-SOX_LDFLAGS             := $(LIBSOX_STATIC_DEPS) $(LIBOPUSFILE_STATIC_DEPS) -framework CoreAudio -lz
+SOX_LDFLAGS             := $(shell arm-pkg-config --libs sox) -framework CoreAudio -lz
 endif
 
 # -Wl,--no-as-needed is required to force linker not to evict libs it thinks we