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,159 @@ 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 L reset where the SPM lattice drops the explicit L
+    // ("Upp"+"v" -> "PPV" instead of "PPv"). See
+    // https://microsoft.visualstudio.com/Edge/_git/edge.onnxruntime-extensions
+    // commit history for the full repro.
+
+    token.clear();
+    token.reserve(piece.size());
+
+    char mode = (*state)->signature_;
+
+    auto is_letter_codepoint = [this](const std::string& cp_utf8) -> bool {
+      if (cp_utf8.empty()) return false;
+      wchar_t codepoint = 0;
+      size_t char_len = 0;
+      if (!DecodeFirstUTF8Codepoint(cp_utf8, codepoint, char_len)) {
+        return false;
+      }
+      (void)char_len;
+      return std::iswalpha(static_cast<wint_t>(codepoint)) != 0;
+    };
 
-    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;
+    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);
+    };
+
+    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;
+      }
+
+      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;
       }
-    } 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
+      if (ch == normalizer::cPunctuation) {
+        // P is a pass-through marker; preserve the active mode.
+        ++i;
+        continue;
+      }
+
+      // Real codepoint: determine UTF-8 byte length from the lead byte,
+      // then apply / propagate the active case mode.
+      size_t cp_len = 1;
+      if ((ch >> 5) == 0x6) {
+        cp_len = 2;
+      } else if ((ch >> 4) == 0xE) {
+        cp_len = 3;
+      } else if ((ch >> 3) == 0x1E) {
+        cp_len = 4;
+      }
+      if (i + cp_len > n) {
+        cp_len = n - i;  // truncated; emit verbatim
+      }
+      std::string cp = piece.substr(i, cp_len);
+      const bool is_letter = is_letter_codepoint(cp);
+
+      if (is_letter && (mode == normalizer::cTitlecase ||
+                        mode == normalizer::cUppercase ||
+                        mode == normalizer::cAllUppercase)) {
+        uppercase_codepoint(cp);
+        token.append(cp);
+        if (mode == normalizer::cTitlecase) {
+          mode = 0;  // T applies to one codepoint only
+        }
+        // U / A persist
+      } else {
+        token.append(cp);
+        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,126 @@ TEST(OrtxTokenizerTest, MarianTokenizer2) {
                                                 30,  30,   30,    278,  31,    31,    311, 289, 278,  0}));
 }
 
