Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 11 additions & 11 deletions Sources/Tokenizers/Decoder.swift
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,8 @@ class WordPieceDecoder: Decoder {
}

func decode(tokens: [String]) -> [String] {
let firstToken = cleanup ? cleanUpTokenization(tokens.first!) : tokens.first!
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)"
Expand Down Expand Up @@ -187,23 +188,22 @@ class ByteFallbackDecoder: Decoder {
return Int(token[startIndex..<endIndex], radix: 16)
}

func flushPendingBytes() {
guard !byteTokens.isEmpty else { return }
let codeUnits = byteTokens.map { UTF8.CodeUnit($0) }
newTokens.append(String(decoding: codeUnits, as: UTF8.self))
byteTokens.removeAll()
}

for token in tokens {
if let byte = parseByte(token) {
byteTokens.append(byte)
} else {
if !byteTokens.isEmpty {
// decode as utf8 and append
let codeUnits = byteTokens.map { UTF8.CodeUnit($0) }
newTokens.append(String(decoding: codeUnits, as: UTF8.self))
byteTokens.removeAll()
}
flushPendingBytes()
newTokens.append(token)
}
}
if !byteTokens.isEmpty {
let codeUnits = byteTokens.map { UTF8.CodeUnit($0) }
newTokens.append(String(decoding: codeUnits, as: UTF8.self))
}
flushPendingBytes()
return newTokens
}
}
Expand Down
15 changes: 15 additions & 0 deletions Tests/TokenizersTests/DecoderTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,21 @@ 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.
Comment thread
pcuenca marked this conversation as resolved.
Resolved
@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: []) == [])
}

@Test("WordPiece decoder with prefix and cleanup")
func wordPieceDecoder() {
let config = Config(["prefix": "##", "cleanup": true])
Expand Down