diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift index f5873808..6b71b70f 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/ConstrainedDecodingStrategy.swift @@ -168,8 +168,10 @@ public struct ConstrainedDecodingStrategy: DecodingStrategy { let fullDecodedText = tokenizer.decode(tokens: generatedTokens.map { Int($0) }) decodeSpan.end() - let common = fullDecodedText.commonPrefix(with: previousDecodedText) - let delta = String(fullDecodedText.dropFirst(common.count)) + // Compare Unicode scalars: a Character-wise prefix re-emits the grapheme a token extends. + let fullScalars = fullDecodedText.unicodeScalars + let common = zip(fullScalars, previousDecodedText.unicodeScalars).prefix(while: { $0 == $1 }).count + let delta = String(String.UnicodeScalarView(fullScalars.dropFirst(common))) if delta.unicodeScalars.contains(where: { $0 == "\u{FFFD}" }) { return "" diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/PipelinedConstrainedDecodingStrategy.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/PipelinedConstrainedDecodingStrategy.swift index cfd5565f..caae2dd0 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/PipelinedConstrainedDecodingStrategy.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/PipelinedConstrainedDecodingStrategy.swift @@ -211,8 +211,10 @@ extension PipelinedConstrainedSequence { // Decode text incrementally generatedTokens.append(tokenId) let fullDecode = tokenizer.decode(tokens: generatedTokens.map { Int($0) }) - let common = fullDecode.commonPrefix(with: previousDecodedText) - let delta = String(fullDecode.dropFirst(common.count)) + // Compare Unicode scalars: a Character-wise prefix re-emits the grapheme a token extends. + let fullScalars = fullDecode.unicodeScalars + let common = zip(fullScalars, previousDecodedText.unicodeScalars).prefix(while: { $0 == $1 }).count + let delta = String(String.UnicodeScalarView(fullScalars.dropFirst(common))) if delta.unicodeScalars.contains(where: { $0 == "\u{FFFD}" }) { consecutiveDecodeFailures += 1 diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/VanillaDecodingStrategy.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/VanillaDecodingStrategy.swift index b6f447b2..7f54e933 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/VanillaDecodingStrategy.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/VanillaDecodingStrategy.swift @@ -192,9 +192,13 @@ extension VanillaDecodingStrategy.VanillaDecodedSequence { let safeEnd = VanillaDecodingStrategy.safeUTF8Prefix(of: fullDecode) let emittable = String(fullDecode[fullDecode.startIndex.. pendingText.count, emittable.hasPrefix(pendingText) { - newText = String(emittable.dropFirst(pendingText.count)) + if emittedScalars.count > pendingScalars.count, emittedScalars.starts(with: pendingScalars) { + newText = String(String.UnicodeScalarView(emittedScalars.dropFirst(pendingScalars.count))) pendingText = emittable } else if emittable != pendingText { newText = emittable @@ -273,8 +277,10 @@ extension VanillaDecodingStrategy.VanillaDecodedSequence { let finalDecode = tokenizer.decode(tokens: newlyGeneratedTokens) let trailing: GenerationResult? - if finalDecode.count > pendingText.count, finalDecode.hasPrefix(pendingText) { - let remaining = String(finalDecode.dropFirst(pendingText.count)) + let finalScalars = finalDecode.unicodeScalars + let pendingScalars = pendingText.unicodeScalars + if finalScalars.count > pendingScalars.count, finalScalars.starts(with: pendingScalars) { + let remaining = String(String.UnicodeScalarView(finalScalars.dropFirst(pendingScalars.count))) trailing = remaining.isEmpty ? nil diff --git a/swift/Tests/LanguageModelsTests/StreamingTextDeltaTests.swift b/swift/Tests/LanguageModelsTests/StreamingTextDeltaTests.swift new file mode 100644 index 00000000..c13993c6 --- /dev/null +++ b/swift/Tests/LanguageModelsTests/StreamingTextDeltaTests.swift @@ -0,0 +1,88 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import TestUtilities +import Testing + +@testable import CoreAILanguageModels + +/// Streamed text deltas must concatenate to the decoded text. A token that extends the +/// previous grapheme cluster (combining mark, ZWJ emoji sequence, Indic vowel sign) leaves +/// the `Character` count unchanged, so a Character-based diff re-emits text. +@Suite("Streaming text deltas") +struct StreamingTextDeltaTests { + /// MockTokenizer emits one token per UTF-8 byte, so every Unicode scalar arrives on its own. + private static let texts = [ + "cafe\u{0301} au lait", + "family \u{1F468}\u{200D}\u{1F469}\u{200D}\u{1F467} ok", + "\u{0B95}\u{0BB3}\u{0BBF}\u{0BB2}\u{0BCD}", + ] + + private let tokenizer = MockTokenizer() + + private func tokens(for text: String) -> [Int32] { + tokenizer.encode(text: text).map(Int32.init) + } + + private func concatenated(_ stream: some AsyncSequence) async throws -> String { + var text = "" + for try await result in stream { + text += result.text + } + return text + } + + @Test("Vanilla decoding streams a grapheme extension once", arguments: texts) + func vanilla(text: String) async throws { + let scripted = tokens(for: text) + let stream = try await VanillaDecodingStrategy().decode( + from: .tokens([1]), + tokenizer: tokenizer, + inferenceEngine: MockEngine(tokens: scripted, vocabSize: nil), + samplingConfiguration: .greedy, + options: InferenceOptions(maxTokens: scripted.count), + stopSequences: StopSequences(for: tokenizer, additionalEosTokenIds: []) + ) + + let streamed = try await concatenated(stream) + #expect(streamed == text) + } + + @Test("Constrained decoding streams a grapheme extension once", arguments: texts) + func constrained(text: String) async throws { + let json = "\"\(text)\"" + let scripted = tokens(for: json) + let stream = try await ConstrainedDecodingStrategy(jsonSchema: #"{"type": "string"}"#, vocabSize: 256) + .decode( + from: .tokens([1]), + tokenizer: tokenizer, + inferenceEngine: MockEngine(tokens: scripted, vocabSize: 256), + samplingConfiguration: .greedy, + options: InferenceOptions(maxTokens: scripted.count), + stopSequences: StopSequences(for: tokenizer, additionalEosTokenIds: []) + ) + + let streamed = try await concatenated(stream) + #expect(streamed == json) + } + + @Test("Pipelined constrained decoding streams a grapheme extension once", arguments: texts) + func pipelinedConstrained(text: String) async throws { + let json = "\"\(text)\"" + let scripted = tokens(for: json) + let stream = try await PipelinedConstrainedDecodingStrategy(jsonSchema: #"{"type": "string"}"#, vocabSize: 256) + .decode( + from: .tokens([1]), + tokenizer: tokenizer, + inferenceEngine: MockConstrainedEngine(scriptedTokens: scripted), + samplingConfiguration: .greedy, + options: InferenceOptions(maxTokens: scripted.count), + stopSequences: StopSequences(for: tokenizer, additionalEosTokenIds: []) + ) + + let streamed = try await concatenated(stream) + #expect(streamed == json) + } +}