SPB Git

spb/zyquo-mlx Public MIT

The local MLX foundry for your Mac — run, fine-tune, quantize, and ship models. Nothing leaves your machine.

macosswiftaimlxfine-tuning
Swift 93.4% Python 3.8% Makefile 2.2% Shell 0.5%

phase3: foundry core — PyBridge (PythonRunner, uv venv env, JSON-protocol scripts), DatasetService, TrainingService/RunStore/MetricsStream, ConversionService (Swift quant + Py fuse/convert), CLI drivers

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed 11 days ago (Jul 30, 2026) parent 358e09b

Showing 17 changed files with +1,666 and −4

modified Package.swift +6 −1
@@ -38,7 +38,12 @@ let package = Package(
38 38 .product(name: "HuggingFace", package: "swift-huggingface"),
39 39 .product(name: "Tokenizers", package: "swift-transformers"),
40 40 ],
41 − path: "Sources/ZyquoMLX"
41 + path: "Sources/ZyquoMLX",
42 + resources: [
43 + // Pinned Python helper scripts (PyBridge) — versioned with the
44 + // app, never fetched remotely (docs/BUILD.md §3.4).
45 + .copy("PyBridge/scripts")
46 + ]
42 47 )
43 48 ]
44 49 )
modified Sources/ZyquoMLX/App/CLI.swift +2 −2
@@ -126,9 +126,9 @@ enum CLI {
126 126 try await engine.unload()
127 127 }
128 128
129 − // MARK: - Arg parsing
129 + // MARK: - Arg parsing (shared with CLIFoundry)
130 130
131 − private static func value(after flag: String, in args: [String]) -> String? {
131 + static func value(after flag: String, in args: [String]) -> String? {
132 132 guard let index = args.firstIndex(of: flag), index + 1 < args.count else { return nil }
133 133 return args[index + 1]
134 134 }
added Sources/ZyquoMLX/App/CLIFoundry.swift +179 −0
@@ -0,0 +1,179 @@
1 +//
2 +// CLIFoundry.swift
3 +// Zyquo MLX
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Mail: contact@spboucher.ai
7 +//
8 +
9 +import Foundation
10 +
11 +/// CLI drivers for the Phase 3 foundry services (training, fusion, datasets).
12 +///
13 +/// ZyquoMLX --train --model <dir> --dataset <jsonl-or-dataset-dir> [--iters N] [--method lora|qlora|dora|full]
14 +/// ZyquoMLX --fuse --model <dir> --adapter <dir> --out <name> [--dequantize]
15 +/// ZyquoMLX --validate-dataset <file.jsonl>
16 +extension CLI {
17 +
18 + static var shouldRunFoundry: Bool {
19 + let args = CommandLine.arguments
20 + return args.contains("--train") || args.contains("--fuse")
21 + || args.contains("--validate-dataset")
22 + }
23 +
24 + static func runFoundry() async -> Int32 {
25 + do {
26 + let args = CommandLine.arguments
27 + if args.contains("--train") {
28 + try await train(args: args)
29 + return 0
30 + }
31 + if args.contains("--fuse") {
32 + try await fuse(args: args)
33 + return 0
34 + }
35 + if let path = value(after: "--validate-dataset", in: args) {
36 + try await validateDataset(path: path)
37 + return 0
38 + }
39 + return 2
40 + } catch {
41 + FileHandle.standardError.write(Data("error: \(error.localizedDescription)\n".utf8))
42 + return 1
43 + }
44 + }
45 +
46 + // MARK: - Train
47 +
48 + private static func train(args: [String]) async throws {
49 + guard let modelPath = value(after: "--model", in: args),
50 + let dataPath = value(after: "--dataset", in: args)
51 + else {
52 + throw CLIError.usage("--train requires --model <dir> and --dataset <jsonl|dir>")
53 + }
54 +
55 + let modelURL = URL(fileURLWithPath: (modelPath as NSString).expandingTildeInPath)
56 + let model = try await ModelStore.shared.describe(directory: modelURL)
57 +
58 + // Accept a raw JSONL (imported on the fly) or an existing dataset dir.
59 + let dataURL = URL(fileURLWithPath: (dataPath as NSString).expandingTildeInPath)
60 + let dataset: Dataset
61 + if dataURL.pathExtension == "jsonl" {
62 + let report = try await DatasetService.shared.importJSONL(
63 + from: dataURL, name: dataURL.deletingPathExtension().lastPathComponent)
64 + dataset = report.dataset
65 + print("dataset: \(dataset.name) [\(dataset.format.displayName)] "
66 + + "train=\(dataset.trainCount) valid=\(dataset.validCount)"
67 + + (report.issues.isEmpty ? "" : " (\(report.issues.count) rows skipped)"))
68 + for issue in report.issues.prefix(5) {
69 + print(" line \(issue.line): \(issue.problem) → \(issue.fix)")
70 + }
71 + } else {
72 + dataset = try PersistenceService.loadJSON(
73 + Dataset.self, from: dataURL.appendingPathComponent("dataset.json"))
74 + }
75 +
76 + var hp = HyperParams()
77 + hp.iterations = value(after: "--iters", in: args).flatMap(Int.init) ?? 100
78 + hp.batchSize = value(after: "--batch-size", in: args).flatMap(Int.init) ?? 2
79 + hp.saveEvery = value(after: "--save-every", in: args).flatMap(Int.init) ?? 50
80 + hp.stepsPerEval = value(after: "--steps-per-eval", in: args).flatMap(Int.init) ?? 50
81 + hp.stepsPerReport = 5
82 +
83 + let methodName = value(after: "--method", in: args) ?? (model.quantization != nil ? "qlora" : "lora")
84 + guard let method = FineTuneMethod(rawValue: methodName) else {
85 + throw CLIError.usage("unknown --method \(methodName)")
86 + }
87 +
88 + let service = TrainingService.shared
89 + let run = try await service.createRun(
90 + name: "cli-\(model.name)-\(method.rawValue)",
91 + baseModel: model, dataset: dataset, method: method, hyperParams: hp)
92 + print("run: \(run.id.uuidString)")
93 + print("method: \(method.displayName) on \(model.name)")
94 + print("verdict: \(MemoryAdvisor.trainingVerdict(for: model, method: method, params: hp).displayName)\n")
95 +
96 + let events = try await service.start(run: run, baseModel: model, dataset: dataset)
97 + for await event in events {
98 + switch event {
99 + case .started(let model, let iterations):
100 + print("training \(model) for \(iterations) iterations…")
101 + case .metric(let m):
102 + if let loss = m.trainLoss {
103 + let speed = m.tokensPerSecond.map { String(format: "%.0f tok/s", $0) } ?? ""
104 + let mem = m.peakMemoryGB.map { String(format: "peak %.1f GB", $0) } ?? ""
105 + print("iter \(m.iteration): train loss \(String(format: "%.3f", loss)) \(speed) \(mem)")
106 + }
107 + if let loss = m.valLoss {
108 + print("iter \(m.iteration): VAL loss \(String(format: "%.3f", loss))")
109 + }
110 + case .checkpointSaved(let name, _):
111 + print("checkpoint saved: \(name)")
112 + case .finished:
113 + print("\ntraining finished ✅")
114 + case .failed(let message):
115 + print("\ntraining FAILED: \(message)")
116 + }
117 + }
118 +
119 + let finished = try await RunStore.shared.load(id: run.id)
120 + print("state: \(finished.state.rawValue), adapters: \(await RunStore.shared.latestAdapter(for: finished)?.path ?? "none")")
121 + }
122 +
123 + // MARK: - Fuse
124 +
125 + private static func fuse(args: [String]) async throws {
126 + guard let modelPath = value(after: "--model", in: args),
127 + let adapterPath = value(after: "--adapter", in: args),
128 + let outName = value(after: "--out", in: args)
129 + else {
130 + throw CLIError.usage("--fuse requires --model <dir>, --adapter <dir>, --out <name>")
131 + }
132 + let model = try await ModelStore.shared.describe(
133 + directory: URL(fileURLWithPath: (modelPath as NSString).expandingTildeInPath))
134 + let adapter = URL(fileURLWithPath: (adapterPath as NSString).expandingTildeInPath)
135 +
136 + let events = try await ConversionService.shared.fuse(
137 + baseModel: model, adapterDirectory: adapter, outputName: outName,
138 + dequantize: args.contains("--dequantize"))
139 + for await event in events {
140 + switch event {
141 + case .stage(let stage, _):
142 + print("stage: \(stage)")
143 + case .finished(let output):
144 + print("fused model → \(output.path) ✅")
145 + case .failed(let message):
146 + throw CLIError.usage("fuse failed: \(message)")
147 + }
148 + }
149 + }
150 +
151 + // MARK: - Dataset validation
152 +
153 + private static func validateDataset(path: String) async throws {
154 + let url = URL(fileURLWithPath: (path as NSString).expandingTildeInPath)
155 + let report = try await DatasetService.shared.importJSONL(
156 + from: url, name: url.deletingPathExtension().lastPathComponent + "-validated")
157 + let d = report.dataset
158 + print("format: \(d.format.displayName)")
159 + print("valid rows: \(d.sampleCount) (train \(d.trainCount) / valid \(d.validCount) / test \(d.testCount))")
160 + print("est. tokens: \(d.totalTokens ?? 0) total, longest sample ≈ \(d.maxSequenceTokens ?? 0)")
161 + if report.issues.isEmpty {
162 + print("issues: none ✅")
163 + } else {
164 + print("issues: \(report.issues.count) rows skipped")
165 + for issue in report.issues.prefix(10) {
166 + print(" line \(issue.line): \(issue.problem) → \(issue.fix)")
167 + }
168 + }
169 + print("written to: \(d.directory.path)")
170 + }
171 +}
172 +
173 +enum CLIError: LocalizedError {
174 + case usage(String)
175 + var errorDescription: String? {
176 + if case .usage(let message) = self { return message }
177 + return nil
178 + }
179 +}
modified Sources/ZyquoMLX/App/Main.swift +3 −0
@@ -16,6 +16,9 @@ enum Main {
16 16 if CLI.shouldRun {
17 17 exit(await CLI.run())
18 18 }
19 + if CLI.shouldRunFoundry {
20 + exit(await CLI.runFoundry())
21 + }
19 22 ZyquoMLXApp.main()
20 23 }
21 24 }
added Sources/ZyquoMLX/Convert/ConversionService.swift +161 −0
@@ -0,0 +1,161 @@
1 +//
2 +// ConversionService.swift
3 +// Zyquo MLX
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Mail: contact@spboucher.ai
7 +//
8 +
9 +import Foundation
10 +import MLX
11 +import MLXLLM
12 +import MLXLMCommon
13 +
14 +/// Progress for conversion/quantization/fusion jobs.
15 +enum ConversionEvent: Sendable {
16 + case stage(String, fraction: Double?)
17 + case finished(outputDirectory: URL)
18 + case failed(message: String)
19 +}
20 +
21 +enum ConversionServiceError: LocalizedError {
22 + case alreadyQuantized
23 + case outputExists(URL)
24 +
25 + var errorDescription: String? {
26 + switch self {
27 + case .alreadyQuantized:
28 + "This model is already quantized — de-quantize first if you want a different recipe."
29 + case .outputExists(let url):
30 + "Output folder already exists: \(url.lastPathComponent). Choose another name or delete it."
31 + }
32 + }
33 +}
34 +
35 +/// Conversion, quantization, and adapter fusion (charter 3.C).
36 +///
37 +/// Execution split per docs/MLX-RESEARCH.md §2/§6:
38 +/// - Quantizing a local safetensors LLM → **native Swift**
39 +/// (`LLMModelFactory.convert`, affine/mxfp4).
40 +/// - Fusing adapters to disk, HF→MLX conversion of arbitrary repos
41 +/// (incl. `.bin`), de-quantize → **Python bridge** (`mlx_lm` pinned).
42 +actor ConversionService {
43 +
44 + static let shared = ConversionService()
45 +
46 + // MARK: - Swift-native quantization
47 +
48 + /// Quantize a local, unquantized safetensors LLM into a new model
49 + /// directory in the library.
50 + func quantize(model: LocalModel, config: QuantConfig, outputName: String)
51 + -> AsyncStream<ConversionEvent>
52 + {
53 + let output = PersistenceService.modelsDirectory
54 + .appendingPathComponent(outputName, isDirectory: true)
55 +
56 + let (stream, continuation) = AsyncStream.makeStream(of: ConversionEvent.self)
57 + Task {
58 + do {
59 + guard model.quantization == nil else { throw ConversionServiceError.alreadyQuantized }
60 + guard !FileManager.default.fileExists(atPath: output.path) else {
61 + throw ConversionServiceError.outputExists(output)
62 + }
63 + let options = ModelConversionOptions(
64 + bits: config.bits,
65 + groupSize: config.groupSize,
66 + mode: config.mode == .mxfp4 ? .mxfp4 : .affine
67 + )
68 + _ = try await LLMModelFactory.shared.convert(
69 + from: model.directory,
70 + to: output,
71 + options: options,
72 + progressHandler: { progress in
73 + continuation.yield(
74 + .stage(progress.stage.rawValue, fraction: progress.fractionCompleted))
75 + }
76 + )
77 + Memory.clearCache()
78 + continuation.yield(.finished(outputDirectory: output))
79 + } catch {
80 + continuation.yield(.failed(message: error.localizedDescription))
81 + }
82 + continuation.finish()
83 + }
84 + return stream
85 + }
86 +
87 + // MARK: - Python-bridge pipelines
88 +
89 + /// Fuse a trained adapter into its base model, producing a standalone
90 + /// MLX model directory in the library.
91 + func fuse(baseModel: LocalModel, adapterDirectory: URL, outputName: String, dequantize: Bool)
92 + async throws -> AsyncStream<ConversionEvent>
93 + {
94 + let output = PersistenceService.modelsDirectory
95 + .appendingPathComponent(outputName, isDirectory: true)
96 + guard !FileManager.default.fileExists(atPath: output.path) else {
97 + throw ConversionServiceError.outputExists(output)
98 + }
99 +
100 + var arguments = [
101 + "--model", baseModel.directory.path,
102 + "--adapter-path", adapterDirectory.path,
103 + "--save-path", output.path,
104 + ]
105 + if dequantize { arguments.append("--dequantize") }
106 +
107 + return try await pythonJob(script: "zyquo_fuse", arguments: arguments, output: output)
108 + }
109 +
110 + /// Convert an HF repo/directory (any weights format mlx-lm handles) to
111 + /// MLX, optionally quantized.
112 + func convert(hfPath: String, outputName: String, quantize: QuantConfig?)
113 + async throws -> AsyncStream<ConversionEvent>
114 + {
115 + let output = PersistenceService.modelsDirectory
116 + .appendingPathComponent(outputName, isDirectory: true)
117 + guard !FileManager.default.fileExists(atPath: output.path) else {
118 + throw ConversionServiceError.outputExists(output)
119 + }
120 +
121 + var arguments = ["--hf-path", hfPath, "--mlx-path", output.path]
122 + if let quantize {
123 + arguments += [
124 + "--quantize",
125 + "--q-bits", String(quantize.bits),
126 + "--q-group-size", String(quantize.groupSize),
127 + "--q-mode", quantize.mode.rawValue,
128 + ]
129 + }
130 +
131 + return try await pythonJob(script: "zyquo_convert", arguments: arguments, output: output)
132 + }
133 +
134 + private func pythonJob(script: String, arguments: [String], output: URL)
135 + async throws -> AsyncStream<ConversionEvent>
136 + {
137 + let pythonStream = try await PythonRunner.shared.stream(
138 + script: script, arguments: arguments)
139 + let (stream, continuation) = AsyncStream.makeStream(of: ConversionEvent.self)
140 + Task {
141 + do {
142 + for try await event in pythonStream {
143 + switch event.event {
144 + case "start":
145 + continuation.yield(.stage(event.string("stage") ?? "working", fraction: nil))
146 + case "done":
147 + continuation.yield(.finished(outputDirectory: output))
148 + case "error":
149 + continuation.yield(.failed(message: event.string("message") ?? "unknown"))
150 + default:
151 + break
152 + }
153 + }
154 + } catch {
155 + continuation.yield(.failed(message: error.localizedDescription))
156 + }
157 + continuation.finish()
158 + }
159 + return stream
160 + }
161 +}
added Sources/ZyquoMLX/Convert/QuantConfig.swift +41 −0
@@ -0,0 +1,41 @@
1 +//
2 +// QuantConfig.swift
3 +// Zyquo MLX
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Mail: contact@spboucher.ai
7 +//
8 +
9 +import Foundation
10 +
11 +/// Quantization recipe (docs/MLX-RESEARCH.md §5: affine supports bits
12 +/// 2/3/4/5/6/8 with group sizes 32/64/128; mxfp4/nvfp4/mxfp8 have fixed
13 +/// bits/groups).
14 +struct QuantConfig: Codable, Hashable, Sendable {
15 + enum Mode: String, Codable, CaseIterable, Sendable {
16 + case affine
17 + case mxfp4
18 + case nvfp4
19 + case mxfp8
20 + }
21 +
22 + var bits: Int = 4
23 + var groupSize: Int = 64
24 + var mode: Mode = .affine
25 +
26 + static let affineBits = [2, 3, 4, 5, 6, 8]
27 + static let affineGroupSizes = [32, 64, 128]
28 +
29 + /// Effective bytes/param for size previews (group-scale overhead included).
30 + var bytesPerParameter: Double {
31 + (Double(bits) + 32.0 / Double(groupSize)) / 8.0
32 + }
33 +
34 + /// Predicted weight size after quantizing a model with the given
35 + /// parameter count (before/after preview in the Convert UI).
36 + func predictedWeightBytes(parameterCount: Int64) -> Int64 {
37 + Int64(Double(parameterCount) * bytesPerParameter)
38 + }
39 +
40 + var label: String { "\(bits)-bit \(mode.rawValue) (g\(groupSize))" }
41 +}
added Sources/ZyquoMLX/Data/DatasetFormats.swift +121 −0
@@ -0,0 +1,121 @@
1 +//
2 +// DatasetFormats.swift
3 +// Zyquo MLX
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Mail: contact@spboucher.ai
7 +//
8 +
9 +import Foundation
10 +
11 +/// A problem found in one dataset row, with a concrete fix
12 +/// (charter 3.A: report malformed rows with fixes).
13 +struct DatasetRowIssue: Identifiable, Codable, Hashable, Sendable {
14 + var id: Int { line }
15 + /// 1-based line number in the source file.
16 + var line: Int
17 + var problem: String
18 + var fix: String
19 +}
20 +
21 +/// One parsed sample, normalized for preview.
22 +struct DatasetSample: Sendable {
23 + var messages: [(role: String, content: String)]
24 + var estimatedTokens: Int
25 +}
26 +
27 +/// Format detection + row validation for the JSONL shapes mlx-lm accepts
28 +/// (docs/TRAINING-RESEARCH.md §4.2 — detection order: completions, chat, text).
29 +enum DatasetFormats {
30 +
31 + static func detect(firstRow row: [String: Any]) -> DatasetFormat? {
32 + if row["prompt"] != nil && row["completion"] != nil { return .completions }
33 + if row["messages"] != nil { return .chat }
34 + if row["text"] != nil { return .text }
35 + return nil
36 + }
37 +
38 + /// Validate one row against a format. Returns an issue or nil when valid.
39 + static func validate(row: [String: Any], format: DatasetFormat, line: Int) -> DatasetRowIssue? {
40 + switch format {
41 + case .chat:
42 + guard let messages = row["messages"] as? [[String: Any]] else {
43 + return DatasetRowIssue(
44 + line: line,
45 + problem: "Missing or non-array \"messages\" key.",
46 + fix: "Use {\"messages\": [{\"role\": \"user\", \"content\": \"…\"}, …]}.")
47 + }
48 + if messages.isEmpty {
49 + return DatasetRowIssue(
50 + line: line, problem: "\"messages\" is empty.",
51 + fix: "Provide at least a user and an assistant message.")
52 + }
53 + let validRoles: Set<String> = ["system", "user", "assistant", "tool"]
54 + for (i, message) in messages.enumerated() {
55 + guard let role = message["role"] as? String, validRoles.contains(role) else {
56 + return DatasetRowIssue(
57 + line: line,
58 + problem: "Message \(i + 1) has a missing/invalid \"role\".",
59 + fix: "Use one of: system, user, assistant, tool.")
60 + }
61 + guard message["content"] is String else {
62 + return DatasetRowIssue(
63 + line: line,
64 + problem: "Message \(i + 1) (\(role)) has no string \"content\".",
65 + fix: "Every message needs a \"content\" string.")
66 + }
67 + }
68 + if (messages.last?["role"] as? String) != "assistant" {
69 + return DatasetRowIssue(
70 + line: line,
71 + problem: "Conversation does not end with an assistant message.",
72 + fix: "The final message must be the assistant reply the model should learn.")
73 + }
74 + return nil
75 +
76 + case .completions:
77 + guard row["prompt"] is String else {
78 + return DatasetRowIssue(
79 + line: line, problem: "Missing string \"prompt\".",
80 + fix: "Use {\"prompt\": \"…\", \"completion\": \"…\"}.")
81 + }
82 + guard row["completion"] is String else {
83 + return DatasetRowIssue(
84 + line: line, problem: "Missing string \"completion\".",
85 + fix: "Use {\"prompt\": \"…\", \"completion\": \"…\"}.")
86 + }
87 + return nil
88 +
89 + case .text:
90 + guard let text = row["text"] as? String, !text.isEmpty else {
91 + return DatasetRowIssue(
92 + line: line, problem: "Missing or empty string \"text\".",
93 + fix: "Use {\"text\": \"…\"} with non-empty content.")
94 + }
95 + return nil
96 + }
97 + }
98 +
99 + /// Normalize a valid row into a preview sample.
100 + static func sample(from row: [String: Any], format: DatasetFormat) -> DatasetSample {
101 + var messages: [(String, String)] = []
102 + switch format {
103 + case .chat:
104 + for message in row["messages"] as? [[String: Any]] ?? [] {
105 + messages.append(
106 + (message["role"] as? String ?? "?", message["content"] as? String ?? ""))
107 + }
108 + case .completions:
109 + messages = [
110 + ("user", row["prompt"] as? String ?? ""),
111 + ("assistant", row["completion"] as? String ?? ""),
112 + ]
113 + case .text:
114 + messages = [("text", row["text"] as? String ?? "")]
115 + }
116 + let characters = messages.reduce(0) { $0 + $1.1.count }
117 + // ~4 chars/token heuristic; precise counts come from the tokenizer at
118 + // training time (trainer truncates over max-seq-length with a warning).
119 + return DatasetSample(messages: messages, estimatedTokens: max(1, characters / 4))
120 + }
121 +}
added Sources/ZyquoMLX/Data/DatasetService.swift +197 −0
@@ -0,0 +1,197 @@
1 +//
2 +// DatasetService.swift
3 +// Zyquo MLX
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Mail: contact@spboucher.ai
7 +//
8 +
9 +import Foundation
10 +
11 +enum DatasetImportError: LocalizedError {
12 + case emptyFile
13 + case undetectableFormat
14 + case allRowsInvalid(issues: [DatasetRowIssue])
15 + case tooFewSamples(count: Int)
16 +
17 + var errorDescription: String? {
18 + switch self {
19 + case .emptyFile:
20 + "The file contains no data rows."
21 + case .undetectableFormat:
22 + "Could not detect the dataset format. Rows must be JSON objects with \"messages\", \"prompt\"+\"completion\", or \"text\" keys."
23 + case .allRowsInvalid(let issues):
24 + "No valid rows found. First problem (line \(issues.first?.line ?? 1)): \(issues.first?.problem ?? "")"
25 + case .tooFewSamples(let count):
26 + "Only \(count) valid samples — at least 8 are needed for a meaningful train/valid split."
27 + }
28 + }
29 +}
30 +
31 +/// Result of a dataset import: what was written plus the validation report.
32 +struct DatasetImportReport: Sendable {
33 + var dataset: Dataset
34 + var issues: [DatasetRowIssue]
35 + var skippedRows: Int
36 +}
37 +
38 +/// Imports, validates, splits, and previews training datasets
39 +/// (charter 3.A). mlx-lm does not auto-split local JSONL
40 +/// (docs/TRAINING-RESEARCH.md §4.5) — the split is ours.
41 +actor DatasetService {
42 +
43 + static let shared = DatasetService()
44 +
45 + private let root: URL
46 +
47 + init(root: URL = PersistenceService.datasetsDirectory) {
48 + self.root = root
49 + }
50 +
51 + // MARK: - Import
52 +
53 + /// Import a JSONL file: validate every row, split into
54 + /// train/valid(/test), and write the dataset directory.
55 + func importJSONL(
56 + from sourceURL: URL,
57 + name: String,
58 + validFraction: Double = 0.1,
59 + testFraction: Double = 0.0,
60 + seed: UInt64 = 0
61 + ) throws -> DatasetImportReport {
62 + let content = try String(contentsOf: sourceURL, encoding: .utf8)
63 + let lines = content.split(separator: "\n", omittingEmptySubsequences: true)
64 + guard !lines.isEmpty else { throw DatasetImportError.emptyFile }
65 +
66 + var format: DatasetFormat?
67 + var validRows: [String] = []
68 + var issues: [DatasetRowIssue] = []
69 + var totalTokens = 0
70 + var maxTokens = 0
71 +
72 + for (index, line) in lines.enumerated() {
73 + let lineNumber = index + 1
74 + guard
75 + let data = line.data(using: .utf8),
76 + let row = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any]
77 + else {
78 + issues.append(
79 + DatasetRowIssue(
80 + line: lineNumber, problem: "Not a valid JSON object.",
81 + fix: "Each line must be one complete JSON object (no trailing commas, no multi-line objects)."))
82 + continue
83 + }
84 +
85 + if format == nil {
86 + format = DatasetFormats.detect(firstRow: row)
87 + guard format != nil else { throw DatasetImportError.undetectableFormat }
88 + }
89 +
90 + if let issue = DatasetFormats.validate(row: row, format: format!, line: lineNumber) {
91 + issues.append(issue)
92 + continue
93 + }
94 +
95 + let sample = DatasetFormats.sample(from: row, format: format!)
96 + totalTokens += sample.estimatedTokens
97 + maxTokens = max(maxTokens, sample.estimatedTokens)
98 + validRows.append(String(line))
99 + }
100 +
101 + guard let detectedFormat = format else { throw DatasetImportError.undetectableFormat }
102 + guard !validRows.isEmpty else { throw DatasetImportError.allRowsInvalid(issues: issues) }
103 + guard validRows.count >= 8 else { throw DatasetImportError.tooFewSamples(count: validRows.count) }
104 +
105 + // Deterministic shuffle then split.
106 + var generator = SeededGenerator(seed: seed)
107 + validRows.shuffle(using: &generator)
108 +
109 + let testCount = Int(Double(validRows.count) * testFraction)
110 + let validCount = max(1, Int(Double(validRows.count) * validFraction))
111 + let trainCount = validRows.count - validCount - testCount
112 +
113 + let directory = root.appendingPathComponent(sanitize(name), isDirectory: true)
114 + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
115 +
116 + let train = validRows[0..<trainCount]
117 + let valid = validRows[trainCount..<(trainCount + validCount)]
118 + let test = validRows[(trainCount + validCount)...]
119 +
120 + try write(rows: Array(train), to: directory.appendingPathComponent("train.jsonl"))
121 + try write(rows: Array(valid), to: directory.appendingPathComponent("valid.jsonl"))
122 + if !test.isEmpty {
123 + try write(rows: Array(test), to: directory.appendingPathComponent("test.jsonl"))
124 + }
125 +
126 + let dataset = Dataset(
127 + id: UUID(),
128 + name: name,
129 + directory: directory,
130 + format: detectedFormat,
131 + trainCount: trainCount,
132 + validCount: validCount,
133 + testCount: test.count,
134 + totalTokens: totalTokens,
135 + maxSequenceTokens: maxTokens,
136 + importedAt: .now
137 + )
138 + try PersistenceService.saveJSON(dataset, to: directory.appendingPathComponent("dataset.json"))
139 +
140 + return DatasetImportReport(
141 + dataset: dataset, issues: issues, skippedRows: issues.count)
142 + }
143 +
144 + // MARK: - Library
145 +
146 + func scan() throws -> [Dataset] {
147 + let fm = FileManager.default
148 + guard fm.fileExists(atPath: root.path) else { return [] }
149 + let entries = try fm.contentsOfDirectory(
150 + at: root, includingPropertiesForKeys: nil, options: .skipsHiddenFiles)
151 + return entries.compactMap {
152 + try? PersistenceService.loadJSON(
153 + Dataset.self, from: $0.appendingPathComponent("dataset.json"))
154 + }
155 + .sorted { $0.importedAt > $1.importedAt }
156 + }
157 +
158 + /// First N samples of a split, parsed for preview.
159 + func preview(dataset: Dataset, split: String = "train", count: Int = 8) throws -> [DatasetSample] {
160 + let url = dataset.directory.appendingPathComponent("\(split).jsonl")
161 + let content = try String(contentsOf: url, encoding: .utf8)
162 + return content.split(separator: "\n").prefix(count).compactMap { line in
163 + guard
164 + let data = line.data(using: .utf8),
165 + let row = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any]
166 + else { return nil }
167 + return DatasetFormats.sample(from: row, format: dataset.format)
168 + }
169 + }
170 +
171 + func delete(_ dataset: Dataset) throws {
172 + try FileManager.default.removeItem(at: dataset.directory)
173 + }
174 +
175 + // MARK: - Helpers
176 +
177 + private func write(rows: [String], to url: URL) throws {
178 + try (rows.joined(separator: "\n") + "\n").write(to: url, atomically: true, encoding: .utf8)
179 + }
180 +
181 + private func sanitize(_ name: String) -> String {
182 + name.replacingOccurrences(of: "[^A-Za-z0-9._-]+", with: "-", options: .regularExpression)
183 + }
184 +}
185 +
186 +/// Deterministic RNG for reproducible splits (SplitMix64).
187 +struct SeededGenerator: RandomNumberGenerator {
188 + private var state: UInt64
189 + init(seed: UInt64) { state = seed &+ 0x9E37_79B9_7F4A_7C15 }
190 + mutating func next() -> UInt64 {
191 + state &+= 0x9E37_79B9_7F4A_7C15
192 + var z = state
193 + z = (z ^ (z >> 30)) &* 0xBF58_476D_1CE4_E5B9
194 + z = (z ^ (z >> 27)) &* 0x94D0_49BB_1331_11EB
195 + return z ^ (z >> 31)
196 + }
197 +}
added Sources/ZyquoMLX/PyBridge/PythonEnvironment.swift +121 −0
@@ -0,0 +1,121 @@
1 +//
2 +// PythonEnvironment.swift
3 +// Zyquo MLX
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Mail: contact@spboucher.ai
7 +//
8 +
9 +import Foundation
10 +
11 +/// Manages the isolated, pinned Python environment in
12 +/// `~/Library/Application Support/ZyquoMLX/py/venv` (docs/BUILD.md §3.4).
13 +/// Provisioned with `uv` on first use; never touches system Python; lives
14 +/// outside the signed bundle (notarization posture, docs/BUILD.md §4).
15 +actor PythonEnvironment {
16 +
17 + static let shared = PythonEnvironment()
18 +
19 + /// Exact pins per docs/MLX-RESEARCH.md §6 (core set; mlx-vlm/whisper/audio
20 + /// install on demand when those capabilities are first used).
21 + static let pythonVersion = "3.12"
22 + static let corePins = ["mlx-lm==0.31.3", "pyyaml==6.0.3"]
23 +
24 + enum EnvironmentError: LocalizedError {
25 + case uvNotFound
26 + case provisioningFailed(String)
27 +
28 + var errorDescription: String? {
29 + switch self {
30 + case .uvNotFound:
31 + "The 'uv' tool is required to set up the Python environment. Install it with: brew install uv"
32 + case .provisioningFailed(let detail):
33 + "Python environment setup failed: \(detail)"
34 + }
35 + }
36 + }
37 +
38 + var venvURL: URL {
39 + PersistenceService.pythonDirectory.appendingPathComponent("venv", isDirectory: true)
40 + }
41 +
42 + var pythonURL: URL {
43 + venvURL.appendingPathComponent("bin/python")
44 + }
45 +
46 + private var markerURL: URL {
47 + PersistenceService.pythonDirectory.appendingPathComponent("provisioned.json")
48 + }
49 +
50 + var isProvisioned: Bool {
51 + FileManager.default.fileExists(atPath: pythonURL.path)
52 + && FileManager.default.fileExists(atPath: markerURL.path)
53 + }
54 +
55 + /// Locate `uv` in the usual installation paths (the app has no shell PATH).
56 + static func findUV() -> URL? {
57 + let candidates = [
58 + "\(NSHomeDirectory())/.local/bin/uv",
59 + "/opt/homebrew/bin/uv",
60 + "/usr/local/bin/uv",
61 + ]
62 + return candidates.first { FileManager.default.fileExists(atPath: $0) }
63 + .map { URL(fileURLWithPath: $0) }
64 + }
65 +
66 + /// Create the venv and install core pins. Idempotent.
67 + func provision(progress: @Sendable (String) -> Void = { _ in }) async throws {
68 + if isProvisioned { return }
69 + guard let uv = Self.findUV() else { throw EnvironmentError.uvNotFound }
70 +
71 + try FileManager.default.createDirectory(
72 + at: PersistenceService.pythonDirectory, withIntermediateDirectories: true)
73 +
74 + progress("Installing Python \(Self.pythonVersion)…")
75 + try runUV(uv, ["python", "install", Self.pythonVersion])
76 +
77 + progress("Creating environment…")
78 + try runUV(uv, ["venv", venvURL.path, "--python", Self.pythonVersion])
79 +
80 + progress("Installing MLX packages…")
81 + try runUV(uv, ["pip", "install", "--python", pythonURL.path] + Self.corePins)
82 +
83 + let marker: [String: String] = [
84 + "python": Self.pythonVersion,
85 + "pins": Self.corePins.joined(separator: " "),
86 + "provisionedAt": ISO8601DateFormatter().string(from: .now),
87 + ]
88 + try PersistenceService.saveJSON(marker, to: markerURL)
89 + progress("Python environment ready.")
90 + }
91 +
92 + /// Install additional pinned packages (e.g. mlx-vlm when VLM Python
93 + /// pipelines are first needed).
94 + func install(pins: [String]) throws {
95 + guard let uv = Self.findUV() else { throw EnvironmentError.uvNotFound }
96 + try runUV(uv, ["pip", "install", "--python", pythonURL.path] + pins)
97 + }
98 +
99 + /// Delete and re-provision (Settings › Python Environment › Repair).
100 + func repair(progress: @Sendable (String) -> Void = { _ in }) async throws {
101 + try? FileManager.default.removeItem(at: venvURL)
102 + try? FileManager.default.removeItem(at: markerURL)
103 + try await provision(progress: progress)
104 + }
105 +
106 + private func runUV(_ uv: URL, _ arguments: [String]) throws {
107 + let process = Process()
108 + process.executableURL = uv
109 + process.arguments = arguments
110 + let errPipe = Pipe()
111 + process.standardError = errPipe
112 + process.standardOutput = Pipe()
113 + try process.run()
114 + process.waitUntilExit()
115 + if process.terminationStatus != 0 {
116 + let data = errPipe.fileHandleForReading.readDataToEndOfFile()
117 + throw EnvironmentError.provisioningFailed(
118 + String(data: data, encoding: .utf8) ?? "exit \(process.terminationStatus)")
119 + }
120 + }
121 +}
added Sources/ZyquoMLX/PyBridge/PythonRunner.swift +188 −0
@@ -0,0 +1,188 @@
1 +//
2 +// PythonRunner.swift
3 +// Zyquo MLX
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Mail: contact@spboucher.ai
7 +//
8 +
9 +import Foundation
10 +
11 +/// One JSON-lines event from a PyBridge helper script.
12 +/// All scripts emit `{"event": "...", ...}` objects, one per line
13 +/// (protocol defined in `PyBridge/scripts/` — docs/TRAINING-RESEARCH.md §5).
14 +struct PythonEvent: @unchecked Sendable {
15 + let event: String
16 + /// JSONSerialization plist values (NSString/NSNumber/…) — value types in
17 + /// practice, hence the unchecked-Sendable annotation.
18 + let payload: [String: Any]
19 +
20 + func double(_ key: String) -> Double? {
21 + (payload[key] as? Double) ?? (payload[key] as? Int).map(Double.init)
22 + }
23 + func int(_ key: String) -> Int? {
24 + (payload[key] as? Int) ?? (payload[key] as? Double).map(Int.init)
25 + }
26 + func string(_ key: String) -> String? { payload[key] as? String }
27 +}
28 +
29 +enum PythonRunnerError: LocalizedError {
30 + case scriptNotFound(String)
31 + case environmentNotProvisioned
32 + case processFailed(exitCode: Int32, stderr: String)
33 +
34 + var errorDescription: String? {
35 + switch self {
36 + case .scriptNotFound(let name):
37 + "Bundled Python script missing: \(name)"
38 + case .environmentNotProvisioned:
39 + "The Python environment is not set up yet. It is provisioned automatically on first use."
40 + case .processFailed(let code, let stderr):
41 + "Python pipeline failed (exit \(code)). \(stderr.suffix(500))"
42 + }
43 + }
44 +}
45 +
46 +/// All Python interaction goes through here (charter engineering standard):
47 +/// spawns the pinned venv's interpreter on a bundled script and streams the
48 +/// JSON-lines protocol back. Cancellation terminates the process.
49 +actor PythonRunner {
50 +
51 + static let shared = PythonRunner()
52 +
53 + /// Resolve a bundled helper script by name.
54 + static func scriptURL(named name: String) throws -> URL {
55 + guard
56 + let url = Bundle.module.url(
57 + forResource: name, withExtension: "py", subdirectory: "scripts")
58 + else {
59 + throw PythonRunnerError.scriptNotFound("\(name).py")
60 + }
61 + return url
62 + }
63 +
64 + /// Run a script and stream its events. The stream throws on nonzero exit
65 + /// (unless an `error` event already surfaced the failure) and finishes
66 + /// after the process exits. Terminating the stream kills the process.
67 + func stream(script: String, arguments: [String]) async throws
68 + -> AsyncThrowingStream<PythonEvent, Error>
69 + {
70 + let env = PythonEnvironment.shared
71 + if await !env.isProvisioned {
72 + try await env.provision()
73 + }
74 + let python = await env.pythonURL
75 + let scriptURL = try Self.scriptURL(named: script)
76 +
77 + let process = Process()
78 + process.executableURL = python
79 + process.arguments = ["-u", scriptURL.path] + arguments
80 + // Isolated, reproducible interpreter environment.
81 + process.environment = [
82 + "PATH": "/usr/bin:/bin",
83 + "HOME": NSHomeDirectory(),
84 + "PYTHONUNBUFFERED": "1",
85 + ]
86 +
87 + let stdout = Pipe()
88 + let stderr = Pipe()
89 + process.standardOutput = stdout
90 + process.standardError = stderr
91 +
92 + return AsyncThrowingStream { continuation in
93 + let stderrBuffer = LockedBuffer()
94 +
95 + stderr.fileHandleForReading.readabilityHandler = { handle in
96 + stderrBuffer.append(handle.availableData)
97 + }
98 +
99 + let lineBuffer = LockedBuffer()
100 + stdout.fileHandleForReading.readabilityHandler = { handle in
101 + let data = handle.availableData
102 + guard !data.isEmpty else { return }
103 + for line in lineBuffer.appendAndExtractLines(data) {
104 + if let event = Self.decode(line: line) {
105 + continuation.yield(event)
106 + }
107 + }
108 + }
109 +
110 + process.terminationHandler = { proc in
111 + stdout.fileHandleForReading.readabilityHandler = nil
112 + stderr.fileHandleForReading.readabilityHandler = nil
113 + // Flush any trailing output.
114 + let rest = stdout.fileHandleForReading.readDataToEndOfFile()
115 + for line in lineBuffer.appendAndExtractLines(rest, flush: true) {
116 + if let event = Self.decode(line: line) {
117 + continuation.yield(event)
118 + }
119 + }
120 + stderrBuffer.append(stderr.fileHandleForReading.readDataToEndOfFile())
121 + if proc.terminationStatus == 0 {
122 + continuation.finish()
123 + } else {
124 + continuation.finish(
125 + throwing: PythonRunnerError.processFailed(
126 + exitCode: proc.terminationStatus,
127 + stderr: stderrBuffer.string))
128 + }
129 + }
130 +
131 + continuation.onTermination = { termination in
132 + if case .cancelled = termination, process.isRunning {
133 + process.terminate()
134 + }
135 + }
136 +
137 + do {
138 + try process.run()
139 + } catch {
140 + continuation.finish(throwing: error)
141 + }
142 + }
143 + }
144 +
145 + private static func decode(line: Data) -> PythonEvent? {
146 + guard
147 + let object = try? JSONSerialization.jsonObject(with: line) as? [String: Any],
148 + let event = object["event"] as? String
149 + else { return nil }
150 + return PythonEvent(event: event, payload: object)
151 + }
152 +}
153 +
154 +/// Thread-safe byte buffer for pipe readability handlers (they run on a
155 +/// private DispatchQueue while the stream lives elsewhere).
156 +private final class LockedBuffer: @unchecked Sendable {
157 + private let lock = NSLock()
158 + private var data = Data()
159 +
160 + func append(_ chunk: Data) {
161 + lock.lock()
162 + defer { lock.unlock() }
163 + data.append(chunk)
164 + }
165 +
166 + /// Append then split out complete `\n`-terminated lines (keeps the tail).
167 + func appendAndExtractLines(_ chunk: Data, flush: Bool = false) -> [Data] {
168 + lock.lock()
169 + defer { lock.unlock() }
170 + data.append(chunk)
171 + var lines: [Data] = []
172 + while let newline = data.firstIndex(of: UInt8(ascii: "\n")) {
173 + lines.append(data.subdata(in: data.startIndex..<newline))
174 + data.removeSubrange(data.startIndex...newline)
175 + }
176 + if flush, !data.isEmpty {
177 + lines.append(data)
178 + data.removeAll()
179 + }
180 + return lines
181 + }
182 +
183 + var string: String {
184 + lock.lock()
185 + defer { lock.unlock() }
186 + return String(data: data, encoding: .utf8) ?? ""
187 + }
188 +}
added Sources/ZyquoMLX/PyBridge/scripts/zyquo_convert.py +69 −0
@@ -0,0 +1,69 @@
1 +#
2 +# zyquo_convert.py
3 +# Zyquo MLX
4 +#
5 +# Author: Simon-Pierre Boucher
6 +# Mail: contact@spboucher.ai
7 +#
8 +# Conversion/quantization driver: HF model → MLX format, optionally
9 +# quantized (mlx_lm.convert semantics — docs/MLX-RESEARCH.md §5.2).
10 +# JSON-lines protocol on stdout.
11 +#
12 +# Usage: zyquo_convert.py --hf-path <repo-or-dir> --mlx-path <dir>
13 +# [--quantize --q-bits 4 --q-group-size 64 --q-mode affine]
14 +# [--dtype float16|bfloat16|float32] [--dequantize]
15 +
16 +import argparse
17 +import json
18 +import sys
19 +import time
20 +
21 +
22 +def emit(obj):
23 + sys.stdout.write(json.dumps(obj) + "\n")
24 + sys.stdout.flush()
25 +
26 +
27 +def main():
28 + parser = argparse.ArgumentParser()
29 + parser.add_argument("--hf-path", required=True)
30 + parser.add_argument("--mlx-path", required=True)
31 + parser.add_argument("--quantize", action="store_true")
32 + parser.add_argument("--q-bits", type=int, default=None)
33 + parser.add_argument("--q-group-size", type=int, default=None)
34 + parser.add_argument("--q-mode", default="affine",
35 + choices=["affine", "mxfp4", "nvfp4", "mxfp8"])
36 + parser.add_argument("--dtype", default=None,
37 + choices=["float16", "bfloat16", "float32"])
38 + parser.add_argument("--dequantize", action="store_true")
39 + args = parser.parse_args()
40 +
41 + try:
42 + emit({"event": "start", "stage": "convert", "hf_path": args.hf_path, "ts": time.time()})
43 +
44 + from mlx_lm import convert
45 +
46 + kwargs = {
47 + "hf_path": args.hf_path,
48 + "mlx_path": args.mlx_path,
49 + "quantize": args.quantize,
50 + "dequantize": args.dequantize,
51 + }
52 + if args.q_bits is not None:
53 + kwargs["q_bits"] = args.q_bits
54 + if args.q_group_size is not None:
55 + kwargs["q_group_size"] = args.q_group_size
56 + if args.quantize:
57 + kwargs["q_mode"] = args.q_mode
58 + if args.dtype:
59 + kwargs["dtype"] = args.dtype
60 +
61 + convert(**kwargs)
62 + emit({"event": "done", "mlx_path": args.mlx_path})
63 + except Exception as exc: # noqa: BLE001
64 + emit({"event": "error", "message": str(exc), "type": type(exc).__name__})
65 + sys.exit(1)
66 +
67 +
68 +if __name__ == "__main__":
69 + main()
added Sources/ZyquoMLX/PyBridge/scripts/zyquo_fuse.py +63 −0
@@ -0,0 +1,63 @@
1 +#
2 +# zyquo_fuse.py
3 +# Zyquo MLX
4 +#
5 +# Author: Simon-Pierre Boucher
6 +# Mail: contact@spboucher.ai
7 +#
8 +# Adapter-fuse driver: fuses a LoRA adapter into its base model and saves a
9 +# standalone MLX model directory (mlx_lm.fuse semantics —
10 +# docs/TRAINING-RESEARCH.md §2.2). JSON-lines protocol on stdout.
11 +#
12 +# Usage: zyquo_fuse.py --model <base> --adapter-path <dir> --save-path <dir>
13 +# [--dequantize]
14 +
15 +import argparse
16 +import json
17 +import sys
18 +import time
19 +
20 +
21 +def emit(obj):
22 + sys.stdout.write(json.dumps(obj) + "\n")
23 + sys.stdout.flush()
24 +
25 +
26 +def main():
27 + parser = argparse.ArgumentParser()
28 + parser.add_argument("--model", required=True)
29 + parser.add_argument("--adapter-path", required=True)
30 + parser.add_argument("--save-path", required=True)
31 + parser.add_argument("--dequantize", action="store_true")
32 + args = parser.parse_args()
33 +
34 + try:
35 + emit({"event": "start", "stage": "fuse", "model": args.model, "ts": time.time()})
36 +
37 + # mlx_lm.fuse exposes only a CLI main(); call it with argv to stay
38 + # on the supported surface of the pinned version.
39 + from mlx_lm import fuse
40 +
41 + argv = [
42 + "--model", args.model,
43 + "--adapter-path", args.adapter_path,
44 + "--save-path", args.save_path,
45 + ]
46 + if args.dequantize:
47 + argv.append("--de-quantize")
48 +
49 + old_argv = sys.argv
50 + sys.argv = ["mlx_lm.fuse"] + argv
51 + try:
52 + fuse.main()
53 + finally:
54 + sys.argv = old_argv
55 +
56 + emit({"event": "done", "save_path": args.save_path})
57 + except Exception as exc: # noqa: BLE001
58 + emit({"event": "error", "message": str(exc), "type": type(exc).__name__})
59 + sys.exit(1)
60 +
61 +
62 +if __name__ == "__main__":
63 + main()
added Sources/ZyquoMLX/PyBridge/scripts/zyquo_train.py +82 −0
@@ -0,0 +1,82 @@
1 +#
2 +# zyquo_train.py
3 +# Zyquo MLX
4 +#
5 +# Author: Simon-Pierre Boucher
6 +# Mail: contact@spboucher.ai
7 +#
8 +# Training driver for the Zyquo MLX Swift app.
9 +#
10 +# Wraps mlx_lm.lora.run() with a custom TrainingCallback and emits a stable
11 +# JSON-lines protocol on stdout (one JSON object per line, event-typed).
12 +# We NEVER rely on mlx-lm's own stdout format (it changed between 0.31.3 and
13 +# main — docs/TRAINING-RESEARCH.md §5). Pinned against mlx-lm==0.31.3.
14 +#
15 +# Usage: zyquo_train.py --config <run-config.yaml>
16 +#
17 +# Events: {"event":"start", ...} {"event":"train", ...} {"event":"val", ...}
18 +# {"event":"save", ...} {"event":"done"} {"event":"error", ...}
19 +
20 +import argparse
21 +import json
22 +import sys
23 +import time
24 +import types
25 +
26 +
27 +def emit(obj):
28 + sys.stdout.write(json.dumps(obj) + "\n")
29 + sys.stdout.flush()
30 +
31 +
32 +class ZyquoCallback:
33 + """Receives the stable mlx-lm TrainingCallback dict payloads
34 + (docs/TRAINING-RESEARCH.md §5.2) and re-emits them as JSON lines."""
35 +
36 + def on_train_loss_report(self, info):
37 + emit({"event": "train", **info, "ts": time.time()})
38 +
39 + def on_val_loss_report(self, info):
40 + emit({"event": "val", **info, "ts": time.time()})
41 +
42 +
43 +def main():
44 + parser = argparse.ArgumentParser()
45 + parser.add_argument("--config", required=True, help="YAML run config (mlx-lm schema)")
46 + args = parser.parse_args()
47 +
48 + try:
49 + import yaml
50 + from mlx_lm import lora
51 +
52 + with open(args.config) as f:
53 + config = yaml.safe_load(f)
54 +
55 + # Build the args namespace exactly like mlx_lm.lora's CLI does:
56 + # defaults first, then config overrides.
57 + run_args = dict(lora.CONFIG_DEFAULTS)
58 + run_args.update(config)
59 + ns = types.SimpleNamespace(**run_args)
60 +
61 + emit({
62 + "event": "start",
63 + "model": ns.model,
64 + "fine_tune_type": ns.fine_tune_type,
65 + "iters": ns.iters,
66 + "batch_size": ns.batch_size,
67 + "learning_rate": ns.learning_rate,
68 + "adapter_path": ns.adapter_path,
69 + })
70 +
71 + lora.run(ns, training_callback=ZyquoCallback())
72 + emit({"event": "done"})
73 + except KeyboardInterrupt:
74 + emit({"event": "error", "message": "cancelled"})
75 + sys.exit(130)
76 + except Exception as exc: # noqa: BLE001 - single funnel to the app
77 + emit({"event": "error", "message": str(exc), "type": type(exc).__name__})
78 + sys.exit(1)
79 +
80 +
81 +if __name__ == "__main__":
82 + main()
added Sources/ZyquoMLX/Training/MetricsStream.swift +68 −0
@@ -0,0 +1,68 @@
1 +//
2 +// MetricsStream.swift
3 +// Zyquo MLX
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Mail: contact@spboucher.ai
7 +//
8 +
9 +import Foundation
10 +
11 +/// Typed training events, decoded from the PyBridge JSON-lines protocol
12 +/// (`zyquo_train.py`; payload fields are the stable mlx-lm callback dicts —
13 +/// docs/TRAINING-RESEARCH.md §5.2).
14 +enum TrainingEvent: Sendable {
15 + case started(model: String, iterations: Int)
16 + case metric(TrainingMetric)
17 + case checkpointSaved(fileName: String, iteration: Int)
18 + case finished
19 + case failed(message: String)
20 +}
21 +
22 +enum MetricsStream {
23 +
24 + /// Decode one Python event into a training event (nil = ignorable).
25 + static func decode(_ event: PythonEvent) -> TrainingEvent? {
26 + switch event.event {
27 + case "start":
28 + return .started(
29 + model: event.string("model") ?? "?",
30 + iterations: event.int("iters") ?? 0)
31 +
32 + case "train":
33 + return .metric(
34 + TrainingMetric(
35 + iteration: event.int("iteration") ?? 0,
36 + trainLoss: event.double("train_loss"),
37 + valLoss: nil,
38 + learningRate: event.double("learning_rate"),
39 + iterationsPerSecond: event.double("iterations_per_second"),
40 + tokensPerSecond: event.double("tokens_per_second"),
41 + trainedTokens: event.int("trained_tokens"),
42 + peakMemoryGB: event.double("peak_memory"),
43 + timestamp: .now))
44 +
45 + case "val":
46 + return .metric(
47 + TrainingMetric(
48 + iteration: event.int("iteration") ?? 0,
49 + trainLoss: nil,
50 + valLoss: event.double("val_loss"),
51 + learningRate: nil,
52 + iterationsPerSecond: nil,
53 + tokensPerSecond: nil,
54 + trainedTokens: nil,
55 + peakMemoryGB: nil,
56 + timestamp: .now))
57 +
58 + case "done":
59 + return .finished
60 +
61 + case "error":
62 + return .failed(message: event.string("message") ?? "unknown error")
63 +
64 + default:
65 + return nil
66 + }
67 + }
68 +}
added Sources/ZyquoMLX/Training/RunStore.swift +125 −0
@@ -0,0 +1,125 @@
1 +//
2 +// RunStore.swift
3 +// Zyquo MLX
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Mail: contact@spboucher.ai
7 +//
8 +
9 +import Foundation
10 +
11 +/// Persists training runs, their metrics history, and checkpoint inventory
12 +/// under `~/Library/Application Support/ZyquoMLX/Runs/<run-id>/`.
13 +///
14 +/// Layout:
15 +/// run.json — the `TrainingRun` (state, config, our iteration counter)
16 +/// metrics.jsonl — append-only `TrainingMetric` history (chart replay)
17 +/// adapters/ — mlx-lm adapter dir (adapters.safetensors, NNNNNNN_… checkpoints)
18 +/// config.yaml — the exact mlx-lm config the run used
19 +actor RunStore {
20 +
21 + static let shared = RunStore()
22 +
23 + private let root: URL
24 +
25 + init(root: URL = PersistenceService.runsDirectory) {
26 + self.root = root
27 + }
28 +
29 + func directory(for id: UUID) -> URL {
30 + root.appendingPathComponent(id.uuidString, isDirectory: true)
31 + }
32 +
33 + func adaptersDirectory(for run: TrainingRun) -> URL {
34 + run.directory.appendingPathComponent("adapters", isDirectory: true)
35 + }
36 +
37 + // MARK: - CRUD
38 +
39 + func create(_ run: TrainingRun) throws {
40 + try FileManager.default.createDirectory(
41 + at: run.directory, withIntermediateDirectories: true)
42 + try save(run)
43 + }
44 +
45 + func save(_ run: TrainingRun) throws {
46 + try PersistenceService.saveJSON(run, to: run.directory.appendingPathComponent("run.json"))
47 + }
48 +
49 + func load(id: UUID) throws -> TrainingRun {
50 + try PersistenceService.loadJSON(
51 + TrainingRun.self, from: directory(for: id).appendingPathComponent("run.json"))
52 + }
53 +
54 + func scan() throws -> [TrainingRun] {
55 + let fm = FileManager.default
56 + guard fm.fileExists(atPath: root.path) else { return [] }
57 + return try fm.contentsOfDirectory(at: root, includingPropertiesForKeys: nil)
58 + .compactMap {
59 + try? PersistenceService.loadJSON(
60 + TrainingRun.self, from: $0.appendingPathComponent("run.json"))
61 + }
62 + .sorted { $0.createdAt > $1.createdAt }
63 + }
64 +
65 + func delete(_ run: TrainingRun) throws {
66 + try FileManager.default.removeItem(at: run.directory)
67 + }
68 +
69 + // MARK: - Metrics history
70 +
71 + func append(metric: TrainingMetric, to run: TrainingRun) {
72 + guard let data = try? JSONEncoder().encode(metric),
73 + let line = String(data: data, encoding: .utf8)
74 + else { return }
75 + let url = run.directory.appendingPathComponent("metrics.jsonl")
76 + if let handle = FileHandle(forWritingAtPath: url.path) {
77 + handle.seekToEndOfFile()
78 + handle.write(Data((line + "\n").utf8))
79 + try? handle.close()
80 + } else {
81 + try? (line + "\n").write(to: url, atomically: true, encoding: .utf8)
82 + }
83 + }
84 +
85 + func metricsHistory(for run: TrainingRun) -> [TrainingMetric] {
86 + guard
87 + let content = try? String(
88 + contentsOf: run.directory.appendingPathComponent("metrics.jsonl"),
89 + encoding: .utf8)
90 + else { return [] }
91 + let decoder = JSONDecoder()
92 + return content.split(separator: "\n").compactMap {
93 + $0.data(using: .utf8).flatMap { try? decoder.decode(TrainingMetric.self, from: $0) }
94 + }
95 + }
96 +
97 + // MARK: - Checkpoints
98 +
99 + /// Scan the adapter directory for numbered checkpoints
100 + /// (`{iter:07d}_adapters.safetensors` — docs/TRAINING-RESEARCH.md §2.1).
101 + func checkpoints(for run: TrainingRun) -> [Checkpoint] {
102 + let dir = adaptersDirectory(for: run)
103 + guard
104 + let files = try? FileManager.default.contentsOfDirectory(
105 + at: dir, includingPropertiesForKeys: [.creationDateKey])
106 + else { return [] }
107 + return files.compactMap { url -> Checkpoint? in
108 + let name = url.lastPathComponent
109 + guard name.hasSuffix("_adapters.safetensors"),
110 + let iteration = Int(name.prefix(7))
111 + else { return nil }
112 + let savedAt =
113 + (try? url.resourceValues(forKeys: [.creationDateKey]))?.creationDate ?? .now
114 + return Checkpoint(fileName: name, iteration: iteration, fileURL: url, savedAt: savedAt)
115 + }
116 + .sorted { $0.iteration < $1.iteration }
117 + }
118 +
119 + /// Latest usable adapter weights for resume/fuse/inference.
120 + func latestAdapter(for run: TrainingRun) -> URL? {
121 + let main = adaptersDirectory(for: run).appendingPathComponent("adapters.safetensors")
122 + if FileManager.default.fileExists(atPath: main.path) { return main }
123 + return checkpoints(for: run).last?.fileURL
124 + }
125 +}
added Sources/ZyquoMLX/Training/TrainingService.swift +229 −0
@@ -0,0 +1,229 @@
1 +//
2 +// TrainingService.swift
3 +// Zyquo MLX
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Mail: contact@spboucher.ai
7 +//
8 +
9 +import Foundation
10 +
11 +enum TrainingServiceError: LocalizedError {
12 + case memoryGate(verdict: MemoryVerdict, suggestions: [String])
13 + case qloraNeedsQuantizedBase
14 + case alreadyRunning(UUID)
15 + case nothingToResume
16 +
17 + var errorDescription: String? {
18 + switch self {
19 + case .memoryGate(_, let suggestions):
20 + "This configuration won't fit in memory on this Mac. Try: "
21 + + suggestions.joined(separator: " · ")
22 + case .qloraNeedsQuantizedBase:
23 + "QLoRA requires a quantized base model. Pick a quantized base (e.g. a -4bit model) or quantize it first in Convert."
24 + case .alreadyRunning(let id):
25 + "Run \(id.uuidString.prefix(8)) is already active."
26 + case .nothingToResume:
27 + "No saved adapter weights found to resume from."
28 + }
29 + }
30 +}
31 +
32 +/// Orchestrates fine-tuning runs (charter 3.B): builds the mlx-lm YAML
33 +/// config, drives `zyquo_train.py` through `PythonRunner`, persists state and
34 +/// metrics via `RunStore`, and exposes a typed event stream. Runs are
35 +/// cancellable; resume is an honest warm start (docs/TRAINING-RESEARCH.md §5.3).
36 +actor TrainingService {
37 +
38 + static let shared = TrainingService()
39 +
40 + private var activeRunID: UUID?
41 + private var activeStreamTask: Task<Void, Never>?
42 +
43 + // MARK: - Configure
44 +
45 + /// Create (and persist) a new run after memory-gating it.
46 + func createRun(
47 + name: String,
48 + baseModel: LocalModel,
49 + dataset: Dataset,
50 + method: FineTuneMethod,
51 + hyperParams: HyperParams
52 + ) async throws -> TrainingRun {
53 + // QLoRA = LoRA on a quantized base; LoRA on a quantized base IS QLoRA —
54 + // keep the user's mental model consistent with reality.
55 + if method == .qlora && baseModel.quantization == nil {
56 + throw TrainingServiceError.qloraNeedsQuantizedBase
57 + }
58 +
59 + let verdict = MemoryAdvisor.trainingVerdict(
60 + for: baseModel, method: method, params: hyperParams)
61 + if verdict == .wontFit {
62 + throw TrainingServiceError.memoryGate(
63 + verdict: verdict,
64 + suggestions: MemoryAdvisor.suggestions(
65 + for: baseModel, method: method, params: hyperParams))
66 + }
67 +
68 + let id = UUID()
69 + let run = TrainingRun(
70 + id: id,
71 + name: name,
72 + baseModelID: baseModel.id,
73 + datasetID: dataset.id,
74 + method: method,
75 + hyperParams: hyperParams,
76 + state: .configured,
77 + directory: await RunStore.shared.directory(for: id),
78 + createdAt: .now,
79 + startedAt: nil,
80 + finishedAt: nil,
81 + completedIterations: 0,
82 + checkpoints: [],
83 + lastError: nil
84 + )
85 + try await RunStore.shared.create(run)
86 + try writeConfig(run: run, baseModel: baseModel, dataset: dataset, resumeFrom: nil)
87 + return run
88 + }
89 +
90 + // MARK: - Start / resume
91 +
92 + /// Start (or warm-start-resume) a run, streaming typed events.
93 + func start(run: TrainingRun, baseModel: LocalModel, dataset: Dataset, resume: Bool = false)
94 + async throws -> AsyncStream<TrainingEvent>
95 + {
96 + if let activeRunID { throw TrainingServiceError.alreadyRunning(activeRunID) }
97 +
98 + var run = run
99 + var resumeAdapter: URL?
100 + if resume {
101 + guard let adapter = await RunStore.shared.latestAdapter(for: run) else {
102 + throw TrainingServiceError.nothingToResume
103 + }
104 + resumeAdapter = adapter
105 + }
106 + try writeConfig(run: run, baseModel: baseModel, dataset: dataset, resumeFrom: resumeAdapter)
107 +
108 + run.state = .running
109 + run.startedAt = run.startedAt ?? .now
110 + run.lastError = nil
111 + try await RunStore.shared.save(run)
112 + activeRunID = run.id
113 +
114 + let configPath = run.directory.appendingPathComponent("config.yaml").path
115 + let pythonStream = try await PythonRunner.shared.stream(
116 + script: "zyquo_train", arguments: ["--config", configPath])
117 +
118 + let (stream, continuation) = AsyncStream.makeStream(of: TrainingEvent.self)
119 + let runSnapshot = run
120 +
121 + activeStreamTask = Task {
122 + var current = runSnapshot
123 + do {
124 + for try await pythonEvent in pythonStream {
125 + guard let event = MetricsStream.decode(pythonEvent) else { continue }
126 + switch event {
127 + case .metric(let metric):
128 + await RunStore.shared.append(metric: metric, to: current)
129 + current.completedIterations = max(
130 + current.completedIterations, metric.iteration)
131 + case .failed(let message):
132 + current.lastError = message
133 + default:
134 + break
135 + }
136 + continuation.yield(event)
137 + }
138 + current.state = current.lastError == nil ? .completed : .failed
139 + } catch is CancellationError {
140 + current.state = .cancelled
141 + } catch {
142 + current.state = .failed
143 + current.lastError = error.localizedDescription
144 + continuation.yield(.failed(message: error.localizedDescription))
145 + }
146 + current.finishedAt = .now
147 + current.checkpoints = await RunStore.shared.checkpoints(for: current)
148 + try? await RunStore.shared.save(current)
149 + self.clearActive()
150 + continuation.finish()
151 + }
152 +
153 + return stream
154 + }
155 +
156 + /// Cancel the active run (terminates the Python process via the stream's
157 + /// onTermination; state persists as cancelled with checkpoints intact for
158 + /// warm-start resume).
159 + func cancel() {
160 + activeStreamTask?.cancel()
161 + }
162 +
163 +
164 + private func clearActive() {
165 + activeRunID = nil
166 + activeStreamTask = nil
167 + }
168 +
169 + // MARK: - Config generation
170 +
171 + /// Serialize the run into mlx-lm's YAML schema
172 + /// (exact keys/defaults per docs/TRAINING-RESEARCH.md §1.2; the
173 + /// config-only keys `lora_parameters` make YAML mandatory).
174 + private func writeConfig(
175 + run: TrainingRun, baseModel: LocalModel, dataset: Dataset, resumeFrom: URL?
176 + ) throws {
177 + let hp = run.hyperParams
178 + let adapterPath = run.directory.appendingPathComponent("adapters").path
179 +
180 + // .qlora maps to upstream "lora" — QLoRA is implied by the quantized base.
181 + let fineTuneType = run.method == .qlora ? "lora" : run.method.rawValue
182 +
183 + var lines: [String] = [
184 + "# Generated by Zyquo MLX — run \(run.id.uuidString)",
185 + "model: \(yamlQuote(baseModel.directory.path))",
186 + "train: true",
187 + "data: \(yamlQuote(dataset.directory.path))",
188 + "fine_tune_type: \(fineTuneType)",
189 + "optimizer: \(hp.optimizer)",
190 + "mask_prompt: \(hp.maskPrompt)",
191 + "num_layers: \(hp.numLayers)",
192 + "batch_size: \(hp.batchSize)",
193 + "iters: \(hp.iterations)",
194 + "val_batches: \(hp.valBatches)",
195 + "learning_rate: \(hp.learningRate)",
196 + "steps_per_report: \(hp.stepsPerReport)",
197 + "steps_per_eval: \(hp.stepsPerEval)",
198 + "grad_accumulation_steps: \(hp.gradAccumulationSteps)",
199 + "save_every: \(hp.saveEvery)",
200 + "adapter_path: \(yamlQuote(adapterPath))",
201 + "max_seq_length: \(hp.maxSeqLength)",
202 + "grad_checkpoint: \(hp.gradCheckpoint)",
203 + "seed: \(hp.seed)",
204 + ]
205 +
206 + if let resumeFrom {
207 + lines.append("resume_adapter_file: \(yamlQuote(resumeFrom.path))")
208 + }
209 +
210 + if run.method != .full {
211 + lines.append("lora_parameters:")
212 + lines.append(" rank: \(hp.rank)")
213 + lines.append(" scale: \(hp.scale)")
214 + lines.append(" dropout: \(hp.dropout)")
215 + if let keys = hp.keys, !keys.isEmpty {
216 + lines.append(" keys: [\(keys.map(yamlQuote).joined(separator: ", "))]")
217 + }
218 + }
219 +
220 + try (lines.joined(separator: "\n") + "\n")
221 + .write(
222 + to: run.directory.appendingPathComponent("config.yaml"),
223 + atomically: true, encoding: .utf8)
224 + }
225 +
226 + private func yamlQuote(_ value: String) -> String {
227 + "\"" + value.replacingOccurrences(of: "\"", with: "\\\"") + "\""
228 + }
229 +}
modified docs/PLAN.md +11 −1
@@ -101,7 +101,17 @@ sources (no guessed names — `#huggingFaceTokenizerLoader` macro requires
101 101 through `perform(nonSendable:)`). CLI POC proves the two-model-type gate with
102 102 real downloads in the app's Models library. Build is arm64-only via
103 103 `--arch arm64` (swiftbuild otherwise builds universal and x86_64 fails).
104 −## Phase 3 — Training, Quantization & Conversion (not started)
104 +## Phase 3 — Training, Quantization & Conversion (in progress)
105 +
106 +- [ ] `PyBridge/PythonRunner` — Process wrapper over the venv, JSON-lines progress protocol, cancellation
107 +- [ ] `PyBridge/PythonEnvironment` — uv-provisioned pinned venv in App Support (status/repair; mlx-lm now, vlm/whisper/audio on demand)
108 +- [ ] `PyBridge/scripts/` — zyquo_train.py (TrainingCallback → JSON), zyquo_fuse.py, zyquo_convert.py
109 +- [ ] `Data/DatasetService` + `DatasetFormats` — import/validate JSONL (chat/completions/text), auto split, malformed-row report, token stats, preview
110 +- [ ] `Training/TrainingService` + `RunStore` + `MetricsStream` — cancellable/resumable runs, persisted state, checkpoints, live metrics
111 +- [ ] `Convert/ConversionService` + `QuantConfig` — Swift-native affine quant for safetensors LLMs; Python bridge for fuse/dequantize
112 +- [ ] CLI: `--train`, `--fuse`, `--validate-dataset` to drive the services (UI in Phase 6)
113 +- [ ] PHASE GATE: real LoRA/QLoRA fine-tune on a tiny dataset → live loss ↓, adapter saved, fused, fused model generates differently
114 +- [ ] Phase checkpoint: build green, gate verified, summary
105 115 ## Phase 4 — Design System & UI (not started)
106 116 ## Phase 5 — App Icon (not started)
107 117 ## Phase 6 — Features (not started)
108 118