From 727ef7bb485ef8d0eab8d77492593b277c06871e Mon Sep 17 00:00:00 2001 From: Tym Rabchuk Date: Tue, 14 Jul 2026 15:21:24 +0000 Subject: [PATCH] Decode Citrinet CTC output through the SentencePiece model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Citrinet runtime resolved CTC token ids against the sidecar vocab.txt copied verbatim from the NeMo archive. That file is NeMo's WordPiece-style compat vocab: it omits (id 0), so every line sits one position off from the true SentencePiece ids, and the runtime decoded every token to the wrong piece — transcription came out as plausible-looking subword salad on every backend. The misalignment passed the vocab_size check because the line loader appended a phantom empty entry for the trailing newline. Load the piece table from tokenizer.model (already exported by model_manager and named in the model package spec) via the framework's SentencePiece support and decode with it, replacing the WordPiece join heuristics and the spec's vocab entries. Existing converted model directories work unchanged. --- CMakeLists.txt | 10 +++ include/engine/models/citrinet_asr/assets.h | 3 +- model_specs/citrinet_asr.json | 6 +- src/models/citrinet_asr/assets.cpp | 35 ++------- src/models/citrinet_asr/runtime.cpp | 34 +-------- tests/unittests/test_asr_standalone_gguf.cpp | 7 +- .../test_citrinet_tokenizer_decode.cpp | 76 +++++++++++++++++++ 7 files changed, 101 insertions(+), 70 deletions(-) create mode 100644 tests/unittests/test_citrinet_tokenizer_decode.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index c7e89cdb2..21da96050 100644 --- a/CMakeLists.txt +++ b/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( diff --git a/include/engine/models/citrinet_asr/assets.h b/include/engine/models/citrinet_asr/assets.h index a2250fdd5..74bed8ea3 100644 --- a/include/engine/models/citrinet_asr/assets.h +++ b/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 #include @@ -97,7 +98,7 @@ struct CitrinetWeights { std::vector fb; std::vector blocks; Conv1dWeights decoder; - std::vector vocab; + std::vector tokenizer_pieces; }; std::shared_ptr load_citrinet_weights_cached(const std::filesystem::path & model_path); diff --git a/model_specs/citrinet_asr.json b/model_specs/citrinet_asr.json index 8ae132ac2..8400fc7f9 100644 --- a/model_specs/citrinet_asr.json +++ b/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" diff --git a/src/models/citrinet_asr/assets.cpp b/src/models/citrinet_asr/assets.cpp index 50c8006c0..b41b5eb16 100644 --- a/src/models/citrinet_asr/assets.cpp +++ b/src/models/citrinet_asr/assets.cpp @@ -61,31 +61,6 @@ std::vector parse_jasper_config(const io::json::Value & value return blocks; } -std::vector load_vocab_file(const std::filesystem::path & path) { - const auto text = io::read_text_file(path); - std::vector 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(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(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()); diff --git a/src/models/citrinet_asr/runtime.cpp b/src/models/citrinet_asr/runtime.cpp index 18dce40e4..39eca33db 100644 --- a/src/models/citrinet_asr/runtime.cpp +++ b/src/models/citrinet_asr/runtime.cpp @@ -494,7 +494,6 @@ std::vector to_time_major_features(const engine::audio::AudioTensor & fea FeaturePack compute_citrinet_features(const std::vector & waveform, const CitrinetWeights & weights); std::vector greedy_ctc_ids(const CitrinetInferenceResult & result, int32_t blank_id); -std::string decode_wordpieces(const std::vector & vocab, const std::vector & 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(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 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 & vocab, const std::vector & ids) { - std::string text; - for (int32_t id : ids) { - if (id < 0 || id >= static_cast(vocab.size())) { - throw std::runtime_error("token id out of vocab range"); - } - const std::string & piece = vocab[static_cast(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( diff --git a/tests/unittests/test_asr_standalone_gguf.cpp b/tests/unittests/test_asr_standalone_gguf.cpp index 722738fca..5de8c8ad4 100644 --- a/tests/unittests/test_asr_standalone_gguf.cpp +++ b/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"), diff --git a/tests/unittests/test_citrinet_tokenizer_decode.cpp b/tests/unittests/test_citrinet_tokenizer_decode.cpp new file mode 100644 index 000000000..8dfa1896c --- /dev/null +++ b/tests/unittests/test_citrinet_tokenizer_decode.cpp @@ -0,0 +1,76 @@ +#include "engine/framework/tokenizers/sentencepiece.h" + +#include +#include +#include +#include +#include + +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 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 entry at id 0, not start at the + // first normal piece. + require(!pieces.empty(), "tokenizer model contained no pieces"); + require(pieces[0].text == "", "expected piece 0 to be "); + 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 shifted; + shifted.reserve(ids.size()); + for (const int32_t id : ids) { + if (static_cast(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; +}