Repository: nvidia-riva/riva-asrlib-decoder (language: python)

Issue:
Error building from source.
Hi @galv, I have the intention to plug the `k2` HLG decoding graph into your amazing decoding algorithm. So I tried to build `riva-asrlib-decoder` from source and here is my environment detected during configuration:

```
$ cmake -DRIVA_ASRLIB_BUILD_PYTHON_BINDINGS=NO ..
-- The CUDA compiler identification is NVIDIA 11.7.99
-- The CXX compiler identification is GNU 11.3.0
-- Detecting CUDA compiler ABI info
-- Detecting CUDA compiler ABI info - done
-- Check for working CUDA compiler: /usr/local/cuda/bin/nvcc - skipped
-- Detecting CUDA compile features
-- Detecting CUDA compile features - done
-- Detecting CXX compiler ABI info
-- Detecting CXX compiler ABI info - done
-- Check for working CXX compiler: /usr/bin/c++ - skipped
-- Detecting CXX compile features
-- Detecting CXX compile features - done
-- Performing Test CMAKE_HAVE_LIBC_PTHREAD
-- Performing Test CMAKE_HAVE_LIBC_PTHREAD - Success
-- Found Threads: TRUE
-- Found CUDAToolkit: /usr/local/cuda/include (found version "11.7.99")
-- Configuring done
-- Generating done
-- Build files have been written to: /home/english_k2/riva-asrlib-decoder/build
```

Now the building process was successful until it reached `batched-mapped-decoder-cuda` compilation, where it failed with errors. The stack trace is long, so I attach the log file.
[build_log.txt](https://github.com/nvidia-riva/riva-asrlib-decoder/files/10326635/build_log.txt)

Thanks in advance for the help.


Discussion / hints:

Fix patch:
diff --git a/Dockerfile b/Dockerfile
index b4a01ba..95dafe6 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -13,6 +13,7 @@ FROM ${FROM_IMAGE_NAME}
 ENV PATH=/usr/local/cuda/bin:${PATH}
 
 ENV NVIDIA_DRIVER_CAPABILITIES=video,compute,utility
+ENV NVIDIA_VISIBLE_DEVICES=all
 
 # CUDA
 COPY --from=cuda /usr/local/cuda /usr/local/cuda
diff --git a/include/riva/asrlib/decoder/batched-mapped-decoder-cuda.h b/include/riva/asrlib/decoder/batched-mapped-decoder-cuda.h
index 338a61e..5a4f60b 100644
--- a/include/riva/asrlib/decoder/batched-mapped-decoder-cuda.h
+++ b/include/riva/asrlib/decoder/batched-mapped-decoder-cuda.h
@@ -48,9 +48,7 @@ class BatchedMappedDecoderCuda {
       std::size_t logits_n_input_frames_valid,
       // this should not be copied by reference. It should be passed as an
       // lvalue or an rvalue because it is always copied.
-      const std::function<void(std::tuple<
-                               std::optional<kaldi::CompactLattice>,
-                               std::optional<kaldi::cuda_decoder::CTMResult>>&)>& callback);
+      const BatchedMappedOnlineDecoderCuda::LatticeCallback& callback);
 
   void ComputeTasks();
   void AcquireTasks();
@@ -96,10 +94,7 @@ class BatchedMappedDecoderCuda {
     std::size_t logits_frame_stride;
     std::size_t logits_n_input_frames_valid;
     CorrelationID corr_id;
-    std::function<void(
-        std::tuple<
-            std::optional<kaldi::CompactLattice>, std::optional<kaldi::cuda_decoder::CTMResult>>&)>
-        callback;
+    BatchedMappedOnlineDecoderCuda::LatticeCallback callback;
 
     std::size_t loglikes_time_offset = std::size_t(0);
   };
diff --git a/include/riva/asrlib/decoder/batched-mapped-online-decoder-cuda.h b/include/riva/asrlib/decoder/batched-mapped-online-decoder-cuda.h
index cf0812f..964d6c9 100644
--- a/include/riva/asrlib/decoder/batched-mapped-online-decoder-cuda.h
+++ b/include/riva/asrlib/decoder/batched-mapped-online-decoder-cuda.h
@@ -20,6 +20,7 @@
 #include <cudadecoder/cuda-decoder-common.h>
 #include <cudadecoder/cuda-decoder.h>
 #include <cudadecoder/lattice-postprocessor.h>
+#include <optional>
 #include <fstext/lattice-utils.h>
 #include <itf/options-itf.h>
 #include <itf/transition-information.h>
