Fix Gemma byte-fallback detokenization for UTF-8 output

cactus/engine/engine.h

@@ -6,6 +6,7 @@
 #include <unordered_set>
 #include <memory>
 #include <cstdint>
+#include <string_view>
 
 #include "../graph/graph.h"
 
@@ -294,6 +295,9 @@ TokenizerRuntimeConfig load_tokenizer_runtime_config(const std::string& config_f
 void load_special_tokens_map(const std::string& config_file, std::unordered_map<std::string, uint32_t>& special_tokens);
 std::vector<std::string> split_with_special_tokens(const std::string& text, const std::unordered_map<std::string, uint32_t>& special_tokens);
 
+bool parse_byte_fallback_piece(std::string_view piece, uint8_t* out_byte);
+std::string reassemble_byte_fallback(const std::string& text);
+
 inline std::string extract_json_string(const std::string& json, size_t& pos) {
     std::string value;
     while (pos < json.size() && json[pos] != '"') {

cactus/engine/engine_bpe.cpp

@@ -520,6 +520,16 @@ std::vector<uint32_t> BPETokenizer::encode(const std::string& text) const {
 }
 
 std::string BPETokenizer::decode(const std::vector<uint32_t>& tokens) const {
+    if (runtime_config_.byte_fallback && tokens.size() == 1) {
+        uint32_t token_id = tokens[0];
+        if (token_id < id_to_token_.size()) {
+            uint8_t byte_value;
+            if (parse_byte_fallback_piece(id_to_token_[token_id], &byte_value)) {
+                return std::string(1, static_cast<char>(byte_value));
+            }
+        }
+    }
+
     if (runtime_config_.decoder == TokenizerRuntimeConfig::Decoder::REPLACE_METASPACE) {
         std::string result;
         result.reserve(tokens.size() * 4);
@@ -534,6 +544,9 @@ std::string BPETokenizer::decode(const std::vector<uint32_t>& tokens) const {
             result.replace(pos, 3, " ");
             pos += 1;
         }
+        if (runtime_config_.byte_fallback) {
+            result = reassemble_byte_fallback(result);
+        }
         return result;
     }
 
@@ -558,7 +571,11 @@ std::string BPETokenizer::decode(const std::vector<uint32_t>& tokens) const {
         }
     }
 
-    return unicode_to_bytes(unicode_result);
+    std::string result = unicode_to_bytes(unicode_result);
+    if (runtime_config_.byte_fallback) {
+        result = reassemble_byte_fallback(result);
+    }
+    return result;
 }
 
 

cactus/engine/engine_sp.cpp

@@ -496,7 +496,11 @@ std::string SPTokenizer::decode(const std::vector<uint32_t>& tokens) const {
             raw += piece;
         }
     }
-    return postprocess_text(raw);
+    std::string result = postprocess_text(raw);
+    if (runtime_config_.byte_fallback) {
+        result = reassemble_byte_fallback(result);
+    }
+    return result;
 }
 
 void SPTokenizer::load_special_tokens(const std::string& config_file) {

cactus/engine/engine_tokenizer.cpp

@@ -1,8 +1,10 @@
 #include "engine.h"
+#include <string_view>
 #include <fstream>
 #include <sstream>
 #include <algorithm>
 #include <cmath>
+#include <cstdio>
 #include <map>
 
 extern "C" {
@@ -1014,5 +1016,82 @@ std::string Tokenizer::format_youtu_style(const std::vector<ChatMessage>& messag
     return result;
 }
 
+
+
+bool parse_byte_fallback_piece(std::string_view piece, uint8_t* out_byte) {
+    if (piece.size() != 6) return false;
+    if (piece[0] != '<' || piece[1] != '0' || piece[2] != 'x' || piece[5] != '>') return false;
+    auto hex_value = [](char c) -> int {
+        if (c >= '0' && c <= '9') return c - '0';
+        if (c >= 'a' && c <= 'f') return 10 + (c - 'a');
+        if (c >= 'A' && c <= 'F') return 10 + (c - 'A');
+        return -1;
+    };
+    int hi = hex_value(piece[3]);
+    int lo = hex_value(piece[4]);
+    if (hi < 0 || lo < 0) return false;
+    *out_byte = static_cast<uint8_t>((hi << 4) | lo);
+    return true;
+}
+
+namespace {
+
+bool is_valid_utf8(const std::string& s) {
+    size_t i = 0;
+    while (i < s.size()) {
+        uint8_t b = static_cast<uint8_t>(s[i]);
+        size_t n;
+        if (b < 0x80) n = 1;
+        else if ((b & 0xE0) == 0xC0) n = 2;
+        else if ((b & 0xF0) == 0xE0) n = 3;
+        else if ((b & 0xF8) == 0xF0) n = 4;
+        else return false;
+        if (i + n > s.size()) return false;
+        for (size_t k = 1; k < n; k++) {
+            if ((static_cast<uint8_t>(s[i + k]) & 0xC0) != 0x80) return false;
+        }
+        i += n;
+    }
+    return true;
+}
+
+}
+
+std::string reassemble_byte_fallback(const std::string& text) {
+    std::string out;
+    out.reserve(text.size());
+    std::string pending;
+
+    auto flush_pending = [&]() {
+        if (pending.empty()) return;
+        if (is_valid_utf8(pending)) {
+            out += pending;
+        } else {
+            for (unsigned char raw : pending) {
+                char buf[7];
+                std::snprintf(buf, sizeof(buf), "<0x%02X>", raw);
+                out += buf;
+            }
+        }
+        pending.clear();
+    };
+
+    size_t i = 0;
+    while (i < text.size()) {
+        uint8_t byte_value;
+        if (i + 6 <= text.size() &&
+            parse_byte_fallback_piece(std::string_view(text.data() + i, 6), &byte_value)) {
+            pending.push_back(static_cast<char>(byte_value));
+            i += 6;
+        } else {
+            flush_pending();
+            out.push_back(text[i]);
+            i++;
+        }
+    }
+    flush_pending();
+    return out;
+}
+
 } // namespace engine
 } // namespace cactus

