Decode Citrinet CTC output through the SentencePiece model

CMakeLists.txt

@@ -694,6 +694,16 @@ if (ENGINE_BUILD_TESTS)
         COMMAND sentencepiece_tokenizer1_test
     )
 
+    add_engine_unittest(citrinet_tokenizer_decode_test tests/unittests/test_citrinet_tokenizer_decode.cpp)
+    target_compile_definitions(citrinet_tokenizer_decode_test PRIVATE
+        ENGINE_TEST_ASSET_ROOT="${ENGINE_UNITTEST_ASSET_ROOT}"
+    )
+
+    add_test(
+        NAME citrinet_tokenizer_decode_test
+        COMMAND citrinet_tokenizer_decode_test
+    )
+
     add_engine_unittest(audio_dsp_test tests/unittests/test_audio_dsp.cpp)
 
     add_test(

include/engine/models/citrinet_asr/assets.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "engine/framework/assets/resource_bundle.h"
+#include "engine/framework/tokenizers/sentencepiece.h"
 
 #include <cstdint>
 #include <filesystem>
@@ -97,7 +98,7 @@ struct CitrinetWeights {
     std::vector<float> fb;
     std::vector<JasperBlockWeights> blocks;
     Conv1dWeights decoder;
-    std::vector<std::string> vocab;
+    std::vector<tokenizers::SentencePiecePiece> tokenizer_pieces;
 };
 
 std::shared_ptr<const CitrinetWeights> load_citrinet_weights_cached(const std::filesystem::path & model_path);

model_specs/citrinet_asr.json

@@ -9,8 +9,7 @@
       },
       "files": {
         "config": "model:citrinet_256_config.json",
-        "tokenizer": "model:citrinet_256_tokenizer.model",
-        "vocab": "model:citrinet_256_vocab.txt"
+        "tokenizer": "model:citrinet_256_tokenizer.model"
       },
       "tensors": {
         "weights": "weights:"
@@ -23,8 +22,7 @@
       },
       "files": {
         "config": "model:citrinet_256_config.json",
-        "tokenizer": "model:citrinet_256_tokenizer.model",
-        "vocab": "model:citrinet_256_vocab.txt"
+        "tokenizer": "model:citrinet_256_tokenizer.model"
       },
       "tensors": {
         "weights": "model:citrinet_256.safetensors"

src/models/citrinet_asr/assets.cpp

@@ -61,31 +61,6 @@ std::vector<JasperBlockConfig> parse_jasper_config(const io::json::Value & value
     return blocks;
 }
 
-std::vector<std::string> load_vocab_file(const std::filesystem::path & path) {
-    const auto text = io::read_text_file(path);
-    std::vector<std::string> vocab;
-    size_t start = 0;
-    while (true) {
-        const size_t end = text.find('\n', start);
-        std::string piece = text.substr(start, end == std::string::npos ? std::string::npos : end - start);
-        if (!piece.empty() && piece.back() == '\r') {
-            piece.pop_back();
-        }
-        vocab.push_back(std::move(piece));
-        if (end == std::string::npos) {
-            break;
-        }
-        start = end + 1;
-        if (start > text.size()) {
-            break;
-        }
-    }
-    if (vocab.empty()) {
-        throw std::runtime_error("empty vocab file: " + path.string());
-    }
-    return vocab;
-}
-
 CitrinetConfig parse_config(const io::json::Value & root) {
     CitrinetConfig cfg;
     cfg.sample_rate = root.require("sample_rate").as_i64();
@@ -144,9 +119,13 @@ CitrinetWeights load_citrinet_weights(engine::assets::ResourceBundle resources)
         "preprocessor.featurizer.fb",
         {1, weights.config.n_mels, weights.config.n_fft / 2 + 1});
 
-    weights.vocab = load_vocab_file(resources.require_file("vocab"));
-    if (static_cast<int64_t>(weights.vocab.size()) != weights.config.vocab_size) {
-        throw std::runtime_error("vocab size mismatch");
+    const auto tokenizer_model_path = resources.require_file("tokenizer");
+    weights.tokenizer_pieces = tokenizers::load_sentencepiece_model(tokenizer_model_path);
+    if (static_cast<int64_t>(weights.tokenizer_pieces.size()) != weights.config.vocab_size) {
+        throw std::runtime_error(
+            "Citrinet tokenizer model has " + std::to_string(weights.tokenizer_pieces.size()) +
+            " pieces but the config expects vocab_size " + std::to_string(weights.config.vocab_size) +
+            ": " + tokenizer_model_path.string());
     }
 
     weights.blocks.resize(weights.config.jasper.size());

src/models/citrinet_asr/runtime.cpp

@@ -494,7 +494,6 @@ std::vector<float> to_time_major_features(const engine::audio::AudioTensor & fea
 
 FeaturePack compute_citrinet_features(const std::vector<float> & waveform, const CitrinetWeights & weights);
 std::vector<int32_t> greedy_ctc_ids(const CitrinetInferenceResult & result, int32_t blank_id);
-std::string decode_wordpieces(const std::vector<std::string> & vocab, const std::vector<int32_t> & ids);
 
 FeaturePack extract_feature_pack_from_audio(
     const runtime::AudioBuffer & audio,
@@ -513,7 +512,7 @@ CitrinetTranscriptionResult make_transcription_result(
     CitrinetInferenceResult inference) {
     auto ids = greedy_ctc_ids(inference, static_cast<int32_t>(weights.config.blank_id));
     CitrinetTranscriptionResult result;
-    result.text = decode_wordpieces(weights.vocab, ids);
+    result.text = tokenizers::decode_sentencepiece(weights.tokenizer_pieces, ids);
     result.token_ids = std::move(ids);
     result.inference = std::move(inference);
     return result;
@@ -635,37 +634,6 @@ std::vector<int32_t> greedy_ctc_ids(const CitrinetInferenceResult & result, int3
     return ids;
 }
 
-bool is_join_punctuation(const std::string & piece) {
-    return piece == "." || piece == "," || piece == "!" || piece == "?" || piece == ":" || piece == ";" ||
-           piece == "'" || piece == "\"" || piece == ")" || piece == "]" || piece == "}" || piece == "-" ||
-           piece == "/" || piece == "\\";
-}
-
-std::string decode_wordpieces(const std::vector<std::string> & vocab, const std::vector<int32_t> & ids) {
-    std::string text;
-    for (int32_t id : ids) {
-        if (id < 0 || id >= static_cast<int32_t>(vocab.size())) {
-            throw std::runtime_error("token id out of vocab range");
-        }
-        const std::string & piece = vocab[static_cast<size_t>(id)];
-        if (piece.rfind("##", 0) == 0) {
-            text += piece.substr(2);
-            continue;
-        }
-        if (text.empty()) {
-            text += piece;
-            continue;
-        }
-        if (!text.empty() && (text.back() == '\'' || text.back() == '-' || is_join_punctuation(piece))) {
-            text += piece;
-            continue;
-        }
-        text += ' ';
-        text += piece;
-    }
-    return text;
-}
-
 }  // namespace
 
 CitrinetInferenceResult infer_runtime_audio(

tests/unittests/test_asr_standalone_gguf.cpp

@@ -95,6 +95,9 @@ void test_citrinet_standalone_gguf() {
     write_dummy_weights(native / "citrinet_256.safetensors");
     write_text(native / "citrinet_256_config.json", R"json({"vocab_file":"citrinet_256_vocab.txt"})json");
     write_text(native / "citrinet_256_tokenizer.model", "test tokenizer sidecar");
+    // Legacy sidecar: the spec no longer references vocab.txt (CTC ids decode
+    // through tokenizer.model), but previously converted model directories
+    // still contain it — loading must ignore the stray file, not choke on it.
     write_text(native / "citrinet_256_vocab.txt", "a\nb\n");
 
     const auto gguf = packed / "renamed-citrinet.gguf";
@@ -113,10 +116,6 @@ void test_citrinet_standalone_gguf() {
         assets.require_file("config").filename().string(),
         std::string("citrinet_256_config.json"),
         "Citrinet embedded config");
-    engine::test::require_eq(
-        assets.require_file("vocab").filename().string(),
-        std::string("citrinet_256_vocab.txt"),
-        "Citrinet embedded vocabulary");
     engine::test::require_eq(
         assets.require_file("tokenizer").filename().string(),
         std::string("citrinet_256_tokenizer.model"),

tests/unittests/test_citrinet_tokenizer_decode.cpp

@@ -0,0 +1,76 @@
+#include "engine/framework/tokenizers/sentencepiece.h"
+
+#include <iostream>
+#include <filesystem>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace {
+
+std::filesystem::path test_asset_path(const std::string & relative) {
+    return std::filesystem::path(ENGINE_TEST_ASSET_ROOT) / relative;
+}
+
+void require(bool condition, const std::string & message) {
+    if (!condition) {
+        throw std::runtime_error(message);
+    }
+}
+
+// Citrinet CTC decoding maps token ids straight into the tokenizer's piece
+// table, so the table must be exactly id-aligned with the SentencePiece
+// model. A table that drops or reorders entries (NeMo's sidecar vocab.txt
+// omits <unk> at id 0) yields plausible-looking but wrong words instead of
+// an error, which is why the runtime decodes through tokenizer.model.
+void test_ctc_id_decode_round_trip() {
+    const auto pieces = engine::tokenizers::load_sentencepiece_model(
+        test_asset_path("tokenizers/tokenizer-1.model"));
+
+    // The id-alignment contract the vocab.txt sidecar violated: the piece
+    // table must include the model's <unk> entry at id 0, not start at the
+    // first normal piece.
+    require(!pieces.empty(), "tokenizer model contained no pieces");
+    require(pieces[0].text == "<unk>", "expected piece 0 to be <unk>");
+    require(
+        pieces[0].type == engine::tokenizers::SentencePieceType::Unknown,
+        "expected piece 0 to have type Unknown");
+
+    const std::string phrase = "hello world how are you";
+    const auto ids = engine::tokenizers::tokenize_sentencepiece(pieces, phrase);
+    require(!ids.empty(), "tokenizer produced no ids for the test phrase");
+
+    const auto decoded = engine::tokenizers::decode_sentencepiece(pieces, ids);
+    require(
+        decoded == phrase,
+        "decode_sentencepiece round trip failed: got '" + decoded + "'");
+
+    // Shifting every id by one simulates a piece table that lost its first
+    // entry, the failure mode of NeMo's sidecar vocab.txt. The result must
+    // be different text, not a silent reproduction of the phrase.
+    std::vector<int32_t> shifted;
+    shifted.reserve(ids.size());
+    for (const int32_t id : ids) {
+        if (static_cast<size_t>(id) + 1 < pieces.size()) {
+            shifted.push_back(id + 1);
+        }
+    }
+    require(shifted.size() == ids.size(), "shifted id set lost entries");
+    const auto misaligned = engine::tokenizers::decode_sentencepiece(pieces, shifted);
+    require(
+        misaligned != phrase,
+        "shifted ids unexpectedly decoded to the original phrase");
+}
+
+}  // namespace
+
+int main() {
+    try {
+        test_ctc_id_decode_round_trip();
+    } catch (const std::exception & ex) {
+        std::cerr << "test_citrinet_tokenizer_decode failed: " << ex.what() << "\n";
+        return 1;
+    }
+    std::cout << "test_citrinet_tokenizer_decode passed\n";
+    return 0;
+}