diff --git a/operators/tokenizer/ugm_kernels.hpp b/operators/tokenizer/ugm_kernels.hpp index bd5f37320..990dfa9a8 100644 --- a/operators/tokenizer/ugm_kernels.hpp +++ b/operators/tokenizer/ugm_kernels.hpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -896,73 +897,149 @@ class SpmUgmDecoder { return {}; } - token = vocab_[id]; - if (case_encoding_ && token.length() == 1) { - if (token[0] == normalizer::cUppercase || token[0] == normalizer::cAllUppercase || - token[0] == normalizer::cTitlecase || token[0] == normalizer::cLowercase || - token[0] == normalizer::cPunctuation) { - (*state)->signature_ = token[0]; - token = ""; - return {}; - } - } + const std::string& piece = vocab_[id]; - const std::string ws = " "; - auto pos = token.find(spm_escaped_space); - if (pos != std::string::npos) { + if (!case_encoding_) { + // Non-Marian unigram path: just rewrite the SPM space marker and + // emit the piece verbatim. + token = piece; + auto pos = token.find(spm_escaped_space); if (pos == 0) { - token = ws + token.substr(spm_escaped_space.length()); - } else if (pos + 3 == token.length()) { - token = token.substr(0, pos) + ws; + token = std::string(" ") + token.substr(spm_escaped_space.length()); + } else if (pos != std::string::npos && + pos + spm_escaped_space.length() == token.length()) { + token = token.substr(0, pos) + std::string(" "); } - } - - if (!case_encoding_) { return {}; } - char signature = 0; - if ((*state)->signature_ != 0) { - signature = (*state)->signature_; - (*state)->signature_ = 0; - } + // Marian case-encoder protocol -- per-piece byte-level state machine. + // + // The Marian case-encoder pre-pass lowercases everything before applying + // markers, so any uppercase ASCII letter we encounter inside a unigram + // piece is by construction a marker (cUppercase 'U', cAllUppercase 'A', + // cTitlecase 'T', cLowercase 'L', cPunctuation 'P'). + // + // The previous implementation only recognized a marker at position 0 of + // a piece, which broke three real-world cases that ship in the trained + // vocab: (1) cross-piece U-runs ("Umc"+"p" -> "MCp" instead of "MCP"), + // (2) mid-piece markers ("iTphone" -> "iTphone" instead of "iPhone"), + // and (3) implicit mode reset after a non-letter boundary, where + // uppercase/titlecase state must not leak past punctuation or into the + // following lowercase run (e.g. "PPV-mp" decoded as "PPV-MP" instead + // of "PPV-mp"). + + token.clear(); + token.reserve(piece.size()); + + char mode = (*state)->signature_; + + auto uppercase_codepoint = [this](std::string& cp_utf8) { + if (cp_utf8.empty()) return; + wchar_t codepoint = 0; + size_t char_len = 0; + if (!DecodeFirstUTF8Codepoint(cp_utf8, codepoint, char_len)) return; + (void)char_len; + // Match TitlecaseFirstCharacter's special-cases. + if (codepoint >= L'\u0430' && codepoint <= L'\u044f') { + codepoint = codepoint - (L'\u0430' - L'\u0410'); + } else if (codepoint == L'\u0451') { + codepoint = L'\u0401'; + } else { + codepoint = static_cast( + std::towupper(static_cast(codepoint))); + } + cp_utf8 = EncodeUTF8(codepoint); + }; - if (signature) { - // Apply transformation from previous token's signature - switch (signature) { - case normalizer::cUppercase: - case normalizer::cAllUppercase: - std::transform(token.begin(), token.end(), token.begin(), ::toupper); - break; - case normalizer::cTitlecase: - TitlecaseFirstCharacter(token); - break; - case normalizer::cLowercase: - case normalizer::cPunctuation: - // No transformation needed - break; + size_t i = 0; + const size_t n = piece.size(); + while (i < n) { + // SPM space marker (\u2581 = U+2581, 3 UTF-8 bytes). + if (i + spm_escaped_space.size() <= n && + std::memcmp(piece.data() + i, spm_escaped_space.data(), + spm_escaped_space.size()) == 0) { + token.push_back(' '); + // U/T modes do not survive a word boundary. cAllUppercase by + // design crosses spaces (matches case_encoder.cc::PostProcess + // A-spans). + if (mode != normalizer::cAllUppercase) { + mode = 0; + } + i += spm_escaped_space.size(); + continue; } - } else if (!token.empty()) { - // Check if current token starts with a signature character - char first_char = token[0]; - if (first_char == normalizer::cUppercase || first_char == normalizer::cAllUppercase || - first_char == normalizer::cTitlecase || first_char == normalizer::cLowercase || - first_char == normalizer::cPunctuation) { - token.erase(0, 1); // Remove signature character - - switch (first_char) { - case normalizer::cUppercase: - case normalizer::cAllUppercase: - std::transform(token.begin(), token.end(), token.begin(), ::toupper); - break; - case normalizer::cTitlecase: - TitlecaseFirstCharacter(token); - break; - // For cLowercase and cPunctuation, no transformation needed + + const unsigned char ch = static_cast(piece[i]); + + // Single-byte ASCII marker dispatch. + if (ch == normalizer::cUppercase) { + mode = normalizer::cUppercase; + ++i; + continue; + } + if (ch == normalizer::cAllUppercase) { + mode = normalizer::cAllUppercase; + ++i; + continue; + } + if (ch == normalizer::cTitlecase) { + mode = normalizer::cTitlecase; + ++i; + continue; + } + if (ch == normalizer::cLowercase) { + mode = 0; + ++i; + continue; + } + if (ch == normalizer::cPunctuation) { + // P is a pass-through marker; preserve the active mode. + ++i; + continue; + } + + // Real codepoint: use DecodeFirstUTF8Codepoint for robust byte- + // length detection (handles malformed / truncated sequences). + wchar_t codepoint = 0; + size_t cp_len = 0; + std::string remaining = piece.substr(i); + if (!DecodeFirstUTF8Codepoint(remaining, codepoint, cp_len) || cp_len == 0) { + // Malformed UTF-8: emit a single byte verbatim and advance. + token.push_back(piece[i]); + ++i; + continue; + } + + const bool is_letter = std::iswalpha(static_cast(codepoint)) != 0; + + if (is_letter && (mode == normalizer::cTitlecase || + mode == normalizer::cUppercase || + mode == normalizer::cAllUppercase)) { + // Uppercase transform needed -- allocate only for this path. + std::string cp = piece.substr(i, cp_len); + uppercase_codepoint(cp); + token.append(cp); + if (mode == normalizer::cTitlecase) { + mode = 0; // T applies to one codepoint only + } + // U / A persist + } else { + // No transform: append directly from piece without allocating. + token.append(piece, i, cp_len); + if (!is_letter && (mode == normalizer::cUppercase || + mode == normalizer::cTitlecase)) { + // Implicit L: the encoder would have terminated the U/T run at + // a non-letter codepoint, but the SPM unigram lattice may have + // merged the explicit L away in scoring. Recover here. + mode = 0; } } + + i += cp_len; } + (*state)->signature_ = mode; return {}; } diff --git a/test/pp_api_test/test_tokenizer_capi.cc b/test/pp_api_test/test_tokenizer_capi.cc index 014daf67a..d54f6767d 100644 --- a/test/pp_api_test/test_tokenizer_capi.cc +++ b/test/pp_api_test/test_tokenizer_capi.cc @@ -298,6 +298,81 @@ TEST(OrtxTokenizerTest, MarianTokenizer2) { 30, 30, 30, 278, 31, 31, 311, 289, 278, 0})); } +// ============================================================================ +// Marian Id2Token bug-fix regression tests +// ============================================================================ + +// Fixture: shares a single NMT tokenizer instance and provides a helper +// that tokenizes + detokenizes a string, returning the round-tripped text. +class MarianId2TokenTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + tokenizer_ = OrtxObjectPtr(OrtxCreateTokenizer, "data/tokenizer/nmt"); + } + static void TearDownTestSuite() { tokenizer_.reset(); } + + // Tokenize |input|, detokenize, and return the result. + static std::string RoundTrip(const char* input) { + const char* inputs[] = {input}; + OrtxObjectPtr token_ids; + OrtxTokenize(tokenizer_.get(), inputs, 1, token_ids.ToBeAssigned()); + EXPECT_EQ(token_ids.Code(), kOrtxOK); + + size_t length = 0; + const extTokenId_t* ids = nullptr; + OrtxTokenId2DArrayGetItem(token_ids.get(), 0, &ids, &length); + EXPECT_GT(length, 0u); + + std::vector ids_vec(ids, ids + length); + OrtxObjectPtr decoded; + OrtxDetokenize1D(tokenizer_.get(), ids_vec.data(), ids_vec.size(), + decoded.ToBeAssigned()); + EXPECT_EQ(decoded.Code(), kOrtxOK); + + const char* text = nullptr; + OrtxStringArrayGetItem(decoded.get(), 0, &text); + return text ? std::string(text) : std::string(); + } + + static OrtxObjectPtr tokenizer_; +}; + +OrtxObjectPtr MarianId2TokenTest::tokenizer_; + +// Bug 1: Mode doesn't propagate across pieces. +// The case-encoder U (uppercase) mode must persist across SPM piece +// boundaries. E.g. "MCP" encodes as pieces like "Umc"+"p", and the U mode +// from the first piece must carry into the second so "p" becomes "P". +TEST_F(MarianId2TokenTest, CrossPieceModePropagate) { + ASSERT_EQ(tokenizer_.Code(), kOrtxOK) << "Failed to create tokenizer."; + EXPECT_EQ(RoundTrip("MCP protocol"), "MCP protocol"); +} + +// Bug 2: Markers mid-piece are ignored. +// When the SPM unigram lattice merges a case marker into the middle of a +// piece (e.g. "iTphone" where T is a titlecase marker), the old decoder +// only checked position 0 and emitted the marker literally. +TEST_F(MarianId2TokenTest, MidPieceMarker) { + ASSERT_EQ(tokenizer_.Code(), kOrtxOK) << "Failed to create tokenizer."; + EXPECT_EQ(RoundTrip("iPhone is great"), "iPhone is great"); +} + +// Bug 3: Implicit mode reset after a non-letter boundary. +// When the SPM lattice drops an explicit L (lowercase) marker at a non-letter +// codepoint boundary (e.g. "-"), the decoder must implicitly reset the mode +// so the following lowercase run is not uppercased. +TEST_F(MarianId2TokenTest, ImplicitLReset) { + ASSERT_EQ(tokenizer_.Code(), kOrtxOK) << "Failed to create tokenizer."; + EXPECT_EQ(RoundTrip("PPV-mp format"), "PPV-mp format"); +} + +// Combined test: exercises all three Id2Token bugs in a single sentence. +TEST_F(MarianId2TokenTest, CombinedBugs) { + ASSERT_EQ(tokenizer_.Code(), kOrtxOK) << "Failed to create tokenizer."; + EXPECT_EQ(RoundTrip("THIS iPhone costs PPV-mp only"), + "THIS iPhone costs PPV-mp only"); +} + // ============================================================================ // Transformers v5 format tests // ============================================================================