From 881abc9c898bc7deed71e55f0a129d0d30362af3 Mon Sep 17 00:00:00 2001 From: james-333i Date: Tue, 25 Aug 2026 12:10:56 -0700 Subject: [PATCH 1/2] Widen the empty-array probe to merged item-start tokens Byte-pair vocabularies put most item-start probability on merged tokens, such as a quote fused with the first word or carrying a leading space. The probe offered only the bare single-character tokens, so the comparison against the closing bracket was noise and small models closed arrays the prompt asked them to fill. Admit every token whose first non-whitespace character starts the item type, and treat any token that trims to the closing bracket as a close. --- .../Shared/StructuredGeneration.swift | 52 +++++++++++++++---- 1 file changed, 42 insertions(+), 10 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index e02d88ab..414a5b21 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -580,31 +580,63 @@ struct ConstrainedJSONGenerator { /// Probe after `[` when `minItems == 0`: model may close immediately or start an item. /// - /// Samples once among `]` and tokens that can start the item type. Choosing `]` means - /// close; any other sample is discarded without decoding so ``generateNode`` can emit - /// the first element from the same backend state. + /// Samples once among closing tokens and tokens that can start the item type. Choosing + /// a close means close; any other sample is discarded without decoding so + /// ``generateNode`` can emit the first element from the same backend state. + /// + /// Byte-pair vocabularies put most item-start probability on merged tokens (a quote + /// fused with the first word, or a leading space), so the probe admits every token + /// whose first non-whitespace character starts the item type, not just the bare + /// single-character token. private mutating func sampleWhetherToCloseEmptyArray( items: GenerationSchema.Node ) async throws -> Bool { - let closeToken = try Self.singleToken(for: "]", backend: backend) + let closeTokens = tokensMatchingTrimmed("]") var allowed = try itemStartTokens(for: items) - allowed.insert(closeToken) - guard !allowed.isEmpty else { + allowed.formUnion(closeTokens) + guard !allowed.isEmpty, !closeTokens.isEmpty else { return false } let token = try await backend.sample(from: allowed) - return token == closeToken + return closeTokens.contains(token) + } + + /// Tokens whose text equals `text` after trimming surrounding whitespace. + private func tokensMatchingTrimmed(_ text: String) -> Set { + var tokens = Set() + for token in 0 ..< backend.vocabSize { + if backend.isSpecialToken(token) { continue } + guard let tokenText = backend.tokenText(token) else { continue } + if tokenText.trimmingCharacters(in: .whitespacesAndNewlines) == text { + tokens.insert(token) + } + } + return tokens + } + + /// Tokens whose first non-whitespace character is `prefix`. + private func tokensStarting(with prefix: Character) -> Set { + var tokens = Set() + for token in 0 ..< backend.vocabSize { + if backend.isSpecialToken(token) { continue } + guard let text = backend.tokenText(token) else { continue } + guard let first = text.drop(while: { $0.isWhitespace }).first else { continue } + if first == prefix { + tokens.insert(token) + } + } + return tokens } /// Tokens that can begin a JSON value for `node` (empty-array probe). private func itemStartTokens(for node: GenerationSchema.Node) throws -> Set { switch node { case .string: - return [quoteToken] + return tokensStarting(with: "\"") case .object: - return [try Self.singleToken(for: "{", backend: backend)] + return tokensStarting(with: "{") case .array: - return [try Self.singleToken(for: "[", backend: backend)] + return tokensStarting(with: "[") case .boolean: var tokens = Set() for literal in ["true", "false"] { From d8eb0cbb11e4fd2b4635766f2c44f8f747b6cce6 Mon Sep 17 00:00:00 2001 From: mattt Date: Fri, 4 Sep 2026 05:29:23 -0700 Subject: [PATCH 2/2] Apply the widened probe to boolean and number items Boolean starts admit any token that is a prefix of true or false after leading whitespace, and number starts look at the first non-whitespace character, matching the string, object, and array cases. Adds a test with merged item-start and close tokens. --- .../Shared/StructuredGeneration.swift | 30 +++++++++----- .../StructuredGenerationTests.swift | 41 +++++++++++++++++++ 2 files changed, 61 insertions(+), 10 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index 414a5b21..18c6736d 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -614,6 +614,21 @@ struct ConstrainedJSONGenerator { return tokens } + /// Tokens whose text, after leading whitespace, is a non-empty prefix of one of `literals`. + private func tokensPrefixing(anyOf literals: [String]) -> Set { + var tokens = Set() + for token in 0 ..< backend.vocabSize { + if backend.isSpecialToken(token) { continue } + guard let text = backend.tokenText(token) else { continue } + let trimmed = text.drop(while: { $0.isWhitespace }) + guard !trimmed.isEmpty else { continue } + if literals.contains(where: { $0.hasPrefix(trimmed) }) { + tokens.insert(token) + } + } + return tokens + } + /// Tokens whose first non-whitespace character is `prefix`. private func tokensStarting(with prefix: Character) -> Set { var tokens = Set() @@ -638,13 +653,7 @@ struct ConstrainedJSONGenerator { case .array: return tokensStarting(with: "[") case .boolean: - var tokens = Set() - for literal in ["true", "false"] { - if let first = try backend.tokenize(literal).first { - tokens.insert(first) - } - } - return tokens + return tokensPrefixing(anyOf: ["true", "false"]) case .number(let numberNode): let numeric = numberNode.integerOnly @@ -653,9 +662,10 @@ struct ConstrainedJSONGenerator { // Only tokens that can start a number (digit or minus — not a bare `.`). return Set( numeric.filter { token in - guard let text = backend.tokenText(token), !text.isEmpty else { return false } - let first = text.first - return first?.isNumber == true || first == "-" + guard let text = backend.tokenText(token), + let first = text.drop(while: { $0.isWhitespace }).first + else { return false } + return first.isNumber || first == "-" } ) case .ref(let typeName): diff --git a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift index eff70b83..2b1963bc 100644 --- a/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/StructuredGenerationTests.swift @@ -588,6 +588,47 @@ struct StructuredGenerationTests { #expect(result == "[]") } + @Test func emptyArrayProbeAdmitsMergedTokens() async throws { + var maps = baseTokenMaps() + // Byte-pair vocabularies carry the item start and the close on merged tokens. + let spacedBracket = 60 + let quoteA = 61 + maps.tokenToText[spacedBracket] = " ]" + maps.tokenToText[quoteA] = "\"a" + let arrayNode = GenerationSchema.ArrayNode( + description: nil, + items: .string(.init(enumChoices: ["a"])), + minItems: nil, + maxItems: 1 + ) + let schema = GenerationSchema.primitive([String].self, node: .array(arrayNode)) + let eosToken = 50 + + // The probe must offer the whitespace-prefixed close, and choosing it closes the array. + let closing = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [spacedBracket] + ) + var closingGenerator = try ConstrainedJSONGenerator(backend: closing, schema: schema) + #expect(try await closingGenerator.generate() == "[]") + + // The probe must offer the merged item start, and choosing it fills the array. + let filling = MockTokenBackend( + tokenToText: maps.tokenToText, + textToTokens: maps.textToTokens, + eosToken: eosToken, + endTokens: [eosToken], + maximumTokens: 64, + samplingQueue: [quoteA] + ) + var fillingGenerator = try ConstrainedJSONGenerator(backend: filling, schema: schema) + #expect(try await fillingGenerator.generate() == "[\"a\"]") + } + @Test func arrayTruncatesUnderBudgetPressure() async throws { let maps = baseTokenMaps() let arrayNode = GenerationSchema.ArrayNode(