From a5c229982636a70299b82cff20f920c35b4a2780 Mon Sep 17 00:00:00 2001 From: john-rocky Date: Fri, 15 May 2026 23:04:02 +0900 Subject: [PATCH 1/5] Fix decoder crashes / drops on empty and trailing-byte inputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two independent bugs in Sources/Tokenizers/Decoder.swift: - WordPieceDecoder.decode(tokens:) force-unwrapped `tokens.first!` and crashed when handed an empty list. An empty list can legitimately reach the decoder via `Tokenizer.decode(tokens:, skipSpecialTokens: true)` when the upstream compactMap or special-token filter drops every entry. Early-return [] in that case to match the reference Rust behavior. - ByteFallbackDecoder.decode(tokens:) accumulated runs of `<0xHH>` byte-fallback tokens and only flushed them to a UTF-8 string when it saw a non-byte token. Inputs that ended on a byte run (trailing emoji, CJK, or `<0x0A>` newline) had those bytes silently dropped. Added a matching post-loop flush so `["<0xC2>", "<0xA9>"]` decodes to `["©"]` and `["hi", "<0xF0>", "<0x9F>", "<0x9A>", "<0x80>"]` decodes to `["hi", "🚀"]`. Tests: two new regression tests covering both bugs. One existing test ("Jinja block whitespace control") was updated — the prior expected string omitted the trailing newline that the rendered template legitimately ends with, because the pre-fix ByteFallbackDecoder dropped the `<0x0A>` byte token. A comment explains the dependency. --- Sources/Tokenizers/Decoder.swift | 28 ++++++++++---- Tests/TokenizersTests/ChatTemplateTests.swift | 11 ++++-- Tests/TokenizersTests/DecoderTests.swift | 37 +++++++++++++++++++ 3 files changed, 65 insertions(+), 11 deletions(-) diff --git a/Sources/Tokenizers/Decoder.swift b/Sources/Tokenizers/Decoder.swift index 49021645..0cb3f626 100644 --- a/Sources/Tokenizers/Decoder.swift +++ b/Sources/Tokenizers/Decoder.swift @@ -84,7 +84,14 @@ class WordPieceDecoder: Decoder { } func decode(tokens: [String]) -> [String] { - let firstToken = cleanup ? cleanUpTokenization(tokens.first!) : tokens.first! + // An empty token list can reach this decoder when the calling chain has + // filtered out every input — e.g. `decode(tokens:, skipSpecialTokens: true)` + // on an id sequence that consisted entirely of special tokens, or after a + // `compactMap { convertIdToToken($0) }` drops every entry as out-of-vocab. + // The reference Rust implementation no-ops in that case rather than + // unwrapping `tokens.first!` (which previously crashed). + guard let first = tokens.first else { return [] } + let firstToken = cleanup ? cleanUpTokenization(first) : first return [firstToken] + tokens.dropFirst().map { token in let token = token.hasPrefix(prefix) ? token.replacingCharacters(in: token.range(of: prefix)!, with: "") : " \(token)" @@ -186,19 +193,26 @@ class ByteFallbackDecoder: Decoder { return Int(token[startIndex..` (e.g. multi-byte UTF-8 for a final emoji or CJK character), + // the previous implementation dropped them on the floor instead of + // appending the decoded UTF-8 string. + flushPendingBytes() return newTokens } } diff --git a/Tests/TokenizersTests/ChatTemplateTests.swift b/Tests/TokenizersTests/ChatTemplateTests.swift index 50f0117e..007d9331 100644 --- a/Tests/TokenizersTests/ChatTemplateTests.swift +++ b/Tests/TokenizersTests/ChatTemplateTests.swift @@ -181,10 +181,13 @@ struct ChatTemplateTests { messages: messages, chatTemplate: whitespaceSensitiveTemplate ) let decoded = tokenizer.decode(tokens: encoded) - let expected = """ - Describe the Swift programming language. - assistant - """ + // The template ends with a `{% endif %}` block on its own line, so the + // rendered text legitimately ends with "\n". Phi's tokenizer encodes that + // trailing newline as a `<0x0A>` byte-fallback token; before + // `ByteFallbackDecoder` was fixed to flush its pending byte buffer at + // end-of-input, the trailing newline was silently dropped and this test + // appeared to pass on the truncated string. + let expected = "Describe the Swift programming language.\nassistant\n" #expect(decoded == expected) #expect(!decoded.hasPrefix("\n")) #expect(!decoded.contains("\n\n")) diff --git a/Tests/TokenizersTests/DecoderTests.swift b/Tests/TokenizersTests/DecoderTests.swift index ddd0a735..7e2227c3 100644 --- a/Tests/TokenizersTests/DecoderTests.swift +++ b/Tests/TokenizersTests/DecoderTests.swift @@ -114,6 +114,43 @@ struct DecoderTests { #expect(decoded.joined() == "How are you?") } + /// Regression coverage: `WordPieceDecoder.decode(tokens: [])` previously crashed + /// with `Fatal error: Unexpectedly found nil while unwrapping an Optional value` + /// because of `tokens.first!`. An empty token list can legitimately reach the + /// decoder when `Tokenizer.decode(tokens:, skipSpecialTokens: true)` is called + /// on an id sequence whose non-special-token ids all fail vocab lookup (the + /// `compactMap` upstream drops them silently), so the decoder has to no-op + /// instead of asserting. + @Test("WordPiece decoder no-ops on empty input") + func wordPieceDecoderEmptyInput() { + let decoder = WordPieceDecoder(config: Config(["prefix": "##", "cleanup": true])) + #expect(decoder.decode(tokens: []) == []) + let decoderNoCleanup = WordPieceDecoder(config: Config(["prefix": "##", "cleanup": false])) + #expect(decoderNoCleanup.decode(tokens: []) == []) + } + + /// Regression coverage: `ByteFallbackDecoder.decode` accumulated runs of + /// `<0xHH>` byte tokens and flushed them only when a non-byte token followed. + /// Inputs that ended with a multi-byte UTF-8 run (e.g. a trailing emoji or + /// CJK character whose final code point fell back to bytes) had those bytes + /// silently dropped from the output. + @Test("ByteFallback decoder flushes trailing byte tokens") + func byteFallbackDecoderTrailingBytes() { + let decoder = ByteFallbackDecoder(config: Config(["type": "ByteFallback"])) + // © is U+00A9 → UTF-8 0xC2 0xA9. + #expect(decoder.decode(tokens: ["<0xC2>", "<0xA9>"]) == ["©"]) + // Trailing bytes after a real token must still flush. + #expect(decoder.decode(tokens: ["abc", "<0xC2>", "<0xA9>"]) == ["abc", "©"]) + // 🚀 is U+1F680 → UTF-8 0xF0 0x9F 0x9A 0x80 (four-byte fallback at end). + #expect(decoder.decode(tokens: ["hi", "<0xF0>", "<0x9F>", "<0x9A>", "<0x80>"]) == ["hi", "🚀"]) + // Existing in-loop flush behavior is preserved. + #expect(decoder.decode(tokens: ["<0xC2>", "<0xA9>", "xyz"]) == ["©", "xyz"]) + // Empty input still returns empty. + #expect(decoder.decode(tokens: []) == []) + // Non-byte-only inputs unchanged. + #expect(decoder.decode(tokens: ["abc", "def"]) == ["abc", "def"]) + } + @Test("WordPiece decoder with prefix and cleanup") func wordPieceDecoder() { let config = Config(["prefix": "##", "cleanup": true]) From 316bc3ba88f0b2f667b9e6863a4cada192f675fb Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Thu, 29 Oct 2026 17:04:29 +0100 Subject: [PATCH 2/5] Apply suggestion from @pcuenca --- Sources/Tokenizers/Decoder.swift | 6 ------ 1 file changed, 6 deletions(-) diff --git a/Sources/Tokenizers/Decoder.swift b/Sources/Tokenizers/Decoder.swift index a8ef16fa..0db7e72c 100644 --- a/Sources/Tokenizers/Decoder.swift +++ b/Sources/Tokenizers/Decoder.swift @@ -85,12 +85,6 @@ class WordPieceDecoder: Decoder { } func decode(tokens: [String]) -> [String] { - // An empty token list can reach this decoder when the calling chain has - // filtered out every input — e.g. `decode(tokens:, skipSpecialTokens: true)` - // on an id sequence that consisted entirely of special tokens, or after a - // `compactMap { convertIdToToken($0) }` drops every entry as out-of-vocab. - // The reference Rust implementation no-ops in that case rather than - // unwrapping `tokens.first!` (which previously crashed). guard let first = tokens.first else { return [] } let firstToken = cleanup ? cleanUpTokenization(first) : first return [firstToken] From 9151315da2494d4b8092bbfda0039bd394a9d6bc Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Thu, 29 Oct 2026 17:07:48 +0100 Subject: [PATCH 3/5] Use helper --- Sources/Tokenizers/Decoder.swift | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/Sources/Tokenizers/Decoder.swift b/Sources/Tokenizers/Decoder.swift index 0db7e72c..fc690197 100644 --- a/Sources/Tokenizers/Decoder.swift +++ b/Sources/Tokenizers/Decoder.swift @@ -203,10 +203,7 @@ class ByteFallbackDecoder: Decoder { newTokens.append(token) } } - if !byteTokens.isEmpty { - let codeUnits = byteTokens.map { UTF8.CodeUnit($0) } - newTokens.append(String(decoding: codeUnits, as: UTF8.self)) - } + flushPendingBytes() return newTokens } } From 268b6280225ea06556621ac35d9ea31265f7beb8 Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Thu, 29 Oct 2026 17:09:33 +0100 Subject: [PATCH 4/5] Remove test similar to existing one --- Tests/TokenizersTests/DecoderTests.swift | 22 ---------------------- 1 file changed, 22 deletions(-) diff --git a/Tests/TokenizersTests/DecoderTests.swift b/Tests/TokenizersTests/DecoderTests.swift index 19c4c94d..e1be5e29 100644 --- a/Tests/TokenizersTests/DecoderTests.swift +++ b/Tests/TokenizersTests/DecoderTests.swift @@ -129,28 +129,6 @@ struct DecoderTests { #expect(decoderNoCleanup.decode(tokens: []) == []) } - /// Regression coverage: `ByteFallbackDecoder.decode` accumulated runs of - /// `<0xHH>` byte tokens and flushed them only when a non-byte token followed. - /// Inputs that ended with a multi-byte UTF-8 run (e.g. a trailing emoji or - /// CJK character whose final code point fell back to bytes) had those bytes - /// silently dropped from the output. - @Test("ByteFallback decoder flushes trailing byte tokens") - func byteFallbackDecoderTrailingBytes() { - let decoder = ByteFallbackDecoder(config: Config(["type": "ByteFallback"])) - // © is U+00A9 → UTF-8 0xC2 0xA9. - #expect(decoder.decode(tokens: ["<0xC2>", "<0xA9>"]) == ["©"]) - // Trailing bytes after a real token must still flush. - #expect(decoder.decode(tokens: ["abc", "<0xC2>", "<0xA9>"]) == ["abc", "©"]) - // 🚀 is U+1F680 → UTF-8 0xF0 0x9F 0x9A 0x80 (four-byte fallback at end). - #expect(decoder.decode(tokens: ["hi", "<0xF0>", "<0x9F>", "<0x9A>", "<0x80>"]) == ["hi", "🚀"]) - // Existing in-loop flush behavior is preserved. - #expect(decoder.decode(tokens: ["<0xC2>", "<0xA9>", "xyz"]) == ["©", "xyz"]) - // Empty input still returns empty. - #expect(decoder.decode(tokens: []) == []) - // Non-byte-only inputs unchanged. - #expect(decoder.decode(tokens: ["abc", "def"]) == ["abc", "def"]) - } - @Test("WordPiece decoder with prefix and cleanup") func wordPieceDecoder() { let config = Config(["prefix": "##", "cleanup": true]) From 6dfad968e3f4693b1fc5900f162558abfb11630e Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Thu, 29 Oct 2026 17:09:57 +0100 Subject: [PATCH 5/5] Apply suggestion from @pcuenca --- Tests/TokenizersTests/DecoderTests.swift | 7 ------- 1 file changed, 7 deletions(-) diff --git a/Tests/TokenizersTests/DecoderTests.swift b/Tests/TokenizersTests/DecoderTests.swift index e1be5e29..e1c61be2 100644 --- a/Tests/TokenizersTests/DecoderTests.swift +++ b/Tests/TokenizersTests/DecoderTests.swift @@ -114,13 +114,6 @@ struct DecoderTests { #expect(decoded.joined() == "How are you?") } - /// Regression coverage: `WordPieceDecoder.decode(tokens: [])` previously crashed - /// with `Fatal error: Unexpectedly found nil while unwrapping an Optional value` - /// because of `tokens.first!`. An empty token list can legitimately reach the - /// decoder when `Tokenizer.decode(tokens:, skipSpecialTokens: true)` is called - /// on an id sequence whose non-special-token ids all fail vocab lookup (the - /// `compactMap` upstream drops them silently), so the decoder has to no-op - /// instead of asserting. @Test("WordPiece decoder no-ops on empty input") func wordPieceDecoderEmptyInput() { let decoder = WordPieceDecoder(config: Config(["prefix": "##", "cleanup": true]))