Emit U+FFFD for malformed byte-fallback runs in BPE Decode

tokenizer/bpe_tokenizer.cc

@@ -332,15 +332,48 @@ struct QueuedMerge {
   }
 };
 
-// Decodes a run of byte-fallback bytes as UTF-8, appending to `out` (invalid
-// sequences become U+FFFD, matching HuggingFace ByteFallback).
+// Returns the length of the well-formed UTF-8 sequence starting at `b[i]`, or 0
+// if there is none (invalid lead or continuation byte, truncated, overlong,
+// surrogate or above U+10FFFF).
+size_t ValidUtf8Len(const std::string& b, size_t i) {
+  const unsigned char c = static_cast<unsigned char>(b[i]);
+  if (c < 0x80) return 1;
+  size_t len;
+  uint32_t min_cp;
+  if (c >= 0xC2 && c <= 0xDF) {
+    len = 2;
+    min_cp = 0x80;
+  } else if (c >= 0xE0 && c <= 0xEF) {
+    len = 3;
+    min_cp = 0x800;
+  } else if (c >= 0xF0 && c <= 0xF4) {
+    len = 4;
+    min_cp = 0x10000;
+  } else {
+    return 0;
+  }
+  if (i + len > b.size()) return 0;
+  uint32_t cp = c & (0x7Fu >> len);
+  for (size_t k = 1; k < len; ++k) {
+    const unsigned char cc = static_cast<unsigned char>(b[i + k]);
+    if ((cc & 0xC0) != 0x80) return 0;
+    cp = (cp << 6) | (cc & 0x3Fu);
+  }
+  if (cp < min_cp || cp > 0x10FFFF || (cp >= 0xD800 && cp <= 0xDFFF)) return 0;
+  return len;
+}
+
+// Decodes a run of byte-fallback bytes as UTF-8, appending to `out`. As in
+// SentencePiece, well-formed sequences are kept and every other byte becomes
+// U+FFFD. (HF ByteFallback instead replaces every byte of a run that is not
+// entirely valid; the two agree on valid runs.)
 void FlushBytes(std::string* bytes, std::string* out) {
   if (bytes->empty()) return;
   size_t i = 0;
   const std::string& b = *bytes;
   while (i < b.size()) {
-    const size_t clen = Utf8Len(static_cast<unsigned char>(b[i]));
-    if (i + clen <= b.size()) {
+    const size_t clen = ValidUtf8Len(b, i);
+    if (clen != 0) {
       out->append(b, i, clen);
       i += clen;
     } else {

tokenizer/bpe_tokenizer_test.cc

@@ -143,6 +143,74 @@ TEST_F(BpeTokenizerTest, DecodeRoundTrip) {
   }
 }
 
+// Byte-fallback runs that are not valid UTF-8, such as a single token of a
+// multi-byte character (streaming callers decode one token at a time), must
+// still decode to valid UTF-8, as in SentencePiece: well-formed sequences are
+// kept and every other byte becomes U+FFFD.
+TEST_F(BpeTokenizerTest, DecodeMalformedByteFallbackMatchesSentencePiece) {
+  // Returns the byte-fallback token ids ("<0xHH>") for the bytes of `s`. These
+  // tokens have consecutive ids in the test vocabulary.
+  const auto bytes = [](const std::string& s) {
+    constexpr int kByte0Id = 238;  // "<0x00>"
+    std::vector<int> ids;
+    for (unsigned char c : s) ids.push_back(kByte0Id + c);
+    return ids;
+  };
+  const std::string kRepl = "\xEF\xBF\xBD";  // U+FFFD
+
+  std::vector<int> truncated_then_word = bytes("\xE2\x80");
+  std::vector<int> word;
+  ASSERT_TRUE(sp_->Encode(" word", &word));
+  truncated_then_word.insert(truncated_then_word.end(), word.begin(),
+                             word.end());
+
+  struct Case {
+    std::vector<int> ids;
+    std::string expected;
+  };
+  const std::vector<Case> cases = {
+      {bytes("\xC3\xA9"), "\xC3\xA9"},                  // valid U+00E9
+      {bytes("\xF0\x9F\x98\x80"), "\xF0\x9F\x98\x80"},  // valid U+1F600
+      {bytes("\xA0"), kRepl},                           // lone continuation
+      {bytes("\xFF"), kRepl},                           // invalid lead byte
+      {bytes("\xC2\xC2"), kRepl + kRepl},               // no continuation
+      {bytes("\xE2\x80"), kRepl + kRepl},               // truncated
+      {bytes("\xC0\xAF"), kRepl + kRepl},               // overlong
+      {bytes("\xE0\x80\x80"), kRepl + kRepl + kRepl},   // overlong
+      {bytes("\xED\xA0\x80"), kRepl + kRepl + kRepl},   // surrogate
+      {bytes("\xF4\x90\x80\x80"), kRepl + kRepl + kRepl + kRepl},  // >U+10FFFF
+      // Valid prefix kept, as in SentencePiece (HF ByteFallback: 3 x U+FFFD).
+      {bytes("\xC3\xA9\xA9"), "\xC3\xA9" + kRepl},  // valid, then stray byte
+      {truncated_then_word, kRepl + kRepl + " word"},
+  };
+  for (const auto& [ids, expected] : cases) {
+    std::string sp_text, bpe_text;
+    ASSERT_TRUE(sp_->Decode(ids, &sp_text));
+    ASSERT_TRUE(bpe_->Decode(ids, &bpe_text));
+    EXPECT_EQ(sp_text, expected) << testing::PrintToString(ids);
+    EXPECT_EQ(bpe_text, expected) << testing::PrintToString(ids);
+  }
+
+  // Every run of one or two byte-fallback tokens.
+  size_t mismatches = 0;
+  for (int b0 = 0; b0 < 256; ++b0) {
+    for (int b1 = -1; b1 < 256; ++b1) {
+      std::string run(1, static_cast<char>(b0));
+      if (b1 >= 0) run.push_back(static_cast<char>(b1));
+      std::string sp_text, bpe_text;
+      ASSERT_TRUE(sp_->Decode(bytes(run), &sp_text));
+      ASSERT_TRUE(bpe_->Decode(bytes(run), &bpe_text));
+      if (bpe_text != sp_text && mismatches++ == 0) {
+        ADD_FAILURE() << "First mismatch on bytes "
+                      << testing::PrintToString(run) << ": bpe "
+                      << testing::PrintToString(bpe_text) << ", sentencepiece "
+                      << testing::PrintToString(sp_text);
+      }
+    }
+  }
+  EXPECT_EQ(mismatches, 0u);
+}
+
 // A tokenizer rebuilt from the compact binary should behave identically to
 // the one parsed from tokenizer.json.
 TEST_F(BpeTokenizerTest, CompactRoundTrip) {