+// ============================================================================
+// Marian Id2Token bug-fix regression tests
+// ============================================================================
+
+// Bug 1: Mode doesn't propagate across pieces.
+// The Marian 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(OrtxTokenizerTest, MarianId2Token_CrossPieceModePropagate) {
+  OrtxObjectPtr<OrtxTokenizer> tokenizer(OrtxCreateTokenizer, "data/tokenizer/nmt");
+  ASSERT_EQ(tokenizer.Code(), kOrtxOK) << "Failed to create tokenizer, stopping the test.";
+
+  // "MCP" triggers cross-piece U-run: the encoder emits "Umc"+"p" (or
+  // similar segmentation) where U mode must survive the piece boundary.
+  const char* input[] = {"MCP protocol"};
+  OrtxObjectPtr<OrtxTokenId2DArray> token_ids;
+  OrtxTokenize(tokenizer.get(), input, 1, token_ids.ToBeAssigned());
+  ASSERT_EQ(token_ids.Code(), kOrtxOK);
+
+  size_t length = 0;
+  const extTokenId_t* ids = nullptr;
+  OrtxTokenId2DArrayGetItem(token_ids.get(), 0, &ids, &length);
+  ASSERT_GT(length, 0u);
+
+  std::vector<extTokenId_t> ids_vec(ids, ids + length);
+  OrtxObjectPtr<OrtxStringArray> decoded_text;
+  OrtxDetokenize1D(tokenizer.get(), ids_vec.data(), ids_vec.size(), decoded_text.ToBeAssigned());
+  ASSERT_EQ(decoded_text.Code(), kOrtxOK);
+
+  const char* text = nullptr;
+  OrtxStringArrayGetItem(decoded_text.get(), 0, &text);
+  EXPECT_STREQ(text, "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(OrtxTokenizerTest, MarianId2Token_MidPieceMarker) {
+  OrtxObjectPtr<OrtxTokenizer> tokenizer(OrtxCreateTokenizer, "data/tokenizer/nmt");
+  ASSERT_EQ(tokenizer.Code(), kOrtxOK) << "Failed to create tokenizer, stopping the test.";
+
+  // "iPhone" triggers mid-piece marker: the encoder lowercases to "iphone"
+  // and inserts a T(itlecase) before the 'p', producing a piece like "iTphone".
+  const char* input[] = {"iPhone is great"};
+  OrtxObjectPtr<OrtxTokenId2DArray> token_ids;
+  OrtxTokenize(tokenizer.get(), input, 1, token_ids.ToBeAssigned());
+  ASSERT_EQ(token_ids.Code(), kOrtxOK);
+
+  size_t length = 0;
+  const extTokenId_t* ids = nullptr;
+  OrtxTokenId2DArrayGetItem(token_ids.get(), 0, &ids, &length);
+  ASSERT_GT(length, 0u);
+
+  std::vector<extTokenId_t> ids_vec(ids, ids + length);
+  OrtxObjectPtr<OrtxStringArray> decoded_text;
+  OrtxDetokenize1D(tokenizer.get(), ids_vec.data(), ids_vec.size(), decoded_text.ToBeAssigned());
+  ASSERT_EQ(decoded_text.Code(), kOrtxOK);
+
+  const char* text = nullptr;
+  OrtxStringArrayGetItem(decoded_text.get(), 0, &text);
+  EXPECT_STREQ(text, "iPhone is great");
+}
+
+// Bug 3: Implicit L reset not recovered.
+// When the SPM unigram lattice drops an explicit L (lowercase) marker at a
+// non-letter codepoint boundary, the decoder must implicitly reset the mode.
+// E.g. "PPV-mp" encodes as "Upp"+"v"+"-"+"mp" and after the hyphen the
+// mode should reset to lowercase, yielding "mp" not "MP".
+TEST(OrtxTokenizerTest, MarianId2Token_ImplicitLReset) {
+  OrtxObjectPtr<OrtxTokenizer> tokenizer(OrtxCreateTokenizer, "data/tokenizer/nmt");
+  ASSERT_EQ(tokenizer.Code(), kOrtxOK) << "Failed to create tokenizer, stopping the test.";
+
+  const char* input[] = {"PPV-mp format"};
+  OrtxObjectPtr<OrtxTokenId2DArray> token_ids;
+  OrtxTokenize(tokenizer.get(), input, 1, token_ids.ToBeAssigned());
+  ASSERT_EQ(token_ids.Code(), kOrtxOK);
+
+  size_t length = 0;
+  const extTokenId_t* ids = nullptr;
+  OrtxTokenId2DArrayGetItem(token_ids.get(), 0, &ids, &length);
+  ASSERT_GT(length, 0u);
+
+  std::vector<extTokenId_t> ids_vec(ids, ids + length);
+  OrtxObjectPtr<OrtxStringArray> decoded_text;
+  OrtxDetokenize1D(tokenizer.get(), ids_vec.data(), ids_vec.size(), decoded_text.ToBeAssigned());
+  ASSERT_EQ(decoded_text.Code(), kOrtxOK);
+
+  const char* text = nullptr;
+  OrtxStringArrayGetItem(decoded_text.get(), 0, &text);
+  EXPECT_STREQ(text, "PPV-mp format");
+}
+
+// Combined test: exercises all three Id2Token bugs in a single sentence.
+TEST(OrtxTokenizerTest, MarianId2Token_CombinedBugs) {
+  OrtxObjectPtr<OrtxTokenizer> tokenizer(OrtxCreateTokenizer, "data/tokenizer/nmt");
+  ASSERT_EQ(tokenizer.Code(), kOrtxOK) << "Failed to create tokenizer, stopping the test.";
+
+  // Combines: cross-piece U-run (THIS), mid-piece marker (iPhone), and
+  // implicit L reset (PPV-mp).
+  const char* input[] = {"THIS iPhone costs PPV-mp only"};
+  OrtxObjectPtr<OrtxTokenId2DArray> token_ids;
+  OrtxTokenize(tokenizer.get(), input, 1, token_ids.ToBeAssigned());
+  ASSERT_EQ(token_ids.Code(), kOrtxOK);
+
+  size_t length = 0;
+  const extTokenId_t* ids = nullptr;
+  OrtxTokenId2DArrayGetItem(token_ids.get(), 0, &ids, &length);
+  ASSERT_GT(length, 0u);
+
+  std::vector<extTokenId_t> ids_vec(ids, ids + length);
+  OrtxObjectPtr<OrtxStringArray> decoded_text;
+  OrtxDetokenize1D(tokenizer.get(), ids_vec.data(), ids_vec.size(), decoded_text.ToBeAssigned());
+  ASSERT_EQ(decoded_text.Code(), kOrtxOK);
+
+  const char* text = nullptr;
+  OrtxStringArrayGetItem(decoded_text.get(), 0, &text);
+  EXPECT_STREQ(text, "THIS iPhone costs PPV-mp only");
+}
+
 // ============================================================================
 // Transformers v5 format tests
 // ============================================================================