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; +}