tests/test_byte_fallback_detokenize.cpp

@@ -0,0 +1,95 @@
+#include "test_utils.h"
+#include "../cactus/engine/engine.h"
+#include <iostream>
+#include <string>
+
+using cactus::engine::reassemble_byte_fallback;
+using cactus::engine::parse_byte_fallback_piece;
+
+static bool eq(const std::string& got, const std::string& want, const char* label) {
+    if (got == want) return true;
+    std::cerr << "[mismatch] " << label << "\n  got : ";
+    for (unsigned char c : got) {
+        if (c >= 0x20 && c < 0x7F) std::cerr << static_cast<char>(c);
+        else std::cerr << "\\x" << std::hex << static_cast<int>(c) << std::dec;
+    }
+    std::cerr << "\n  want: ";
+    for (unsigned char c : want) {
+        if (c >= 0x20 && c < 0x7F) std::cerr << static_cast<char>(c);
+        else std::cerr << "\\x" << std::hex << static_cast<int>(c) << std::dec;
+    }
+    std::cerr << "\n";
+    return false;
+}
+
+static bool test_parse_piece_valid() {
+    uint8_t out = 0;
+    if (!parse_byte_fallback_piece("<0xEA>", &out) || out != 0xEA) return false;
+    if (!parse_byte_fallback_piece("<0x00>", &out) || out != 0x00) return false;
+    if (!parse_byte_fallback_piece("<0xff>", &out) || out != 0xFF) return false;
+    return true;
+}
+
+static bool test_parse_piece_rejects() {
+    uint8_t out = 0;
+    if (parse_byte_fallback_piece("<0xEA", &out)) return false;
+    if (parse_byte_fallback_piece("<0xGG>", &out)) return false;
+    if (parse_byte_fallback_piece("hello", &out)) return false;
+    if (parse_byte_fallback_piece("<0xEA>x", &out)) return false;
+    return true;
+}
+
+static bool test_korean_three_byte_reassembly() {
+    std::string in = std::string("\xEC\x9E\xA0") + " <0xEA><0xB9><0xB0> " + std::string("\xEC\x88\x98");
+    std::string want = std::string("\xEC\x9E\xA0") + " " + std::string("\xEA\xB9\xB0") + " " + std::string("\xEC\x88\x98");
+    return eq(reassemble_byte_fallback(in), want, "korean_three_byte_reassembly");
+}
+
+static bool test_polish_two_byte_reassembly() {
+    std::string in = "<0xC4><0x85><0xC4><0x99>";
+    std::string want = std::string("\xC4\x85") + std::string("\xC4\x99");
+    return eq(reassemble_byte_fallback(in), want, "polish_two_byte_reassembly");
+}
+
+static bool test_ascii_byte_fallback_space() {
+    return eq(reassemble_byte_fallback("A<0x20>B"), "A B", "ascii_byte_fallback_space");
+}
+
+static bool test_invalid_partial_preserves_literal() {
+    return eq(reassemble_byte_fallback("hi<0xEA>x"), "hi<0xEA>x", "invalid_partial_preserves_literal");
+}
+
+static bool test_trailing_incomplete_preserved() {
+    return eq(reassemble_byte_fallback("hi<0xEA><0xB9>"), "hi<0xEA><0xB9>", "trailing_incomplete_preserved");
+}
+
+static bool test_empty_passthrough() {
+    return eq(reassemble_byte_fallback(""), "", "empty_passthrough") &&
+           eq(reassemble_byte_fallback("hello world"), "hello world", "no_byte_tokens_passthrough");
+}
+
+static bool test_two_runs_separated_by_text() {
+    std::string in = "<0xEA><0xB9><0xB0>X<0xEC><0x88><0x98>";
+    std::string want = std::string("\xEA\xB9\xB0") + "X" + std::string("\xEC\x88\x98");
+    return eq(reassemble_byte_fallback(in), want, "two_runs_separated_by_text");
+}
+
+static bool test_lowercase_hex_accepted() {
+    return eq(reassemble_byte_fallback("<0xea><0xb9><0xb0>"), std::string("\xEA\xB9\xB0"), "lowercase_hex_accepted");
+}
+
+int main() {
+    TestUtils::TestRunner runner("Byte-fallback detokenize");
+    runner.run_test("parse_piece_valid", test_parse_piece_valid());
+    runner.run_test("parse_piece_rejects", test_parse_piece_rejects());
+    runner.run_test("korean_three_byte_reassembly", test_korean_three_byte_reassembly());
+    runner.run_test("polish_two_byte_reassembly", test_polish_two_byte_reassembly());
+    runner.run_test("ascii_byte_fallback_space", test_ascii_byte_fallback_space());
+    runner.run_test("invalid_partial_preserves_literal", test_invalid_partial_preserves_literal());
+    runner.run_test("trailing_incomplete_preserved", test_trailing_incomplete_preserved());
+    runner.run_test("empty_and_passthrough", test_empty_passthrough());
+    runner.run_test("two_runs_separated_by_text", test_two_runs_separated_by_text());
+    runner.run_test("lowercase_hex_accepted", test_lowercase_hex_accepted());
+    runner.print_summary();
+    return runner.all_passed() ? 0 : 1;
+}