@@ -117,9 +118,11 @@ struct BatchedMappedOnlineDecoderCudaConfig {
 class BatchedMappedOnlineDecoderCuda {
  public:
   using CorrelationID = uint64_t;
-  using LatticeCallback = std::function<void(
-      std::tuple<
-          std::optional<kaldi::CompactLattice>, std::optional<kaldi::cuda_decoder::CTMResult>>&)>;
+  using ReturnType =
+    std::tuple<std::optional<kaldi::CompactLattice>,
+               std::optional<kaldi::cuda_decoder::CTMResult>,
+               std::optional<std::vector<kaldi::cuda_decoder::NBestResult>>>;
+  using LatticeCallback = std::function<void(ReturnType&)>;
   BatchedMappedOnlineDecoderCuda(
       const BatchedMappedOnlineDecoderCudaConfig& config, const fst::Fst<fst::StdArc>& decode_fst,
       std::unique_ptr<kaldi::TransitionInformation>&& trans_information)
@@ -160,11 +163,10 @@ class BatchedMappedOnlineDecoderCuda {
     if (num_worker_threads > 0) {
       thread_pool_ = std::make_unique<kaldi::futures_thread_pool>(num_worker_threads);
     }
-    // TODO: Uncomment this
-    // if (config_.num_decoder_copy_threads > 0) {
-    //   cuda_decoder_->SetThreadPoolAndStartCPUWorkers(
-    //       thread_pool_.get(), config_.num_decoder_copy_threads);
-    // }
+    if (config_.num_decoder_copy_threads > 0) {
+      cuda_decoder_->SetThreadPoolAndStartCPUWorkers(
+          thread_pool_.get(), config_.num_decoder_copy_threads);
+    }
   }
 
   ~BatchedMappedOnlineDecoderCuda()
@@ -196,7 +198,6 @@ class BatchedMappedOnlineDecoderCuda {
       channels = &channels_;
     }
     ListIChannelsInBatch(corr_ids, &channels_);
-    // KALDI_LOG << "GALVEZ:CHANNELS:" << channels_[0];
 
     std::vector<int> list_channels_first_chunk;
     for (std::size_t i = 0; i < is_first_chunk.size(); ++i) {
@@ -348,6 +349,7 @@ class BatchedMappedOnlineDecoderCuda {
 
 
  private:
+  // could use this as an opportunity to clear saved variables for the channel
   void FinalizeDecoding(int32 ichannel, const LatticeCallback* callback)
   {
     kaldi::Lattice lat;
@@ -373,17 +375,17 @@ class BatchedMappedOnlineDecoderCuda {
 
     // this is wasteful, since GetCTM calls GetPostprocessedLattice
     // but doesn't return the result.
-    // okay, so I am basically doing both here...
-    std::tuple<std::optional<kaldi::CompactLattice>, std::optional<kaldi::cuda_decoder::CTMResult>>
-        result;
+    ReturnType result;
     if (config_.use_lattice_postprocessor) {
         kaldi::CompactLattice clat;
         lattice_postprocessor_->GetPostprocessedLattice(dlat, &clat);
         kaldi::cuda_decoder::CTMResult ctm;
+        // is this okay? I am modifying the refernece twice...
         lattice_postprocessor_->GetCTM(dlat, &ctm);
-        result = {std::make_optional(clat), std::make_optional(ctm)}; 
+        std::vector<kaldi::cuda_decoder::NBestResult> nbest = lattice_postprocessor_->GetNBestList(dlat);
+        result = {std::make_optional(clat), std::make_optional(ctm), std::make_optional(nbest)};
     } else {
-        result = {std::make_optional(dlat), std::nullopt};
+        result = {std::make_optional(dlat), std::nullopt, std::nullopt};
     }
 
     // if ptr set and if callback func callable
@@ -395,14 +397,6 @@ class BatchedMappedOnlineDecoderCuda {
     n_lattice_callbacks_not_done_.fetch_sub(1, std::memory_order_release);
   }
 
-  // static wrapper for thread pool
-  static void FinalizeDecodingWrapper(void* obj, uint64_t ichannel64, void* callback_ptr)
-  {
-    int32 ichannel = static_cast<int32>(ichannel64);
-    const LatticeCallback* callback = static_cast<const LatticeCallback*>(callback_ptr);
-    static_cast<BatchedMappedOnlineDecoderCuda*>(obj)->FinalizeDecoding(ichannel, callback);
-  }
-
   void ListIChannelsInBatch(const std::vector<CorrelationID>& corr_ids, std::vector<int>* channels)
   {
     channels->clear();
diff --git a/pyproject.toml b/pyproject.toml
index 174cc43..27127fc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,6 +17,7 @@ environment = 'RIVA_ASRLIB_DECODER_CMAKE_ARGS="-DCMAKE_BUILD_TYPE=Release" CUDAA
 # it uses setup.py, not pyproject.toml
 before-test = "pip install Cython"
 test-command = "pytest {project}/src/riva/asrlib/decoder/test_decoder.py"
-test-requires = ["pytest", "kaldi-io", "more-itertools", "nemo_toolkit[asr]"]
-# fails because onnx doesn't provide python3.6 wheels: https://pypi.org/project/onnx/#files
-test-skip = "cp36-manylinux_x86_64"
+test-requires = ["pytest", "kaldi-io", "more-itertools", "nemo_toolkit[asr]", "torchaudio"]
+# 3.6 fails because onnx doesn't provide python3.6 wheels: https://pypi.org/project/onnx/#files
+# 3.11 fails becasue of numba: Cannot install on Python version 3.11.1; only versions >=3.7,<3.11 are supported
+test-skip = "{cp36-manylinux_x86_64,cp311-manylinux_x86_64}"
diff --git a/setup.py b/setup.py
index 11b42af..b24ad40 100644
--- a/setup.py
+++ b/setup.py
@@ -95,7 +95,7 @@ def build_extension(self, ext: setuptools.extension.Extension):
 setuptools.setup(
     python_requires='>=3.7',
     name='riva-asrlib-decoder',
-    version='0.1.0',
+    version='0.2.0',
     author='NVIDIA',
     author_email='dgalvez@nvidia.com',
     keywords='ASR, CUDA, WFST, Decoder',
@@ -132,6 +132,7 @@ def build_extension(self, ext: setuptools.extension.Extension):
             "kaldi-io",
             "more-itertools",
             "nemo_toolkit[asr]",
+            "torchaudio",
         ]
     },
     zip_safe=False,
diff --git a/src/riva/asrlib/decoder/batched-mapped-decoder-cuda.cc b/src/riva/asrlib/decoder/batched-mapped-decoder-cuda.cc
index f1b1871..d6bdc04 100644
--- a/src/riva/asrlib/decoder/batched-mapped-decoder-cuda.cc
+++ b/src/riva/asrlib/decoder/batched-mapped-decoder-cuda.cc
@@ -43,9 +43,7 @@ BatchedMappedDecoderCuda::~BatchedMappedDecoderCuda()
 void
 BatchedMappedDecoderCuda::DecodeWithCallback(
     const float* d_logits, std::size_t logits_frame_stride, std::size_t logits_n_input_frames_valid,
-    const std::function<void(std::tuple<
-                             std::optional<kaldi::CompactLattice>,
-                             std::optional<kaldi::cuda_decoder::CTMResult>>&)>& callback)
+    const BatchedMappedOnlineDecoderCuda::LatticeCallback& callback)
 {
   UtteranceTask task;
   // at 5000 files/s, expected to overflow in ~116 million years
@@ -81,9 +79,7 @@ BatchedMappedDecoderCuda::AcquireTasks()
     auto& callback = task.callback;
 
     cuda_online_pipeline_.SetLatticeCallback(
-        task.corr_id, [this, callback](std::tuple<
-                                       std::optional<kaldi::CompactLattice>,
-                                       std::optional<kaldi::cuda_decoder::CTMResult>>& result) {
+        task.corr_id, [this, callback](BatchedMappedOnlineDecoderCuda::ReturnType& result) {
           if (callback)
             callback(result);
           n_tasks_not_done_.fetch_sub(1, std::memory_order_release);
diff --git a/src/riva/asrlib/decoder/bin/offline-cuda-decode-binary.cc b/src/riva/asrlib/decoder/bin/offline-cuda-decode-binary.cc
index 6c46b55..fd38e22 100644
--- a/src/riva/asrlib/decoder/bin/offline-cuda-decode-binary.cc
+++ b/src/riva/asrlib/decoder/bin/offline-cuda-decode-binary.cc
@@ -150,9 +150,7 @@ main(int argc, char** argv)
     // must copy the captured variables. Otherwise, they will be lost.
     auto callback = [key, d_loglikes, &output_writer_m, &clat_writer, &ctm_output_path,
                      &word_syms = cuda_pipeline.GetSymbolTable()](
-                        std::tuple<
-                            std::optional<kaldi::CompactLattice>,
-                            std::optional<kaldi::cuda_decoder::CTMResult>>& result) {
+                        BatchedMappedOnlineDecoderCuda::ReturnType& result) {
       CU_SAFE_CALL(cudaFree(d_loglikes));
 
       if (std::get<0>(result).has_value()) {
diff --git a/src/riva/asrlib/decoder/python_decoder.cc b/src/riva/asrlib/decoder/python_decoder.cc
index d66b635..96ed8ec 100644
--- a/src/riva/asrlib/decoder/python_decoder.cc
+++ b/src/riva/asrlib/decoder/python_decoder.cc
@@ -261,9 +261,7 @@ PybindBatchedMappedDecoderCuda(py::module& m)
           // stride of each row is stride. Always greater than number of cols
           auto write_results =
               [i, &clat_writer, &keys](
-                  std::tuple<
-                      std::optional<kaldi::CompactLattice>,
-                      std::optional<kaldi::cuda_decoder::CTMResult>>& asr_results) {
+                  riva::asrlib::BatchedMappedOnlineDecoderCuda::ReturnType& asr_results) {
               const kaldi::CompactLattice& lattice = std::get<0>(asr_results).value();
               clat_writer.Write(keys[i], lattice);
               };
@@ -278,10 +276,10 @@ PybindBatchedMappedDecoderCuda(py::module& m)
 
 
   pyclass.def(
-      "decode",
+      "decode_mbr",
       [](PyClass& cuda_pipeline, const DLManagedTensor* managed_logits,
          const DLManagedTensor* managed_logits_lengths)
-          -> std::vector<std::vector<std::tuple<std::string, float, float>>> {
+          -> std::vector<std::vector<std::tuple<std::string, float, float, float>>> {
         // contiguousness might not mean what I think it means. It may just mean
         // stride has no padding.
 
@@ -313,7 +311,7 @@ PybindBatchedMappedDecoderCuda(py::module& m)
         }
         // logits should be batch x time x logits
         int64_t batch_size = logits_lengths.shape[0];
-        std::vector<std::vector<std::tuple<std::string, float, float>>> results(batch_size);
+        std::vector<std::vector<std::tuple<std::string, float, float, float>>> results(batch_size);
         for (int64_t i = 0; i < batch_size; ++i) {
           int64_t valid_time_steps = index<int64_t>(logits_lengths, i);
 
@@ -324,15 +322,14 @@ PybindBatchedMappedDecoderCuda(py::module& m)
           // stride of each row is stride. Always greater than number of cols
           auto place_results =
               [i, &results, &word_syms = cuda_pipeline.GetSymbolTable()](
-                  std::tuple<
-                      std::optional<kaldi::CompactLattice>,
-                      std::optional<kaldi::cuda_decoder::CTMResult>>& asr_results) {
+                  riva::asrlib::BatchedMappedOnlineDecoderCuda::ReturnType& asr_results) {
                 const kaldi::cuda_decoder::CTMResult& ctm_result = std::get<1>(asr_results).value();
                 for (size_t iword = 0; iword < ctm_result.times_seconds.size(); ++iword) {
                   results[i].emplace_back(
                       word_syms.Find(ctm_result.words[iword]),
                       ctm_result.times_seconds[iword].first,
-                      ctm_result.times_seconds[iword].second);
+                      ctm_result.times_seconds[iword].second,
+                      ctm_result.conf[iword]);
                 }
               };
           cuda_pipeline.DecodeWithCallback(
@@ -347,6 +344,75 @@ PybindBatchedMappedDecoderCuda(py::module& m)
         cuda_pipeline.WaitForAllTasks();
         return results;
       });
+
+  pyclass.def("decode_map",
+              [](PyClass& cuda_pipeline, const DLManagedTensor* managed_logits,
+                 const DLManagedTensor* managed_logits_lengths, size_t n)
+              -> std::vector<std::vector<std::tuple<std::string, float, float, float>>>
+              {
+        const DLTensor& logits = managed_logits->dl_tensor;
+        const DLTensor& logits_lengths = managed_logits_lengths->dl_tensor;
+        if (logits.ndim != 3) {
+          throw std::invalid_argument("Expected a 3D logits tensor");
+        }
+        if (logits.dtype.code != kDLFloat || logits.dtype.bits != 32 || logits.dtype.lanes != 1) {
+          throw std::invalid_argument("Expected a float32 logits tensor");
+        }
+        // TODO: Consider setting device id based on this... Could use RAII like
+        // in K2 for that.
+        if (logits.device.device_type != kDLCUDA) {
+          throw std::invalid_argument("Expected logits tensor to be on GPU");
+        }
+        if (logits_lengths.ndim != 1) {
+          throw std::invalid_argument("Expected a 1D logits lengths tensor");
+        }
+        if (logits_lengths.dtype.code != kDLInt || logits_lengths.dtype.bits != 64 ||
+            logits_lengths.dtype.lanes != 1) {
+          throw std::invalid_argument("Expected a 64-bit signed integer logits lengths tensor");
+        }
+        if (logits_lengths.device.device_type != kDLCPU) {
+          throw std::invalid_argument("Expected lengths tensor to be on CPU");
+        }
+        // logits should be batch x time x logits
+        int64_t batch_size = logits_lengths.shape[0];
+        std::vector<std::vector<std::tuple<std::string, float, float, float>>> results(batch_size);
+        for (int64_t i = 0; i < batch_size; ++i) {
+          int64_t valid_time_steps = index<int64_t>(logits_lengths, i);
+
+          // this may not be right... Yes, it seems quite wrong...
+          const float* single_sample_logits_start = address<float>(logits, i, 0);
+          // number of rows is number of frames
+          // number of cols is number of logits
+          // stride of each row is stride. Always greater than number of cols
+          auto place_results =
+              [i, &results, &word_syms = cuda_pipeline.GetSymbolTable()](
+                  riva::asrlib::BatchedMappedOnlineDecoderCuda::ReturnType& asr_results) {
+                const std::vector<kaldi::cuda_decoder::NBestResult>& nbest_results = std::get<2>(asr_results).value();
+                std::vector<
+                std::tuple<float,
+                           std::vector<std::string>,
+                           std::vector<std::pair<float, float>>
+                >> result_this_utt;
+                for (const kaldi::cuda_decoder::NBestResult& nbest_result: nbest_results) {
+                  std::vector<std::string> words; words.reserve(nbest_result.words.size());
+                  for (auto&& word_id: nbest_result.words) {
+                    words.emplace_back(word_syms.Find(word_id));
+                  }
+                  result_this_utt.emplace_back(nbest_result.score, words, std::move(nbest_result.times_seconds));
+                }
+          };
+          cuda_pipeline.DecodeWithCallback(
+              single_sample_logits_start,
+              // single_sample_logits.data_ptr<float>(),
+              stride(logits, 1),
+              // single_sample_logits.stride(1),
+              index<int64_t>(logits_lengths, i),
+              // size(0)
+              place_results);
+        }
+        cuda_pipeline.WaitForAllTasks();
+        return results;
+});
 }
 }  // anonymous namespace
 


Test patch:
diff --git a/src/riva/asrlib/decoder/test_decoder.py b/src/riva/asrlib/decoder/test_decoder.py
index 936d843..1a979ff 100644
--- a/src/riva/asrlib/decoder/test_decoder.py
+++ b/src/riva/asrlib/decoder/test_decoder.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import itertools
+import multiprocessing
 import os
 import pathlib
 import shutil
@@ -41,7 +42,7 @@ def test_decoder(self):
         test_data_dir = pathlib.Path(__file__).parent.resolve() / "test_data"
         download_test_data(test_data_dir)
 
-        logits_ark = str("ark:" / test_data_dir / "logits.ark")
+        logits_ark = str("ark,bg:" / test_data_dir / "logits.ark")
 
         _, matrix = next(kaldi_io.read_mat_ark(logits_ark))
         num_tokens_including_blank = matrix.shape[1]
@@ -53,21 +54,24 @@ def test_decoder(self):
         config.online_opts.decoder_opts.lattice_beam = 8.0
         config.online_opts.decoder_opts.max_active = 7000
         config.online_opts.determinize_lattice = True
-        config.online_opts.max_batch_size = 400
-        config.online_opts.num_channels = 800
+        config.online_opts.max_batch_size = 200
+        config.online_opts.num_channels = config.online_opts.max_batch_size * 2
         config.online_opts.frame_shift_seconds = 0.03
+        config.online_opts.num_post_processing_worker_threads = multiprocessing.cpu_count()
+        config.online_opts.num_decoder_copy_threads = 4
         decoder = BatchedMappedDecoderCuda(
             config, str(test_data_dir / "TLG.fst"), str(test_data_dir / "words.txt"), num_tokens_including_blank
         )
 
-        for batch in more_itertools.chunked(kaldi_io.read_mat_ark(logits_ark), 1):
+        for batch in more_itertools.chunked(kaldi_io.read_mat_ark(logits_ark), config.online_opts.max_batch_size):
             sequences = []
             sequence_lengths = []
             for key, matrix in batch:
                 sequences.append(torch.from_numpy(matrix.copy()))
                 sequence_lengths.append(matrix.shape[0])
+            # This could be faster with pinned memory, but this unit test is fast enough right now.
             padded_sequence = torch.nn.utils.rnn.pad_sequence(sequences, batch_first=True).cuda()
             sequence_lengths_tensor = torch.tensor(sequence_lengths, dtype=torch.long)
 
-            for result in decoder.decode(padded_sequence, sequence_lengths_tensor):
-                pass
+            # ignore output
+            decoder.decode_mbr(padded_sequence, sequence_lengths_tensor)
diff --git a/src/riva/asrlib/decoder/test_graph_construction.py b/src/riva/asrlib/decoder/test_graph_construction.py
index a92400f..074b86c 100644
--- a/src/riva/asrlib/decoder/test_graph_construction.py
+++ b/src/riva/asrlib/decoder/test_graph_construction.py
@@ -14,69 +14,78 @@
 # limitations under the License.
 
 import gzip
+from contextlib import ExitStack
 import json
+import multiprocessing
 import os
 import pathlib
+import pytest
 import shutil
 import subprocess
 import tarfile
 import tempfile
+import time
 import unittest
 import zipfile
 
 import more_itertools
 import nemo.collections.asr as nemo_asr
+import numpy as np
 from ruamel.yaml import YAML
 import torch
+import torchaudio
+import torchmetrics
 from tqdm import tqdm
 
 import riva.asrlib.decoder
 from riva.asrlib.decoder.python_decoder import BatchedMappedDecoderCuda, BatchedMappedDecoderCudaConfig
 
-class GraphConstructionTest(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.temp_dir = os.path.abspath("tmp_graph_construction")
-        os.makedirs(cls.temp_dir, exist_ok=True)
+class TestGraphConstruction:
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        # return_type = namedtuple(["temp_dir", "words_path", "gzipped_lm_path", "nemo_model_path", "units_txt", "num_tokens_including_blank", "dataset_map"])
+        self.temp_dir = os.path.abspath("tmp_graph_construction")
+        os.makedirs(self.temp_dir, exist_ok=True)
 
-        lm_zip_file = os.path.join(cls.temp_dir, "speechtotext_english_lm_deployable_v1.0.zip")
+        lm_zip_file = os.path.join(self.temp_dir, "speechtotext_english_lm_deployable_v1.0.zip")
         if not os.path.exists(lm_zip_file):
             subprocess.check_call(
                 f"wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/tao/speechtotext_english_lm/versions/deployable_v1.0/zip -O {lm_zip_file}",
                 shell=True,
             )
             with zipfile.ZipFile(lm_zip_file, 'r') as zip_ref:
-                zip_ref.extractall(cls.temp_dir)
+                zip_ref.extractall(self.temp_dir)
 
-        am_zip_file = os.path.join(cls.temp_dir, "stt_en_conformer_ctc_small_1.6.0.zip")
+        am_zip_file = os.path.join(self.temp_dir, "stt_en_conformer_ctc_small_1.6.0.zip")
         if not os.path.exists(am_zip_file):
             subprocess.check_call(
                 f"wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_en_conformer_ctc_small/versions/1.6.0/zip -O {am_zip_file}",
                 shell=True,
             )
             with zipfile.ZipFile(am_zip_file, 'r') as zip_ref:
-                zip_ref.extractall(cls.temp_dir)
+                zip_ref.extractall(self.temp_dir)
 
         # Work around: At the time of writing this test, the words.txt
         # file downloaded from NGC is simply a git lfs stub file, not
-        # the actual file itself, so overwrite cls.words_path by
+        # the actual file itself, so overwrite self.words_path by
         # exracting the symbol table from the arpa file
-        lm_path = os.path.join(cls.temp_dir, "3-gram.pruned.3e-7.arpa")
-        cls.words_path = os.path.join(cls.temp_dir, "words.mixed_lm.3-gram.pruned.3e-7.txt")
-        temp_words_path = os.path.join(cls.temp_dir, "words_with_ids.txt")
-        subprocess.check_call(
-            [
-                os.path.join(riva.asrlib.decoder.__path__[0], "scripts/prepare_TLG_fst/bin/arpa2fst"),
-                f"--write-symbol-table={temp_words_path}",
-                lm_path,
-                "/dev/null",
-            ]
-        )
-        cls.gzipped_lm_path = lm_path + ".gz"
-        with open(lm_path, 'rb') as f_in, gzip.open(cls.gzipped_lm_path, 'wb') as f_out:
+        lm_path = os.path.join(self.temp_dir, "3-gram.pruned.3e-7.arpa")
+        self.words_path = os.path.join(self.temp_dir, "words.mixed_lm.3-gram.pruned.3e-7.txt")
+        temp_words_path = os.path.join(self.temp_dir, "words_with_ids.txt")
+        if not os.path.exists(temp_words_path):
+            subprocess.check_call(
+                [
+                    os.path.join(riva.asrlib.decoder.__path__[0], "scripts/prepare_TLG_fst/bin/arpa2fst"),
+                    f"--write-symbol-table={temp_words_path}",
+                    lm_path,
+                    "/dev/null",
+                ]
+            )
+        self.gzipped_lm_path = lm_path + ".gz"
+        with open(lm_path, 'rb') as f_in, gzip.open(self.gzipped_lm_path, 'wb') as f_out:
             shutil.copyfileobj(f_in, f_out)
 
-        with open(temp_words_path, "r") as words_with_ids_fh, open(cls.words_path, "w") as words_fh:
+        with open(temp_words_path, "r") as words_with_ids_fh, open(self.words_path, "w") as words_fh:
             for word_with_id in words_with_ids_fh:
                 word = word_with_id.split()[0].lower()
                 if word in {"<eps>", "
", "
", "<unk>"}:
@@ -84,32 +93,46 @@ def setUpClass(cls):
                 words_fh.write(word)
                 words_fh.write("\n")
 
-        cls.nemo_model_path = os.path.join(cls.temp_dir, "stt_en_conformer_ctc_small.nemo")
-        config_yaml = os.path.join(cls.temp_dir, "model_config.yaml")
+        self.nemo_model_path = os.path.join(self.temp_dir, "stt_en_conformer_ctc_small.nemo")
+        config_yaml = os.path.join(self.temp_dir, "model_config.yaml")
 
         yaml = YAML(typ='safe')
-        with tarfile.open(cls.nemo_model_path, "r:gz") as tar_fh:
+        with tarfile.open(self.nemo_model_path, "r:gz") as tar_fh:
             with tar_fh.extractfile("./model_config.yaml") as fh:
                 data = yaml.load(fh)
-        cls.units_txt = os.path.join(cls.temp_dir, "units.txt")
-        with open(cls.units_txt, "w") as fh:
+        self.units_txt = os.path.join(self.temp_dir, "units.txt")
+        with open(self.units_txt, "w") as fh:
             for unit in data["decoder"]["vocabulary"]:
                 fh.write(f"{unit}\n")
 
-        cls.num_tokens_including_blank = len(data["decoder"]["vocabulary"]) + 1
-        assert cls.num_tokens_including_blank == 1025
+        self.num_tokens_including_blank = len(data["decoder"]["vocabulary"]) + 1
+        assert self.num_tokens_including_blank == 1025
+
+        librispeech_test_clean = torchaudio.datasets.LIBRISPEECH(self.temp_dir, "test-clean", download=True)
+        librispeech_test_other = torchaudio.datasets.LIBRISPEECH(self.temp_dir, "test-other", download=True)
+        self.dataset_map = {
+            "test_clean": librispeech_test_clean,
+            "test_other": librispeech_test_other,
+        }
 
-    @classmethod
-    def tearDownClass(cls):
-        pass
-        # shutil.rmtree(cls.temp_dir)
 
     def test_eesen_ctc_topo(self):
         self.create_TLG("ctc_eesen", os.path.join(self.temp_dir, "ctc_eesen"))
 
-    def test_vanilla_ctc_topo(self):
-        self.create_TLG("ctc", os.path.join(self.temp_dir, "ctc"))
-        self.run_decoder(os.path.join(self.temp_dir, "ctc"))
+    # TODO: Debug why these WERs are a bit higher than the ones reported here for "N-gram LM"
+    # https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/stt_en_conformer_ctc_small
+    # TODO: Investigate
+    @pytest.mark.parametrize("dataset, expected_wer, half_precision",
+                             [# ("test_clean", 0.03509205721241631, False),
+                              ("test_clean", 0.035263237979306146, True),
+                              # Causes a crash when I run twice...
+                              # ("test_other", 0.07034369447681639, False)
+                             ])
+    def test_vanilla_ctc_topo(self, dataset, expected_wer, half_precision):
+        work_dir = os.path.join(self.temp_dir, "ctc")
+        # self.create_TLG("ctc", work_dir)
+        wer = self.run_decoder(work_dir, self.dataset_map[dataset], half_precision)
+        assert wer <= expected_wer
 
     def test_compact_ctc_topo(self):
         self.create_TLG("ctc_compact", os.path.join(self.temp_dir, "ctc_compact"))
@@ -152,18 +175,33 @@ def create_TLG(self, topo, work_dir):
             ]
         )
 
-    def run_decoder(self, graph_path: str):
-        asr_model = nemo_asr.models.ASRModel.restore_from(self.nemo_model_path, map_location=torch.device("cuda"))
-
+    @staticmethod
+    def create_decoder_config():
         config = BatchedMappedDecoderCudaConfig()
         config.n_input_per_chunk = 50
         config.online_opts.decoder_opts.default_beam = 17.0
         config.online_opts.decoder_opts.lattice_beam = 8.0
-        config.online_opts.decoder_opts.max_active = 7000
+        config.online_opts.decoder_opts.max_active = 10_000
+        # From when I was testing some penalties discussed in wenet. These can be added back later.
+        # config.online_opts.decoder_opts.blank_penalty = 0.01
+        # config.online_opts.decoder_opts.blank_ilabel = 1024
+        # config.online_opts.decoder_opts.length_penalty = -4.5
         config.online_opts.determinize_lattice = True
-        config.online_opts.max_batch_size = 400
-        config.online_opts.num_channels = 800
-        config.online_opts.frame_shift_seconds = 0.03
+        config.online_opts.max_batch_size = 200
+        config.online_opts.num_channels = config.online_opts.max_batch_size * 2
+        config.online_opts.frame_shift_seconds = 0.04
+        config.online_opts.lattice_postprocessor_opts.acoustic_scale = 10.0
+        config.online_opts.lattice_postprocessor_opts.lm_scale = 6.0
+        config.online_opts.lattice_postprocessor_opts.word_ins_penalty = 0.0
+        config.online_opts.num_decoder_copy_threads = 1
+        config.online_opts.num_post_processing_worker_threads = multiprocessing.cpu_count() - config.online_opts.num_decoder_copy_threads
+
+        return config
+
+    def run_decoder(self, graph_path: str, dataset: torch.utils.data.IterableDataset, half_precision: bool):
+        asr_model = nemo_asr.models.ASRModel.restore_from(self.nemo_model_path, map_location=torch.device("cuda"))
+
+        config = self.create_decoder_config()
         decoder = BatchedMappedDecoderCuda(
             config,
             os.path.join(graph_path, "graph/graph_ctc_3-gram.pruned.3e-7/TLG.fst"),
@@ -171,37 +209,187 @@ def run_decoder(self, graph_path: str):
             self.num_tokens_including_blank
         )
 
-        dir = os.path.join(graph_path, "graph/graph_ctc_3-gram.pruned.3e-7")
+        data_loader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=config.online_opts.num_channels,
+            # num_workers=1,
+            collate_fn=collate_fn,
+            pin_memory=True)
+
+        references = []
+        results = []
+        all_greedy_predictions = []
+
+        total_audio_length_samples = 0
+
+        # wer = torchmetrics.WordErrorRate()
+        asr_model.preprocessor.featurizer.dither = 0.0
+        asr_model.preprocessor.featurizer.pad_to = 0
+        asr_model.eval()
+        asr_model.encoder.freeze()
+        asr_model.decoder.freeze()
+        torch.cuda.cudart().cudaProfilerStart()
 
-        os.makedirs("ctm", exist_ok=True)
+        with ExitStack() as stack:
+            stack.enter_context(torch.inference_mode())
+            if half_precision:
+                stack.enter_context(torch.autocast("cuda"))
+            start_time = time.time_ns()
+            for batch in data_loader:
+                torch.cuda.nvtx.range_push("batch")
+                input_signal, input_signal_length, target = batch
+                total_audio_length_samples += torch.sum(input_signal_length)
+                references.extend(target)
+                torch.cuda.nvtx.range_push("H2D")
+                input_signal = input_signal.cuda()
+                input_signal_length = input_signal_length.cuda()
+                torch.cuda.nvtx.range_pop()
+                torch.cuda.nvtx.range_push("forward")
+                log_probs, lengths, _ = asr_model.forward(input_signal=input_signal, input_signal_length=input_signal_length)
+                torch.cuda.nvtx.range_pop()
+                # from IPython import embed; embed()
+                # greedy_predictions, _ = asr_model.decoding.ctc_decoder_predictions_tensor(log_probs, decoder_lengths=lengths, return_hypotheses=False)
+                # all_greedy_predictions.extend(greedy_predictions)
 
-        import kaldi_io
-        for key, matrix in kaldi_io.read_mat_scp("ark:test_data/logits.ark"):  # "scp:/home/dgalvez/scratch/code/asr/riva-asrlib-decoder/test/logits.scp"):
-            sequences = [torch.from_numpy(matrix.copy())]
-            sequence_lengths = [sequences[0].shape[0]]
-            padded_sequence = torch.nn.utils.rnn.pad_sequence(sequences, batch_first=True).cuda()
-            sequence_lengths_tensor = torch.tensor(sequence_lengths, dtype=torch.long)
-            if config.online_opts.use_lattice_postprocessor:
-                write_ctm_output(key, decoder.decode(padded_sequence, sequence_lengths_tensor)[0])
-            decoder.decode_write_lattice(padded_sequence, sequence_lengths_tensor, [key], "ark:|gzip -c > ctm/lat.1.gz")
+                torch.cuda.nvtx.range_push("D2H lengths")
+                cpu_lengths = lengths.to(torch.int64).to('cpu')
+                torch.cuda.nvtx.range_pop()
+                torch.cuda.nvtx.range_push("beam search decoder")
+                results.extend(decoder.decode_mbr(log_probs.to(torch.float32), cpu_lengths))
+                torch.cuda.nvtx.range_pop()
+                torch.cuda.nvtx.range_pop()
+            end_time = time.time_ns()
+        torch.cuda.cudart().cudaProfilerStop()
+        run_time_seconds = (end_time - start_time) / 1_000_000_000
+        input_time_seconds = total_audio_length_samples / 16_000
+        print("RTFx:", input_time_seconds / run_time_seconds)
+        print("run time:", run_time_seconds)
+        predictions = []
+        for result in results:
+            predictions.append(" ".join(piece[0] for piece in result))
+        references = [s.lower() for s in references]
+        my_wer = wer(references, predictions)
+        print("beam search WER:", my_wer)
+        # print("greedy WER:", wer(references, all_greedy_predictions))
+        return my_wer[0]
 
 def write_ctm_output(key, result):
     for word, start, end in result:
         print(f"{key} 1 {start:.2f} {end - start:.2f} {word} 1.0")
 
-        # manifest = "/mnt/disks/sda_hdd/librispeech/dev_clean.json"
-        # paths = []
-        # with open(manifest) as fh:
-        #     for line in fh:
-        #         entry = json.loads(line)
-        #         paths.append(entry["audio_filepath"])
-
-        # for path in paths:
-        #     logprobs = asr_model.transcribe([path], batch_size=1, logprobs=True)
-        #     sequences = [torch.from_numpy(logprobs[0]).cuda()]
-        #     sequence_lengths = [logprobs[0].shape[0]]
-        #     padded_sequence = torch.nn.utils.rnn.pad_sequence(sequences, batch_first=True)
-        #     sequence_lengths_tensor = torch.tensor(sequence_lengths, dtype=torch.long)
-        #     for result in decoder.decode(padded_sequence,
-        #                                  sequence_lengths_tensor):
-        #         print(result)
+def collate_fn(batch):
+    # A data tuple has the form:
+    # waveform, sample_rate, label, speaker_id, utterance_number
+
+    tensors = []
+    targets = []
+    lengths = []
+
+    # Gather in lists, and encode labels as indices
+    for waveform, _, label, *_ in batch:
+        waveform = waveform.squeeze()
+        tensors += [waveform]
+        # targets += [torch.zeros(1)]
+        targets.append(label)
+        lengths.append(waveform.size(0))
+        # targets += [label_to_index(label)]
+
+    # Group the list of tensors into a batched tensor
+    tensors = torch.nn.utils.rnn.pad_sequence(tensors, batch_first=True)
+    lengths = torch.tensor(lengths, dtype=torch.long)
+    # targets = torch.stack(targets)
+
+    return tensors, lengths, targets
+
+def trace_back_stats(r, h, d):
+    i = len(r)
+    j = len(h)
+    insertions = 0
+    substitutions = 0
+    deletions = 0
+    while True:
+        if i == 0 and j == 0:
+            break
+        elif i >= 1 and j >= 1 and d[i][j] == d[i-1][j-1] and r[i-1] == h[j-1]:
+            i = i - 1
+            j = j - 1
+        elif j >= 1 and d[i][j] == d[i][j-1]+1:
+            insertions += 1
+            i = i
+            j = j - 1
+        elif i >= 1 and j >= 1 and d[i][j] == d[i-1][j-1]+1:
+            substitutions += 1
+            i = i - 1
+            j = j - 1
+        else:
+            deletions += 1
+            i = i - 1
+            j = j
+    return insertions, substitutions, deletions
+
+def wer(references, hypotheses):
+    total_wer = 0
+    total_insertions = 0
+    total_substitutions = 0
+    total_deletions = 0
+    total_words = 0
+    for ref, hyp in zip(references, hypotheses):
+        transcript_wer, i, s, d = wer_single(ref.split(), hyp.split())
+        total_words += len(ref.split())
+        total_wer += transcript_wer
+        total_insertions += i
+        total_substitutions += s
+        total_deletions += d
+    return total_wer / total_words, total_insertions, total_substitutions, total_deletions
+
+def wer_single(r, h):
+    """
+    Calculation of WER with Levenshtein distance.
+
+    Works only for iterables up to 2**16 - 1 elements (uint16).
+    O(nm) time and space complexity.
+
+    Parameters
+    ----------
+    r : list
+    h : list
+
+    Returns
+    -------
+    int
+
+    Examples
+    --------
+    >>> wer("who is there".split(), "is there".split())
+    1
+    >>> wer("who is there".split(), "".split())
+    3
+    >>> wer("".split(), "who is there".split())
+    3
+    """
+    # initialisation
+
+    d = np.zeros((len(r) + 1, len(h) + 1), dtype=np.uint16)
+    for i in range(len(r) + 1):
+        for j in range(len(h) + 1):
+            if i == 0:
+                d[0][j] = j
+            elif j == 0:
+                d[i][0] = i
+
+    # computation
+    for i in range(1, len(r) + 1):
+        for j in range(1, len(h) + 1):
+            if r[i - 1].lower() == h[j - 1].lower():
+                d[i][j] = d[i - 1][j - 1]
+            else:
+                substitution = d[i - 1][j - 1] + 1
+                insertion = d[i][j - 1] + 1
+                deletion = d[i - 1][j] + 1
+                d[i][j] = min(substitution, insertion, deletion)
+
+    insertions, substitions, deletions = trace_back_stats(r, h, d)
+
+    assert insertions + substitions + deletions == d[len(r)][len(h)]
+
+    return d[len(r)][len(h)], insertions, substitions, deletions