Fix Id2Token Bugs

operators/tokenizer/ugm_kernels.hpp

@@ -11,6 +11,7 @@
 #include <string_view>
 #include <vector>
 #include <cfloat>
+#include <cstring>
 #include <functional>
 #include <unordered_map>
 #include <cwctype>
@@ -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<wchar_t>(
+            std::towupper(static_cast<wint_t>(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<unsigned char>(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<wint_t>(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 {};
   }
 

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<OrtxTokenizer>(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<OrtxTokenId2DArray> 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<extTokenId_t> ids_vec(ids, ids + length);
+    OrtxObjectPtr<OrtxStringArray> 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<OrtxTokenizer> tokenizer_;
+};
+
+OrtxObjectPtr<OrtxTokenizer> 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
 // ============================================================================