Repository navigation
Constrained generation always picks the first anyOf variant #289
Description
Local models that use constrained generation (MLX, Core ML, and llama.cpp) always generate the first variant of an anyOf schema, because generateNode(_:) in StructuredGeneration.swift takes variants[0]. So a nullable union like anyOf: [person, .null] (#288) can never produce null, and a union of object types always produces the first one.
The model should choose the branch. One approach is to mask the next token to the start tokens of every variant, as itemStartTokens(for:) already does for array items, and continue with the variant whose start tokens include the sampled token. Variants that start the same way, like two object types, would need to look ahead to the first key.
Activity
2570165831 commentedon Oct 6, 2026
I'd like to take this. Plan: resolve refs and flatten nested unions, then sample once (without decoding, like the empty-array probe) from the union of each variant's
itemStartTokens(for:)and continue with the matching variant. That covers[T, null]and unions of different JSON types. When the remaining variants are all objects, emit{and let the key mask offer the keys of every live variant, dropping variants that don't declare the chosen key. Since #301 gives every case atypekey, a shared key whose values are string choices would narrow by the chosen value. Two questions: for variants that start alike but aren't objects (two strings, integer vs number, two arrays), is falling back to the first matching variant OK for now? And should the discriminator part wait for #301? Tests would useMockTokenBackendwith scripted sampling.mattt commentedon Oct 7, 2026
Thanks for offering to take this, @2570165831. Your plan sounds good to me. To your questions:
enuminside the schema, so narrowing on a shared key's string choices covers constants as well. You can test that with hand-built schemas now, and Generate enums with associated values as Foundation Models does #301'stypekeys will work the same way once it lands.Let the model choose the variant of an anyOf (huggingface#320)