Skip to content

Commit 4021e45

Browse files
petrukha-ivanpetrukha-ivan
committed
Add convenient generation methods
1 parent 3891928 commit 4021e45

File tree

9 files changed

+443
-177
lines changed

9 files changed

+443
-177
lines changed

‎Sources/MLXStructured/Generate.swift‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
//
2+
// Generate.swift
3+
// MLXStructured
4+
//
5+
// Created by Ivan Petrukha on 27.09.2025.
6+
//
7+
8+
import Foundation
9+
import JSONSchema
10+
import MLXLMCommon
11+
import MLX
12+
13+
#if canImport(FoundationModels)
14+
import FoundationModels
15+
#endif
16+
17+
public func generate(
18+
input: LMInput,
19+
parameters: GenerateParameters = GenerateParameters(),
20+
context: ModelContext,
21+
grammar: Grammar,
22+
didGenerate: ([Int]) -> GenerateDisposition = { _ in .more }
23+
) async throws -> GenerateResult {
24+
let sampler = parameters.sampler()
25+
let processor = try await GrammarMaskedLogitProcessor.from(configuration: context.configuration, grammar: grammar)
26+
let iterator = try TokenIterator(input: input, model: context.model, processor: processor, sampler: sampler)
27+
let result = generate(input: input, context: context, iterator: iterator, didGenerate: didGenerate)
28+
return result
29+
}
30+
31+
public func generate<Content: Decodable>(
32+
input: LMInput,
33+
parameters: GenerateParameters = GenerateParameters(),
34+
context: ModelContext,
35+
schema: JSONSchema,
36+
generating: Content.Type,
37+
didGenerate: ([Int]) -> GenerateDisposition = { _ in .more }
38+
) async throws -> (GenerateResult, Content) {
39+
let grammar = try Grammar.schema(schema)
40+
let sampler = parameters.sampler()
41+
let processor = try await GrammarMaskedLogitProcessor.from(configuration: context.configuration, grammar: grammar)
42+
let iterator = try TokenIterator(input: input, model: context.model, processor: processor, sampler: sampler)
43+
let result = generate(input: input, context: context, iterator: iterator, didGenerate: didGenerate)
44+
let content = try JSONDecoder().decode(Content.self, from: Data(result.output.utf8))
45+
return (result, content)
46+
}
47+
48+
#if compiler(>=6.2)
49+
@available(macOS 26.0, iOS 26.0, *)
50+
public func generate<Content: Generable>(
51+
input: LMInput,
52+
parameters: GenerateParameters = GenerateParameters(),
53+
context: ModelContext,
54+
generating: Content.Type,
55+
didGenerate: ([Int]) -> GenerateDisposition = { _ in .more }
56+
) async throws -> (GenerateResult, Content) {
57+
let sampler = parameters.sampler()
58+
let grammar = try Grammar.generable(Content.self)
59+
let processor = try await GrammarMaskedLogitProcessor.from(configuration: context.configuration, grammar: grammar)
60+
let iterator = try TokenIterator(input: input, model: context.model, processor: processor, sampler: sampler)
61+
let result = generate(input: input, context: context, iterator: iterator, didGenerate: didGenerate)
62+
let content = try Content(GeneratedContent(json: result.output))
63+
return (result, content)
64+
}
65+
66+
@available(macOS 26.0, iOS 26.0, *)
67+
public func generate<Content: Generable>(
68+
input: LMInput,
69+
parameters: GenerateParameters = GenerateParameters(),
70+
context: ModelContext,
71+
generating: Content.Type
72+
) async throws -> AsyncStream<Content.PartiallyGenerated> {
73+
let sampler = parameters.sampler()
74+
let grammar = try Grammar.generable(Content.self)
75+
let processor = try await GrammarMaskedLogitProcessor.from(configuration: context.configuration, grammar: grammar)
76+
let iterator = try TokenIterator(input: input, model: context.model, processor: processor, sampler: sampler)
77+
let stream = generate(input: input, context: context, iterator: iterator)
78+
return AsyncStream { continuation in
79+
80+
let task = Task {
81+
var output = ""
82+
for await generation in stream {
83+
if let chunk = generation.chunk {
84+
output.append(chunk)
85+
let generatedContent = try GeneratedContent(json: output)
86+
let partiallyGenerated = try Content.PartiallyGenerated(generatedContent)
87+
continuation.yield(partiallyGenerated)
88+
}
89+
}
90+
91+
continuation.finish()
92+
}
93+
94+
continuation.onTermination = { _ in
95+
task.cancel()
96+
}
97+
}
98+
}
99+
#endif

