diff --git a/families/llama/runtime/CMakeLists.txt b/families/llama/runtime/CMakeLists.txt index e0dfbae4e9..09284950c7 100644 --- a/families/llama/runtime/CMakeLists.txt +++ b/families/llama/runtime/CMakeLists.txt @@ -35,7 +35,8 @@ install(TARGETS trtmc_model_llama ) if(TRTMC_BUILD_TESTS) - foreach(test_name IN ITEMS test_llama_chat_template test_llama_native_kv_cache) + foreach(test_name IN ITEMS test_llama_chat_template test_llama_native_kv_cache + test_llama_pretokenizer_variant) add_executable(${test_name} ${PROJECT_SOURCE_DIR}/families/llama/tests/cpp/${test_name}.cpp ) diff --git a/families/llama/runtime/bpe_tokenizer.cpp b/families/llama/runtime/bpe_tokenizer.cpp index 860defc1df..742a4e7d3c 100644 --- a/families/llama/runtime/bpe_tokenizer.cpp +++ b/families/llama/runtime/bpe_tokenizer.cpp @@ -1272,19 +1272,52 @@ class BpeTokenizer final : public ITokenizer { return pretok::Variant::kLlama; } - // Detect variant from the Split inside a Sequence pre_tokenizer. + // Detect variant from the Split step(s) inside a Sequence pre_tokenizer. + // Some checkpoints isolate digit-grouping (\p{N}{1,3}) as its own Split + // step, ahead of the step whose regex actually identifies the variant + // (e.g. Qwen3's `[^\r\n...`). That earlier step's regex never matches + // any known variant signature and falls through to kLlama, so scanning + // stops before reaching the real classifier. Scan every Split step for + // classification instead of only the first; take the digit-group size + // independently, from whichever step's own regex encodes one, since the + // classifying step and the digit-grouping step are not guaranteed to be + // the same step. static pretok::Variant detect_split_variant(const nlohmann::json& pt, int& digit_group_out) { digit_group_out = 0; if (!pt.contains("pretokenizers")) return pretok::Variant::kLlama; + + pretok::Variant variant = pretok::Variant::kLlama; + bool variant_found = false; for (auto& sub : pt["pretokenizers"]) { if (sub.value("type", "") != "Split") continue; if (!sub.contains("pattern") || !sub["pattern"].contains("Regex")) continue; - return classify_split(sub, digit_group_out); + const auto regex = sub["pattern"]["Regex"].get(); + + if (digit_group_out == 0) { + if (int group = parse_digit_group(regex); group > 0) + digit_group_out = group; + } + + int step_digit_group = 0; // discarded; digit_group_out is tracked separately above + pretok::Variant step_variant = classify_split(sub, step_digit_group); + if (step_variant == pretok::Variant::kLlama) + continue; + if (variant_found && step_variant != variant) { + throw std::runtime_error( + "Ambiguous Sequence pre_tokenizer: Split steps classify to different " + "variants (" + + std::to_string(static_cast(variant)) + " and " + + std::to_string(static_cast(step_variant)) + ")"); + } + if (!variant_found) { + variant = step_variant; + variant_found = true; + } } - return pretok::Variant::kLlama; + return variant; } static bool is_space_split_pre_tokenizer(const nlohmann::json& pt, const std::string& pt_type) { diff --git a/families/llama/tests/cpp/test_llama_pretokenizer_variant.cpp b/families/llama/tests/cpp/test_llama_pretokenizer_variant.cpp new file mode 100644 index 0000000000..f4e89dd0d0 --- /dev/null +++ b/families/llama/tests/cpp/test_llama_pretokenizer_variant.cpp @@ -0,0 +1,191 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// detect_split_variant() previously returned on the first Split step found +// in a Sequence pre_tokenizer, even when that step's regex classifies as +// the generic kLlama default. Checkpoints that isolate digit-grouping +// (\p{N}{1,3}) as its own Split step ahead of the step that actually +// identifies the variant (e.g. Qwen3's `[^\r\n...` signature) were +// misclassified as kLlama, silently disabling variant-specific +// pre-tokenization behavior (trailing-newline attachment, grouped digit +// runs) for the whole checkpoint. +// +// The first two checks below use the pre_tokenizer field exactly as +// published in openbmb/MiniCPM5-2B's tokenizer.json (a real, public +// checkpoint that triggers this misclassification), so a passing result +// here reflects the real checkpoint's own regex shapes, not a synthetic +// stand-in. The third check is a synthetic fixture verifying the +// ambiguous-classification case (two Split steps that each classify to a +// different, non-kLlama variant) throws rather than silently picking one. + +#include "families/llama/runtime/tokenizer.h" + +#include +#include +#include +#include +#include + +namespace { + +int failures = 0; + +void check_ids(const std::vector& actual, const std::vector& expected, + const char* name) { + if (actual == expected) + return; + std::cerr << "FAIL: " << name << " expected"; + for (const int32_t token : expected) + std::cerr << ' ' << token; + std::cerr << " got"; + for (const int32_t token : actual) + std::cerr << ' ' << token; + std::cerr << '\n'; + ++failures; +} + +// MiniCPM5-2B's actual pre_tokenizer field (openbmb/MiniCPM5-2B, +// tokenizer.json): a digit-grouping Split step first, then the real +// classifying (Qwen3-shaped) Split step second, then ByteLevel. +constexpr const char* kMiniCPM5PreTokenizer = R"( + "pre_tokenizer": { + "type": "Sequence", + "pretokenizers": [ + { + "type": "Split", + "pattern": {"Regex": "\\p{N}{1,3}"}, + "behavior": "Isolated", + "invert": false + }, + { + "type": "Split", + "pattern": { + "Regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}+| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" + }, + "behavior": "Isolated", + "invert": false + }, + { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": false + } + ] + })"; + +// Trailing newlines after punctuation should stay attached to the same +// pre-token (Qwen3-variant behavior) rather than split off as their own +// whitespace token. Only reachable if the second Split step's regex is +// actually used to classify the variant, not the first. +void check_newline_attachment() { + const std::string tokenizer_json = std::string(R"({ + "model": { + "type": "BPE", + "vocab": { + ".": 0, + "Ċ": 1, + "ĊĊ": 2, + ".ĊĊ": 3 + }, + "merges": ["Ċ Ċ", ". ĊĊ"] + },)") + kMiniCPM5PreTokenizer + R"( + })"; + + auto tokenizer = trtmc::CreateBpeTokenizer(tokenizer_json.data(), tokenizer_json.size(), false); + if (!tokenizer) { + std::cerr << "FAIL: check_newline_attachment: tokenizer was not created\n"; + ++failures; + return; + } + check_ids(tokenizer->encode(".\n\n"), {3}, "check_newline_attachment"); +} + +// Multi-digit runs should be grouped into chunks of up to 3 digits +// (\p{N}{1,3}), each its own pre-token, not left as one unbounded digit +// run (the kLlama-misclassified behavior) and not split into individual +// one-digit pre-tokens (the failure mode of a fix that scans for the +// variant but couples digit-group size to the same winning step, which +// does not carry the {1,3} grouping in this checkpoint's real regex). +void check_digit_grouping() { + const std::string tokenizer_json = std::string(R"({ + "model": { + "type": "BPE", + "vocab": { + "1": 0, + "5": 1, + "0": 2, + "15": 3, + "150": 4, + "1500": 5 + }, + "merges": ["1 5", "15 0", "150 0"] + },)") + kMiniCPM5PreTokenizer + R"( + })"; + + auto tokenizer = trtmc::CreateBpeTokenizer(tokenizer_json.data(), tokenizer_json.size(), false); + if (!tokenizer) { + std::cerr << "FAIL: check_digit_grouping: tokenizer was not created\n"; + ++failures; + return; + } + // Correct (grouped by 3): ["150", "0"] -> ids {4, 2}. + // kLlama misclassification (unpatched bug): one ungrouped word "1500" -> {5}. + // Coupled/naive fix (digit_group taken from the same step as the + // variant, which for this checkpoint's real second step is 0): four + // one-digit words -> {0, 1, 2, 2}. + check_ids(tokenizer->encode("1500"), {4, 2}, "check_digit_grouping"); +} + +// Synthetic fixture: two Split steps that each classify to a different, +// non-kLlama variant. Not a real checkpoint shape -- this exercises the +// ambiguous-classification guard rather than reproducing a reported bug. +void check_ambiguous_variants_throw() { + const std::string tokenizer_json = R"({ + "model": { + "type": "BPE", + "vocab": {"a": 0}, + "merges": [] + }, + "pre_tokenizer": { + "type": "Sequence", + "pretokenizers": [ + { + "type": "Split", + "pattern": {"Regex": "[^(\\s)]+"}, + "behavior": "Isolated", + "invert": false + }, + { + "type": "Split", + "pattern": {"Regex": "[^\\r\\n\\p{L}\\p{N}]?\\p{L}+"}, + "behavior": "Isolated", + "invert": false + } + ] + } + })"; + + bool threw = false; + try { + trtmc::CreateBpeTokenizer(tokenizer_json.data(), tokenizer_json.size(), false); + } catch (const std::runtime_error&) { + threw = true; + } + if (!threw) { + std::cerr + << "FAIL: check_ambiguous_variants_throw: expected std::runtime_error, none thrown\n"; + ++failures; + } +} + +} // namespace + +int main() { + check_newline_attachment(); + check_digit_grouping(); + check_ambiguous_variants_throw(); + return failures == 0 ? 0 : 1; +}