‎Sources/MLXStructured/Grammar/Grammar+Generable.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import FoundationModels
1414
#if compiler(>=6.2)
1515
@available(macOS 26.0, iOS 26.0, *)
1616
public extension Grammar {
17-
static func schema<Content: Generable>(generable type: Content.Type, indent: Int? = nil) throws -> Grammar {
17+
static func generable<Content: Generable>(_ type: Content.Type, indent: Int? = nil) throws -> Grammar {
1818
let encoder = JSONEncoder()
1919
let data = try encoder.encode(type.generationSchema)
2020
let string = String(decoding: data, as: UTF8.self)

‎Sources/MLXStructuredCLI/Commands/CodableExample.swift‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
//
2+
// CodableExample.swift
3+
// MLXStructured
4+
//
5+
// Created by Ivan Petrukha on 04.10.2025.
6+
//
7+
8+
import Foundation
9+
import ArgumentParser
10+
import JSONSchema
11+
import MLXStructured
12+
import MLXLMCommon
13+
14+
private struct MovieRecord: Codable {
15+
let title: String
16+
let year: Int
17+
let genres: [String]
18+
let director: String
19+
let actors: [String]
20+
}
21+
22+
private extension MovieRecord {
23+
24+
static let instruction = """
25+
Instruction: Extract movie record from the text according to schema: \(schema)
26+
"""
27+
28+
static let sample = """
29+
Text: The Dark Knight (2008) is a superhero crime film directed by Christopher Nolan. Starring Christian Bale, Heath Ledger, and Michael Caine.
30+
"""
31+
32+
static let schema = JSONSchema.object(
33+
description: "Movie record",
34+
properties: [
35+
"title": .string(),
36+
"year": .integer(minimum: 1900, maximum: 2026),
37+
"genres": .array(items: .string(), maxItems: 3),
38+
"director": .string(),
39+
"actors": .array(items: .string(), maxItems: 5)
40+
], required: [
41+
"title",
42+
"year",
43+
"genres",
44+
"director",
45+
"actors"
46+
]
47+
)
48+
}
49+
50+
struct CodableExample: AsyncParsableCommand {
51+
52+
static let configuration = CommandConfiguration(
53+
commandName: "codable",
54+
abstract: "Generate codable content according to JSON Schema."
55+
)
56+
57+
@OptionGroup
58+
var model: ModelArguments
59+
60+
func run() async throws {
61+
let context = try await model.modelContext()
62+
let prompt = MovieRecord.instruction + "\n" + MovieRecord.sample
63+
let input = try await context.processor.prepare(input: UserInput(prompt: prompt))
64+
let (result, model) = try await MLXStructured.generate(input: input, context: context, schema: MovieRecord.schema, generating: MovieRecord.self)
65+
print("Generation result:", result.output)
66+
print("Generated model:", model)
67+
}
68+
}

‎Sources/MLXStructuredCLI/Commands/GenerableExample.swift‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
//
2+
// GenerableExample.swift
3+
// MLXStructured
4+
//
5+
// Created by Ivan Petrukha on 04.10.2025.
6+
//
7+
8+
import Foundation
9+
import FoundationModels
10+
import ArgumentParser
11+
import MLXStructured
12+
import MLXLMCommon
13+
14+
@Generable
15+
@available(macOS 26.0, iOS 26.0, *)
16+
private struct MovieRecord: Codable {
17+
18+
@Guide(description: "Movie title")
19+
let title: String
20+
21+
@Guide(description: "Release year", .range(1900...2026))
22+
let year: Int
23+
24+
@Guide(description: "List of genres", .count(1...3))
25+
let genres: [String]
26+
27+
@Guide(description: "Director name")
28+
let director: String
29+
30+
@Guide(description: "List of principal actors", .count(1...5))
31+
let actors: [String]
32+
}
33+
34+
@available(macOS 26.0, iOS 26.0, *)
35+
private extension MovieRecord {
36+
37+
static let instruction = """
38+
Instruction: Extract movie record from the text according to schema: \(MovieRecord.generationSchema)
39+
"""
40+
41+
static let sample = """
42+
Text: The Dark Knight (2008) is a superhero crime film directed by Christopher Nolan. Starring Christian Bale, Heath Ledger, and Michael Caine.
43+
"""
44+
}
45+
46+
struct GenerableExample: AsyncParsableCommand {
47+
48+
static let configuration = CommandConfiguration(
49+
commandName: "generable",
50+
abstract: "Generate @Generable type."
51+
)
52+
53+
@OptionGroup
54+
var model: ModelArguments
55+
56+
func run() async throws {
57+
guard #available(macOS 26.0, iOS 26.0, *) else {
58+
fatalError("Generable examples available from macOS 26 only")
59+
}
60+
61+
let context = try await model.modelContext()
62+
let prompt = MovieRecord.instruction + "\n" + MovieRecord.sample
63+
let input = try await context.processor.prepare(input: UserInput(prompt: prompt))
64+
let (result, model) = try await MLXStructured.generate(input: input, context: context, generating: MovieRecord.self)
65+
print("Generation result:", result.output)
66+
print("Generated model:", model)
67+
}
68+
}
69+
70+
struct GenerableStreamExample: AsyncParsableCommand {
71+
72+
static let configuration = CommandConfiguration(
73+
commandName: "generable-stream",
74+
abstract: "Generate @Generable type using stream and partially generated content."
75+
)
76+
77+
@OptionGroup
78+
var model: ModelArguments
79+
80+
func run() async throws {
81+
guard #available(macOS 26.0, iOS 26.0, *) else {
82+
fatalError("Generable examples available from macOS 26 only")
83+
}
84+
85+
let context = try await model.modelContext()
86+
let prompt = MovieRecord.instruction + "\n" + MovieRecord.sample
87+
let input = try await context.processor.prepare(input: UserInput(prompt: prompt))
88+
let stream = try await MLXStructured.generate(input: input, context: context, generating: MovieRecord.self)
89+
for await content in stream {
90+
print("Partially generated:", content)
91+
}
92+
}
93+
}

‎Sources/MLXStructuredCLI/Commands/StructuralExample.swift‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
//
2+
// StructuralExample.swift
3+
// MLXStructured
4+
//
5+
// Created by Ivan Petrukha on 04.10.2025.
6+
//
7+
8+
import Foundation
9+
import ArgumentParser
10+
import MLXStructured
11+
import MLXLMCommon
12+
13+
struct StructuralExample: AsyncParsableCommand {
14+
15+
static let configuration = CommandConfiguration(
16+
commandName: "structural",
17+
abstract: "Generate text according to complex structural grammar."
18+
)
19+
20+
@OptionGroup
21+
var model: ModelArguments
22+
23+
func run() async throws {
24+
let context = try await model.modelContext()
25+
let includePrefix = Bool.random()
26+
let grammar = try Grammar {
27+
SequenceFormat {
28+
if includePrefix {
29+
ConstTextFormat(text: "According to my knowledge, my answer is ")
30+
}
31+
OrFormat {
32+
ConstTextFormat(text: "YES")
33+
ConstTextFormat(text: "NO")
34+
}
35+
}
36+
}
37+
let prompt = "Is it true that London is a capital of a Great Britain?"
38+
let input = try await context.processor.prepare(input: UserInput(prompt: prompt))
39+
let result = try await MLXStructured.generate(input: input, context: context, grammar: grammar)
40+
print("Generation result:", result.output)
41+
}
42+
}

0 commit comments

Comments
 (0)