diff --git a/Package.swift b/Package.swift index 799976b..9995d51 100644 --- a/Package.swift +++ b/Package.swift @@ -75,6 +75,10 @@ let package = Package( .executableTarget(name: "GuardrailDemo", dependencies: ["FluidUse"], exclude: ["README.md"]), .executableTarget(name: "FrontDoorDemo", dependencies: ["FluidUse"], exclude: ["README.md"]), .executableTarget(name: "ShortReplyCheck", dependencies: ["FluidUse"]), + .executableTarget(name: "CodeWriterCheck", dependencies: ["FluidUse"]), + .executableTarget( + name: "CodeWriterDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh"], + resources: [.copy("Resources")]), .executableTarget(name: "ShortReplyDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh", "mock-feed"]), .executableTarget(name: "GLiClassServe", dependencies: ["FluidUse"]), .executableTarget( diff --git a/README.md b/README.md index d556357..d1e23d8 100644 --- a/README.md +++ b/README.md @@ -294,6 +294,25 @@ print(draft.reply, draft.timing.totalSeconds) // "Cong `Sources/ShortReplyDemo/demo.sh --x` launches it with a macmon + log terminal ([Sources/ShortReplyDemo](Sources/ShortReplyDemo/README.md)). +## Python writer + +`CodeWriterManager` writes Python from a plain-English request with +[qwen2.5-coder-0.5b-coreml](https://huggingface.co/FluidInference/qwen2.5-coder-0.5b-coreml), Qwen2.5-Coder-0.5B-Instruct +(Apache-2.0) in one 944 MB package: the prompt runs on the Neural Engine, the writing on the GPU at 36–51 tokens/s on +an M5 Pro, streamed as it goes. Pass@1: HumanEval 89/164 (PyTorch: 90/164), MBPP 119/257. + +```swift +let directory = try await CodeWriterModelStore.ensure() // pinned Hugging Face snapshot +let writer = try await CodeWriterManager.load(from: directory) +let result = try await writer.write(task: "Write a function to reverse each string in a list.") { text in + print(text) // the text so far, after every token +} +print(result.code, result.timing.tokensPerSecond) +``` + +`CodeWriterDemo` writes ten short tasks live into an editor and runs each one's asserts with `python3` +([Sources/CodeWriterDemo](Sources/CodeWriterDemo/README.md)). + ## Demo ```bash diff --git a/Sources/CodeWriterCheck/main.swift b/Sources/CodeWriterCheck/main.swift new file mode 100644 index 0000000..8f0da38 --- /dev/null +++ b/Sources/CodeWriterCheck/main.swift @@ -0,0 +1,135 @@ +import FluidUse +import Foundation + +// Parity + latency check for the Swift code-writer host: replays HumanEval prompts through CodeWriterManager, compares +// each token stream with a reference run (Python, same prompt, greedy) and runs the HumanEval tests with python3. +// swift run -c release CodeWriterCheck [limit] +// `hf` downloads the pinned snapshot through CodeWriterModelStore. + +let arguments = CommandLine.arguments +guard arguments.count >= 4 else { + FileHandle.standardError.write( + Data("usage: CodeWriterCheck [limit]\n".utf8)) + exit(2) +} +let limit = arguments.count > 4 ? Int(arguments[4]) ?? .max : .max +let system = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant." + +struct Problem: Decodable { + let taskID: String + let prompt: String + let entryPoint: String + let test: String + + enum CodingKeys: String, CodingKey { + case taskID = "task_id" + case prompt + case entryPoint = "entry_point" + case test + } +} + +struct Reference: Decodable { + let taskID: String + let tokens: [Int] + let passed: Bool + + enum CodingKeys: String, CodingKey { + case taskID = "task_id" + case tokens + case passed + } +} + +func lines(_ path: String) throws -> [Data] { + try String(contentsOf: URL(fileURLWithPath: path), encoding: .utf8) + .split(whereSeparator: \.isNewline).map { Data($0.utf8) } +} + +let problems = try lines(arguments[2]).prefix(limit).map { try JSONDecoder().decode(Problem.self, from: $0) } +let references = Dictionary( + uniqueKeysWithValues: try lines(arguments[3]).map { try JSONDecoder().decode(Reference.self, from: $0) } + .map { ($0.taskID, $0) }) + +/// Mirrors `humaneval.py`: keep the prompt's imports/helpers whether the model wrote the body or the whole function. +func program(_ problem: Problem, _ code: String) -> String { + let body = + code.contains("def \(problem.entryPoint)") + ? problem.prompt + "\n pass\n\n" + code : problem.prompt + code + return body + "\n\n" + problem.test + "\n\ncheck(\(problem.entryPoint))\n" +} + +func passes(_ source: String) throws -> Bool { + let file = FileManager.default.temporaryDirectory.appendingPathComponent("codewriter-\(UUID().uuidString).py") + try source.write(to: file, atomically: true, encoding: .utf8) + defer { try? FileManager.default.removeItem(at: file) } + let process = Process() + process.executableURL = URL(fileURLWithPath: "/usr/bin/env") + process.arguments = ["python3", file.path] + process.standardOutput = FileHandle.nullDevice + process.standardError = FileHandle.nullDevice + try process.run() + let deadline = Date().addingTimeInterval(10) + while process.isRunning, Date() < deadline { usleep(20_000) } + if process.isRunning { + process.terminate() + return false + } + return process.terminationStatus == 0 +} + +@available(macOS 15.0, *) +func run() async throws { + let directory = + arguments[1] == "hf" + ? try await CodeWriterModelStore.ensure { file, bytes in if bytes > 0 { print("downloaded \(file)") } } + : URL(fileURLWithPath: arguments[1]) + let manager = try await CodeWriterManager.load(from: directory) + try await manager.warmUp() + var passed = 0 + var referencePassed = 0 + var identical = 0 + var prefills: [Double] = [] + var perToken: [Double] = [] + for (index, problem) in problems.enumerated() { + let user = "Complete the following python code:\n```python\n\(problem.prompt)\n```" + let completion = try await manager.write(task: user, system: system, maxNewTokens: 512) + let ok = try passes(program(problem, completion.code)) + passed += ok ? 1 : 0 + prefills.append(completion.timing.prefillSeconds) + if completion.timing.generatedTokens > 0 { + perToken.append(completion.timing.decodeSeconds / Double(completion.timing.generatedTokens)) + } + var note = "" + if let reference = references[problem.taskID] { + referencePassed += reference.passed ? 1 : 0 + if completion.tokens == reference.tokens { + identical += 1 + } else { + let diverge = + zip(completion.tokens, reference.tokens).enumerated().first { $0.element.0 != $0.element.1 }?.offset + ?? min(completion.tokens.count, reference.tokens.count) + note = " differs@\(diverge) ref=\(reference.passed ? "PASS" : "fail")" + } + } + print( + "\(index + 1)/\(problems.count) \(problem.taskID) \(ok ? "PASS" : "fail") " + + "\(completion.timing.generatedTokens) tok\(note)") + } + func median(_ values: [Double]) -> Double { + let sorted = values.sorted() + return sorted.isEmpty ? 0 : sorted[sorted.count / 2] + } + print("pass@1 \(passed)/\(problems.count); reference \(referencePassed); identical token streams \(identical)") + print( + String( + format: "prefill p50: %.1f ms decode p50: %.1f ms/token (%.0f tok/s)", median(prefills) * 1000, + median(perToken) * 1000, 1 / max(median(perToken), 1e-9))) +} + +if #available(macOS 15.0, *) { + try await run() +} else { + FileHandle.standardError.write(Data("CodeWriterCheck needs macOS 15\n".utf8)) + exit(2) +} diff --git a/Sources/CodeWriterDemo/CodeWriterDemoApp.swift b/Sources/CodeWriterDemo/CodeWriterDemoApp.swift new file mode 100644 index 0000000..bdde663 --- /dev/null +++ b/Sources/CodeWriterDemo/CodeWriterDemoApp.swift @@ -0,0 +1,25 @@ +import AppKit +import SwiftUI + +@available(macOS 15.0, *) +@main +struct CodeWriterDemoApp: App { + @StateObject private var model = WriterModel() + + init() { + setvbuf(stdout, nil, _IOLBF, 0) + // Bare SwiftPM executables start as background processes; make this one a regular windowed app. + NSApplication.shared.setActivationPolicy(.regular) + NSApplication.shared.activate(ignoringOtherApps: true) + } + + var body: some Scene { + WindowGroup("Python writer — on-device") { + ContentView() + .environmentObject(model) + .frame(minWidth: 1000, minHeight: 680) + .onAppear { model.start() } + } + .defaultSize(width: 1280, height: 860) + } +} diff --git a/Sources/CodeWriterDemo/ContentView.swift b/Sources/CodeWriterDemo/ContentView.swift new file mode 100644 index 0000000..2e402e3 --- /dev/null +++ b/Sources/CodeWriterDemo/ContentView.swift @@ -0,0 +1,308 @@ +import SwiftUI + +@available(macOS 15.0, *) +struct ContentView: View { + @EnvironmentObject private var model: WriterModel + + var body: some View { + VStack(spacing: 0) { + Header() + Divider() + HSplitView { + TaskList().frame(minWidth: 340, idealWidth: 400, maxWidth: 520) + VStack(spacing: 0) { + Editor() + Divider() + Console().frame(height: 150) + } + .frame(minWidth: 520) + } + Divider() + CustomTaskBar() + } + .background(Color(nsColor: .windowBackgroundColor)) + } +} + +@available(macOS 15.0, *) +struct Header: View { + @EnvironmentObject private var model: WriterModel + + var body: some View { + HStack(alignment: .center, spacing: 18) { + VStack(alignment: .leading, spacing: 4) { + Text("Python writer").font(.title2.bold()) + Text("Qwen2.5-Coder 0.5B · Neural Engine + GPU · offline, nothing leaves this Mac") + .font(.caption).foregroundStyle(.secondary) + HStack(spacing: 14) { + Link(destination: URL(string: "https://github.com/FluidInference/FluidUse")!) { + Label( + "github.com/FluidInference/FluidUse", systemImage: "chevron.left.forwardslash.chevron.right" + ) + } + Link(destination: URL(string: "https://huggingface.co/FluidInference/qwen2.5-coder-0.5b-coreml")!) { + Label("huggingface.co/FluidInference/qwen2.5-coder-0.5b-coreml", systemImage: "shippingbox") + } + } + .font(.caption.weight(.medium)) + } + Spacer() + VStack(alignment: .trailing, spacing: 2) { + HStack(alignment: .firstTextBaseline, spacing: 6) { + Text("\(model.passedCount)").font(.system(size: 34, weight: .bold, design: .rounded)) + .foregroundStyle(.green).contentTransition(.numericText()) + Text("/ \(model.doneCount) written pass their tests") + .font(.system(size: 15, weight: .semibold, design: .rounded)) + .contentTransition(.numericText()) + } + Text(status).font(.caption).foregroundStyle(.secondary) + } + HStack(spacing: 8) { + Button { + model.play() + } label: { + Label(model.paused ? "Resume" : "Play", systemImage: "play.fill") + } + .keyboardShortcut(.return, modifiers: [.command]) + .disabled(!model.canPlay) + Button { + model.pause() + } label: { + Label("Pause", systemImage: "pause.fill") + } + .disabled(!model.isRunning || model.paused) + Button { + model.replay() + } label: { + Label("Replay", systemImage: "arrow.counterclockwise") + } + .keyboardShortcut("r", modifiers: [.command]) + .disabled(model.writerMissing) + } + } + .padding(16) + } + + private var status: String { + let speed = model.tokensPerSecond > 0 ? String(format: " · %.0f tokens/s", model.tokensPerSecond) : "" + switch model.step { + case .loading(let message): return message + case .idle: return model.paused ? "Paused" + speed : "Ready · press Play (⌘↩)" + case .writing: return (model.paused ? "Pausing after this task" : "Writing…") + speed + case .finished: return "Done · \(model.tasks.count) tasks" + speed + case .failed(let message): return message + } + } +} + +@available(macOS 15.0, *) +struct TaskList: View { + @EnvironmentObject private var model: WriterModel + + var body: some View { + ScrollViewReader { proxy in + ScrollView { + LazyVStack(spacing: 0) { + ForEach(model.tasks) { task in + TaskRow( + task: task, status: model.statuses[task.id] ?? .pending, active: model.current == task.id + ) + .id(task.id) + .onTapGesture { model.select(task) } + Divider().opacity(0.4) + } + } + } + .onChange(of: model.current) { + if let current = model.current { withAnimation { proxy.scrollTo(current, anchor: .center) } } + } + } + } +} + +@available(macOS 15.0, *) +struct TaskRow: View { + let task: CodingTask + let status: WriterModel.Status + let active: Bool + + var body: some View { + HStack(alignment: .top, spacing: 10) { + StatusIcon(status: status).frame(width: 22) + VStack(alignment: .leading, spacing: 4) { + Text(task.topic).font(.system(size: 10, weight: .bold, design: .rounded)).foregroundStyle(.cyan) + .padding(.horizontal, 6).padding(.vertical, 2) + .background(Capsule().fill(Color.cyan.opacity(0.15))) + Text(task.prompt).font(.system(size: 13)).foregroundStyle(.primary.opacity(0.9)) + .fixedSize(horizontal: false, vertical: true) + } + Spacer(minLength: 0) + } + .padding(.horizontal, 14).padding(.vertical, 9) + .background(active ? Color.accentColor.opacity(0.14) : .clear) + .overlay(alignment: .leading) { if active { Rectangle().fill(Color.accentColor).frame(width: 3) } } + .contentShape(Rectangle()) + } +} + +@available(macOS 15.0, *) +struct StatusIcon: View { + let status: WriterModel.Status + + var body: some View { + switch status { + case .pending: + Image(systemName: "circle").foregroundStyle(.secondary) + case .writing: + Image(systemName: "pencil.line").foregroundStyle(Color.accentColor).symbolEffect(.pulse) + case .testing: + ProgressView().controlSize(.small) + case .passed: + Image(systemName: "checkmark.circle.fill").foregroundStyle(.green) + case .failed: + Image(systemName: "xmark.circle.fill").foregroundStyle(.orange) + case .broken: + Image(systemName: "exclamationmark.triangle.fill").foregroundStyle(.red) + } + } +} + +@available(macOS 15.0, *) +struct Editor: View { + @EnvironmentObject private var model: WriterModel + + static let background = Color(red: 0.12, green: 0.12, blue: 0.14) + static let gutterBackground = Color(red: 0.10, green: 0.10, blue: 0.12) + private let lineHeight: CGFloat = 21 + + var body: some View { + VStack(alignment: .leading, spacing: 0) { + TabBar(title: headerTitle) + let lines = PythonHighlighter.lines(model.code) + let typing = model.isWriting + ScrollViewReader { proxy in + ScrollView(.vertical) { + VStack(alignment: .leading, spacing: 0) { + ForEach(Array(lines.enumerated()), id: \.offset) { index, line in + let isLast = index == lines.count - 1 + HStack(spacing: 0) { + Text("\(index + 1)") + .foregroundStyle(isLast && typing ? Color(white: 0.85) : Color(white: 0.40)) + .frame(width: 44, alignment: .trailing) + .padding(.trailing, 14) + .frame(maxHeight: .infinity) + .background(Self.gutterBackground) + HStack(spacing: 0) { + Text(line).lineLimit(1) + if isLast && typing { Cursor() } + } + .padding(.leading, 12) + Spacer(minLength: 0) + } + .frame(height: lineHeight) + .background(isLast && typing ? Color.white.opacity(0.05) : .clear) + } + Color.clear.frame(height: 1).id("end") + } + .font(.system(size: 14, design: .monospaced)) + .frame(maxWidth: .infinity, alignment: .topLeading) + } + .defaultScrollAnchor(.top) + .onChange(of: lines.count) { proxy.scrollTo("end", anchor: .bottom) } + } + .background(alignment: .leading) { Self.gutterBackground.frame(width: 58) } + } + .background(Self.background) + .environment(\.colorScheme, .dark) + } + + private var headerTitle: String? { + if let custom = model.customTitle { return custom } + return model.tasks.first { $0.id == model.current }?.prompt + } +} + +/// Editor tab strip: one open file, with the task it is answering on the right. +@available(macOS 15.0, *) +struct TabBar: View { + let title: String? + + var body: some View { + HStack(spacing: 0) { + HStack(spacing: 6) { + Image(systemName: "chevron.left.forwardslash.chevron.right").font(.system(size: 10, weight: .bold)) + .foregroundStyle(.yellow) + Text("solution.py").font(.system(size: 12)) + Image(systemName: "xmark").font(.system(size: 8, weight: .bold)).foregroundStyle(.secondary) + } + .padding(.horizontal, 14).padding(.vertical, 8) + .background(Editor.background) + .overlay(alignment: .top) { Rectangle().fill(Color.accentColor).frame(height: 2) } + Spacer() + if let title { + Text(title).font(.caption).foregroundStyle(.secondary).lineLimit(1).truncationMode(.tail) + .padding(.horizontal, 12) + } + } + .background(Color(red: 0.09, green: 0.09, blue: 0.10)) + } +} + +/// Blinking block cursor at the end of the line being written. +struct Cursor: View { + @State private var visible = true + + var body: some View { + Rectangle().fill(Color(white: 0.9)).frame(width: 2, height: 17) + .opacity(visible ? 1 : 0) + .onAppear { + withAnimation(.easeInOut(duration: 0.5).repeatForever()) { visible = false } + } + } +} + +@available(macOS 15.0, *) +struct Console: View { + @EnvironmentObject private var model: WriterModel + + var body: some View { + ScrollView { + VStack(alignment: .leading, spacing: 4) { + Text("$ python3 solution.py").foregroundStyle(.secondary) + ForEach(model.checks) { check in + HStack(alignment: .top, spacing: 8) { + Image(systemName: check.passed ? "checkmark" : "xmark") + .foregroundStyle(check.passed ? .green : .orange).frame(width: 14) + Text(check.test).foregroundStyle(.primary.opacity(0.85)).lineLimit(2) + } + } + if let error = model.runError { + Text(error).foregroundStyle(.red) + } else if model.customTitle != nil, !model.isBusyCustom, !model.code.isEmpty { + Text("✓ valid Python (typed task: no tests)").foregroundStyle(.green) + } + } + .font(.system(size: 12, design: .monospaced)) + .frame(maxWidth: .infinity, alignment: .leading).padding(12) + } + .background(Color(red: 0.08, green: 0.08, blue: 0.09)) + .environment(\.colorScheme, .dark) + } +} + +@available(macOS 15.0, *) +struct CustomTaskBar: View { + @EnvironmentObject private var model: WriterModel + + var body: some View { + HStack { + Image(systemName: "text.cursor").foregroundStyle(.secondary) + TextField("Or type your own: “Write a function that …”", text: $model.customTask) + .textFieldStyle(.plain).font(.title3) + .onSubmit { model.writeCustom() } + .disabled(model.isRunning || model.isBusyCustom || model.writerMissing) + if model.isBusyCustom { ProgressView().controlSize(.small) } + } + .padding(.horizontal, 16).padding(.vertical, 12) + } +} diff --git a/Sources/CodeWriterDemo/DemoLog.swift b/Sources/CodeWriterDemo/DemoLog.swift new file mode 100644 index 0000000..03b534b --- /dev/null +++ b/Sources/CodeWriterDemo/DemoLog.swift @@ -0,0 +1,22 @@ +import Foundation + +/// Colour-coded console lines for the side terminal (`demo.sh` tails stdout): model load, each task, test results. +enum DemoLog { + static func line(_ text: String, color: Int? = nil, bold: Bool = false) { + let time = timeFormatter.string(from: Date()) + var styled = text + if let color { styled = "\u{1B}[38;5;\(color)m\(styled)\u{1B}[0m" } + if bold { styled = "\u{1B}[1m\(styled)\u{1B}[0m" } + print("\u{1B}[2m\(time)\u{1B}[0m \(styled)") + } + + static func event(_ text: String) { line(text, color: 220, bold: true) } + + static func model(_ text: String) { line("[model] " + text, color: 81) } + + private static let timeFormatter: DateFormatter = { + let formatter = DateFormatter() + formatter.dateFormat = "HH:mm:ss.SSS" + return formatter + }() +} diff --git a/Sources/CodeWriterDemo/PythonHighlighter.swift b/Sources/CodeWriterDemo/PythonHighlighter.swift new file mode 100644 index 0000000..8c99eb0 --- /dev/null +++ b/Sources/CodeWriterDemo/PythonHighlighter.swift @@ -0,0 +1,96 @@ +import SwiftUI + +/// Dark-editor colouring for Python. +enum PythonHighlighter { + static let keyword = Color(red: 0.99, green: 0.37, blue: 0.64) + static let builtin = Color(red: 0.36, green: 0.85, blue: 1.0) + static let call = Color(red: 0.40, green: 0.72, blue: 0.64) + static let declaration = Color(red: 0.25, green: 0.63, blue: 0.75) + static let string = Color(red: 0.99, green: 0.42, blue: 0.36) + static let number = Color(red: 0.82, green: 0.75, blue: 0.41) + static let comment = Color(red: 0.47, green: 0.53, blue: 0.60) + static let decorator = Color(red: 0.99, green: 0.56, blue: 0.25) + static let plain = Color(white: 0.92) + + private static let keywords: Set = [ + "def", "return", "if", "elif", "else", "for", "while", "in", "not", "and", "or", "is", "import", "from", "as", + "class", "try", "except", "finally", "raise", "with", "yield", "lambda", "pass", "break", "continue", "global", + "nonlocal", "assert", "del", "None", "True", "False", "async", "await", "match", "case", + ] + private static let builtins: Set = [ + "len", "range", "print", "int", "str", "float", "list", "dict", "set", "tuple", "bool", "sorted", "sum", "min", + "max", "abs", "enumerate", "zip", "map", "filter", "any", "all", "isinstance", "self", "reversed", "round", + "ord", "chr", "type", "object", "super", + ] + private static let token = try! NSRegularExpression( + pattern: + #"(#.*$)|([rbfu]{0,2}(?:"(?:\\.|[^"\\])*"|'(?:\\.|[^'\\])*'))|(@[A-Za-z_][A-Za-z0-9_.]*)|(\b\d+(?:\.\d+)?\b)|([A-Za-z_][A-Za-z0-9_]*)"# + ) + + /// One coloured line per source line (a triple-quoted string spanning lines is tracked across lines). + static func lines(_ source: String) -> [AttributedString] { + var result: [AttributedString] = [] + var inDocstring: String? + for line in source.components(separatedBy: "\n") { + if let quote = inDocstring { + result.append(piece(line, string)) + if line.contains(quote) { inDocstring = nil } + continue + } + if let quote = ["\"\"\"", "'''"].first(where: { line.contains($0) }) { + result.append(piece(line, string)) + if (line.components(separatedBy: quote).count - 1) % 2 == 1 { inDocstring = quote } + continue + } + result.append(code(line)) + } + return result + } + + static func code(_ line: String) -> AttributedString { + var result = AttributedString() + let ns = line as NSString + var cursor = 0 + var previousWord = "" + for match in token.matches(in: line, range: NSRange(location: 0, length: ns.length)) { + if match.range.location > cursor { + result.append( + piece(ns.substring(with: NSRange(location: cursor, length: match.range.location - cursor)), plain)) + } + let text = ns.substring(with: match.range) + let color: Color + if match.range(at: 1).location != NSNotFound { + color = comment + } else if match.range(at: 2).location != NSNotFound { + color = string + } else if match.range(at: 3).location != NSNotFound { + color = decorator + } else if match.range(at: 4).location != NSNotFound { + color = number + } else if keywords.contains(text) { + color = keyword + } else if previousWord == "def" || previousWord == "class" { + color = declaration + } else if builtins.contains(text) { + color = builtin + } else if match.range.location + match.range.length < ns.length, + ns.character(at: match.range.location + match.range.length) == 0x28 // "(" + { + color = call + } else { + color = plain + } + result.append(piece(text, color)) + if match.range(at: 5).location != NSNotFound { previousWord = text } + cursor = match.range.location + match.range.length + } + if cursor < ns.length { result.append(piece(ns.substring(from: cursor), plain)) } + return result + } + + private static func piece(_ text: String, _ color: Color) -> AttributedString { + var part = AttributedString(text) + part.foregroundColor = color + return part + } +} diff --git a/Sources/CodeWriterDemo/PythonRunner.swift b/Sources/CodeWriterDemo/PythonRunner.swift new file mode 100644 index 0000000..449cb79 --- /dev/null +++ b/Sources/CodeWriterDemo/PythonRunner.swift @@ -0,0 +1,111 @@ +import Foundation + +/// Runs model-written Python with the system `python3`: each assert separately, so the console can show which passed. +enum PythonRunner { + struct Outcome: Sendable { + /// One entry per assert, in order; nil when the code itself failed before the tests ran. + let checks: [(test: String, passed: Bool)] + /// Error text when the code did not load (syntax error, exception at import time, timeout). + let error: String? + let seconds: Double + + var passed: Bool { error == nil && !checks.isEmpty && checks.allSatisfy(\.passed) } + } + + /// The harness: load the code, then run every assert in its own try block and print `PASS`/`FAIL` per test. + static func harness(code: String, imports: [String], tests: [String]) -> String { + let quoted = tests.map { " " + pythonLiteral($0) + "," }.joined(separator: "\n") + return """ + import sys + _source = \(pythonLiteral((imports + [code]).joined(separator: "\n"))) + _tests = [ + \(quoted) + ] + try: + exec(compile(_source, "solution.py", "exec"), globals()) + except BaseException as error: + print("LOAD " + type(error).__name__ + ": " + str(error).splitlines()[0] if str(error) else "LOAD " + type(error).__name__) + sys.exit(0) + for _test in _tests: + try: + exec(_test, globals()) + print("PASS") + except BaseException as error: + print("FAIL " + type(error).__name__) + """ + } + + static func run(code: String, imports: [String], tests: [String], timeout: Double = 10) async -> Outcome { + let started = Date() + let script = harness(code: code, imports: imports, tests: tests) + let output: String + do { + output = try await execute(script, timeout: timeout) + } catch { + return Outcome(checks: [], error: error.localizedDescription, seconds: Date().timeIntervalSince(started)) + } + let lines = output.split(whereSeparator: \.isNewline).map(String.init) + if let load = lines.first(where: { $0.hasPrefix("LOAD ") }) { + return Outcome(checks: [], error: String(load.dropFirst(5)), seconds: Date().timeIntervalSince(started)) + } + let results = lines.filter { $0.hasPrefix("PASS") || $0.hasPrefix("FAIL") } + let checks = zip(tests, results).map { (test: $0, passed: $1 == "PASS") } + let error = checks.count == tests.count ? nil : "the tests did not finish" + return Outcome(checks: checks, error: error, seconds: Date().timeIntervalSince(started)) + } + + /// Whether `code` parses as Python (for tasks typed in by hand, which have no tests). + static func compiles(_ code: String) async -> String? { + let script = """ + import ast + try: + ast.parse(\(pythonLiteral(code))) + print("OK") + except SyntaxError as error: + print("SyntaxError line " + str(error.lineno) + ": " + str(error.msg)) + """ + let output = (try? await execute(script, timeout: 10)) ?? "python3 did not run" + let line = output.trimmingCharacters(in: .whitespacesAndNewlines) + return line == "OK" ? nil : line + } + + private static func execute(_ script: String, timeout: Double) async throws -> String { + let file = FileManager.default.temporaryDirectory.appendingPathComponent( + "python-writer-\(UUID().uuidString).py") + try script.write(to: file, atomically: true, encoding: .utf8) + defer { try? FileManager.default.removeItem(at: file) } + let process = Process() + process.executableURL = URL(fileURLWithPath: "/usr/bin/env") + process.arguments = ["python3", file.path] + let pipe = Pipe() + process.standardOutput = pipe + process.standardError = FileHandle.nullDevice + try process.run() + let deadline = Date().addingTimeInterval(timeout) + while process.isRunning { + if Date() > deadline { + process.terminate() + throw RunnerError.timeout(timeout) + } + try await Task.sleep(for: .milliseconds(20)) + } + return String(decoding: pipe.fileHandleForReading.readDataToEndOfFile(), as: UTF8.self) + } + + /// A Python string literal for arbitrary text (JSON string syntax is valid Python for these escapes). + static func pythonLiteral(_ text: String) -> String { + let data = (try? JSONSerialization.data(withJSONObject: [text], options: [.withoutEscapingSlashes])) ?? Data() + let array = String(decoding: data, as: UTF8.self) + return String(array.dropFirst().dropLast()) + } + + enum RunnerError: Error, LocalizedError { + case timeout(Double) + + var errorDescription: String? { + switch self { + case .timeout(let seconds): return "timed out after \(Int(seconds)) s" + } + } + } +} diff --git a/Sources/CodeWriterDemo/README.md b/Sources/CodeWriterDemo/README.md new file mode 100644 index 0000000..01284ef --- /dev/null +++ b/Sources/CodeWriterDemo/README.md @@ -0,0 +1,33 @@ +# Python writer + +A 0.5B model on this Mac writes Python from plain-English requests and the code runs: Qwen2.5-Coder-0.5B-Instruct +on Core ML ([qwen2.5-coder-0.5b-coreml](https://huggingface.co/FluidInference/qwen2.5-coder-0.5b-coreml)), the +prompt on the Neural Engine and the writing on the GPU. Nothing leaves the machine. + +```bash +swift run -c release CodeWriterDemo # downloads the model (~1 GB) on first launch +Sources/CodeWriterDemo/demo.sh --autostart # same app, plus one terminal: macmon above the live log +CODE_WRITER_MODEL_DIR= swift run -c release CodeWriterDemo # a local copy instead of the download +``` + +**Play** (⌘↩) works down a list of ten short tasks, one per Python feature (built-ins, list comprehension, slicing, +dict, string methods, conversion, tuples, loops, `sorted`, `all()`). Each one is written live into the editor, then +`python3` runs the task's three asserts and the console shows each result. **Pause** stops after the current task; +**Replay** (⌘R) starts over. Type your own request in the bar at the bottom: it is written the same way and checked +to parse as Python (it has no tests). Running the model's code needs `python3` on the `PATH`. + +The tasks come from MBPP (sanitized test split, CC-BY-4.0), with the first assert in the prompt so the model knows +the function name. The list is **hand-picked for the recording**: one original pick (sort a matrix by row sum) failed +and was replaced with one that passes. All ten pass; the unbiased numbers are below. + +M5 Pro, macOS 27, greedy decoding: + +| | | +|---|---| +| The ten demo tasks | 10/10 pass all asserts, ~50 s in total (95–330 tokens each) | +| Prompt (448 tokens, Neural Engine) | ~52 ms | +| Writing (GPU) | 36–51 tokens/s | +| MBPP sanitized test split (257) | 119/257 = 46.3% | +| HumanEval (164) | 89/164 = 54.3% (PyTorch fp32: 90/164 = 54.9%) | + +`CodeWriterCheck` replays HumanEval through the Swift host and compares every token stream with a reference run. diff --git a/Sources/CodeWriterDemo/Resources/tasks.json b/Sources/CodeWriterDemo/Resources/tasks.json new file mode 100644 index 0000000..b5a7c74 --- /dev/null +++ b/Sources/CodeWriterDemo/Resources/tasks.json @@ -0,0 +1,112 @@ +[ + { + "id": "Mbpp/62", + "topic": "built-ins", + "prompt": "Write a python function to find smallest number in a list.", + "tests": [ + "assert smallest_num([10, 20, 1, 45, 99]) == 1", + "assert smallest_num([1, 2, 3]) == 1", + "assert smallest_num([45, 46, 50, 60]) == 45" + ], + "imports": [] + }, + { + "id": "Mbpp/412", + "topic": "list comprehension", + "prompt": "Write a python function to remove odd numbers from a given list.", + "tests": [ + "assert remove_odd([1,2,3]) == [2]", + "assert remove_odd([2,4,6]) == [2,4,6]", + "assert remove_odd([10,20,3]) == [10,20]" + ], + "imports": [] + }, + { + "id": "Mbpp/456", + "topic": "slicing", + "prompt": "Write a function to reverse each string in a given list of string values.", + "tests": [ + "assert reverse_string_list(['Red', 'Green', 'Blue', 'White', 'Black'])==['deR', 'neerG', 'eulB', 'etihW', 'kcalB']", + "assert reverse_string_list(['john','amal','joel','george'])==['nhoj','lama','leoj','egroeg']", + "assert reverse_string_list(['jack','john','mary'])==['kcaj','nhoj','yram']" + ], + "imports": [] + }, + { + "id": "Mbpp/88", + "topic": "dict", + "prompt": "Write a function to get the frequency of all the elements in a list, returned as a dictionary.", + "tests": [ + "assert freq_count([10,10,10,10,20,20,20,20,40,40,50,50,30])==({10: 4, 20: 4, 40: 2, 50: 2, 30: 1})", + "assert freq_count([1,2,3,4,3,2,4,1,3,1,4])==({1:3, 2:2,3:3,4:3})", + "assert freq_count([5,6,7,4,9,10,4,5,6,7,9,5])==({10:1,5:3,6:2,7:2,4:2,9:2})" + ], + "imports": [] + }, + { + "id": "Mbpp/411", + "topic": "string methods", + "prompt": "Write a function to convert the given snake case string to camel case string.", + "tests": [ + "assert snake_to_camel('android_tv') == 'AndroidTv'", + "assert snake_to_camel('google_pixel') == 'GooglePixel'", + "assert snake_to_camel('apple_watch') == 'AppleWatch'" + ], + "imports": [] + }, + { + "id": "Mbpp/99", + "topic": "conversion", + "prompt": "Write a function to convert the given decimal number to its binary equivalent, represented as a string with no leading zeros.", + "tests": [ + "assert decimal_to_binary(8) == '1000'", + "assert decimal_to_binary(18) == '10010'", + "assert decimal_to_binary(7) == '111'" + ], + "imports": [] + }, + { + "id": "Mbpp/257", + "topic": "tuples", + "prompt": "Write a function that takes in two numbers and returns a tuple with the second number and then the first number.", + "tests": [ + "assert swap_numbers(10,20)==(20,10)", + "assert swap_numbers(15,17)==(17,15)", + "assert swap_numbers(100,200)==(200,100)" + ], + "imports": [] + }, + { + "id": "Mbpp/96", + "topic": "loops", + "prompt": "Write a python function to find the number of divisors of a given integer.", + "tests": [ + "assert divisor(15) == 4", + "assert divisor(12) == 6", + "assert divisor(9) == 3" + ], + "imports": [] + }, + { + "id": "Mbpp/397", + "topic": "sorted", + "prompt": "Write a function to find the median of three numbers.", + "tests": [ + "assert median_numbers(25,55,65)==55.0", + "assert median_numbers(20,10,30)==20.0", + "assert median_numbers(15,45,75)==45.0" + ], + "imports": [] + }, + { + "id": "Mbpp/115", + "topic": "all()", + "prompt": "Write a function to check whether all dictionaries in a list are empty or not.", + "tests": [ + "assert empty_dit([{},{},{}])==True", + "assert empty_dit([{1,2},{},{}])==False", + "assert empty_dit({})==True" + ], + "imports": [] + } +] diff --git a/Sources/CodeWriterDemo/WriterModel.swift b/Sources/CodeWriterDemo/WriterModel.swift new file mode 100644 index 0000000..50324b6 --- /dev/null +++ b/Sources/CodeWriterDemo/WriterModel.swift @@ -0,0 +1,266 @@ +import FluidUse +import Foundation +import SwiftUI + +/// One entry in the task list: a plain-English request, the Python feature it exercises, and its asserts. +struct CodingTask: Decodable, Identifiable { + let id: String + let topic: String + let prompt: String + let tests: [String] + let imports: [String] + + /// What the model sees: the request plus the first assert, so it knows the function name and signature. + var userMessage: String { "\(prompt)\nYour code should pass this test:\n\(tests[0])" } +} + +/// Plays the task list through Qwen2.5-Coder-0.5B on Core ML: each task is written live, then its asserts run with +/// `python3`. Pause stops after the current task; Replay starts over. +@available(macOS 15.0, *) +@MainActor +final class WriterModel: ObservableObject { + enum Step: Equatable { + case loading(String) + case idle + case writing + case finished + case failed(String) + } + + enum Status: Equatable { + case pending + case writing + case testing + case passed(Int, Int) + case failed(Int, Int) + case broken(String) + } + + struct Check: Identifiable { + let id: Int + let test: String + let passed: Bool + } + + @Published private(set) var step: Step = .loading("Starting…") + @Published private(set) var paused = false + @Published private(set) var statuses: [String: Status] = [:] + @Published private(set) var current: String? + @Published private(set) var code = "" + @Published private(set) var checks: [Check] = [] + @Published private(set) var runError: String? + @Published private(set) var tokensPerSecond: Double = 0 + @Published private(set) var lastSeconds: Double = 0 + @Published var customTask = "" + @Published private(set) var customTitle: String? + + let tasks: [CodingTask] + private var writer: CodeWriterManager? + private var showTask: Task? + + var passedCount: Int { statuses.values.filter { if case .passed = $0 { return true } else { return false } }.count } + var doneCount: Int { + statuses.values.filter { + switch $0 { + case .passed, .failed, .broken: return true + default: return false + } + }.count + } + var isRunning: Bool { step == .writing } + var writerMissing: Bool { writer == nil } + /// True while tokens are streaming into the editor (cursor shown). + @Published private(set) var isWriting = false + var canPlay: Bool { (step == .idle || step == .finished || paused) && writer != nil && !isBusyCustom } + @Published private(set) var isBusyCustom = false + + init() { + guard let url = Bundle.module.url(forResource: "tasks", withExtension: "json", subdirectory: "Resources"), + let data = try? Data(contentsOf: url), + let tasks = try? JSONDecoder().decode([CodingTask].self, from: data) + else { + self.tasks = [] + step = .failed("Resources/tasks.json is missing or invalid") + return + } + self.tasks = tasks + for task in tasks { statuses[task.id] = .pending } + } + + /// `CODE_WRITER_MODEL_DIR` or `--model=` for a local folder; otherwise the pinned Hugging Face snapshot. + static func localModelDirectory() -> URL? { + if let path = ProcessInfo.processInfo.environment["CODE_WRITER_MODEL_DIR"], !path.isEmpty { + return URL(fileURLWithPath: (path as NSString).expandingTildeInPath) + } + if let argument = CommandLine.arguments.first(where: { $0.hasPrefix("--model=") }) { + return URL(fileURLWithPath: (String(argument.dropFirst(8)) as NSString).expandingTildeInPath) + } + return nil + } + + func start() { + guard writer == nil, case .loading = step else { return } + Task { + do { + var directory = Self.localModelDirectory() + if directory == nil { + step = .loading("Downloading \(CodeWriterModelStore.repository) (about 1 GB, first launch only)…") + DemoLog.model("fetching \(CodeWriterModelStore.repository)") + directory = try await CodeWriterModelStore.ensure { file, bytes in + if bytes > 0 { DemoLog.model("downloaded \(file)") } + } + } + guard let directory else { return } + DemoLog.model("loading \(directory.path)") + step = .loading("Loading Qwen2.5-Coder 0.5B (first launch compiles for the Neural Engine)…") + let started = Date() + let writer = try await CodeWriterManager.load(from: directory) + step = .loading("Warming up…") + try await writer.warmUp() + self.writer = writer + DemoLog.model(String(format: "ready in %.1f s", Date().timeIntervalSince(started))) + step = .idle + if CommandLine.arguments.contains("--autostart") { play() } + } catch { + DemoLog.line("load failed: \(error.localizedDescription)", color: 196) + step = .failed(error.localizedDescription) + } + } + } + + func play() { + guard canPlay else { return } + paused = false + if step == .finished { reset() } + step = .writing + showTask = Task { await runShow() } + } + + func pause() { + guard isRunning else { return } + paused = true + DemoLog.event("⏸ pausing after this task") + } + + func replay() { + showTask?.cancel() + showTask = nil + reset() + paused = false + step = .idle + play() + } + + private func reset() { + for task in tasks { statuses[task.id] = .pending } + current = nil + code = "" + checks = [] + runError = nil + customTitle = nil + } + + private func runShow() async { + DemoLog.event("▶ writing \(tasks.count) tasks") + for task in tasks where statuses[task.id] == .pending { + if Task.isCancelled { return } + if paused { + step = .idle + return + } + await write(task) + } + guard !Task.isCancelled else { return } + step = .finished + DemoLog.event("■ \(passedCount)/\(tasks.count) tasks pass all their tests") + } + + private func write(_ task: CodingTask) async { + guard let writer else { return } + current = task.id + customTitle = nil + code = "" + checks = [] + runError = nil + statuses[task.id] = .writing + DemoLog.line("✎ \(task.topic): \(task.prompt)", color: 81) + do { + isWriting = true + defer { isWriting = false } + let completion = try await writer.write(task: task.userMessage) { text in + let shown = CodeWriterManager.extractCode(text) + Task { @MainActor in self.code = shown } + } + code = completion.code + tokensPerSecond = completion.timing.tokensPerSecond + lastSeconds = completion.timing.totalSeconds + DemoLog.model( + String( + format: "%d tokens in %.2f s (prefill %.0f ms on the Neural Engine, %.0f tok/s on the GPU)", + completion.timing.generatedTokens, completion.timing.totalSeconds, + completion.timing.prefillSeconds * 1000, completion.timing.tokensPerSecond)) + statuses[task.id] = .testing + let outcome = await PythonRunner.run(code: completion.code, imports: task.imports, tests: task.tests) + checks = outcome.checks.enumerated().map { + Check(id: $0.offset, test: $0.element.test, passed: $0.element.passed) + } + runError = outcome.error + let passed = outcome.checks.filter(\.passed).count + if let error = outcome.error, outcome.checks.isEmpty { + statuses[task.id] = .broken(error) + DemoLog.line(" ✗ \(error)", color: 196) + } else if outcome.passed { + statuses[task.id] = .passed(passed, task.tests.count) + DemoLog.line(" ✓ \(passed)/\(task.tests.count) tests pass", color: 46) + } else { + statuses[task.id] = .failed(passed, task.tests.count) + DemoLog.line(" ✗ \(passed)/\(task.tests.count) tests pass", color: 208) + } + } catch { + statuses[task.id] = .broken(error.localizedDescription) + runError = error.localizedDescription + DemoLog.line(" ✗ \(error.localizedDescription)", color: 196) + } + } + + /// A task typed by hand: written the same way; no tests, so the result is only checked to parse as Python. + func writeCustom() { + let prompt = customTask.trimmingCharacters(in: .whitespacesAndNewlines) + guard let writer, !prompt.isEmpty, !isRunning, !isBusyCustom else { return } + isBusyCustom = true + current = nil + customTitle = prompt + code = "" + checks = [] + runError = nil + DemoLog.line("✎ (typed) \(prompt)", color: 81) + Task { + defer { isBusyCustom = false } + do { + isWriting = true + let completion = try await writer.write(task: prompt) { text in + let shown = CodeWriterManager.extractCode(text) + Task { @MainActor in self.code = shown } + } + isWriting = false + code = completion.code + tokensPerSecond = completion.timing.tokensPerSecond + lastSeconds = completion.timing.totalSeconds + runError = await PythonRunner.compiles(completion.code) + DemoLog.line( + runError == nil ? " ✓ valid Python" : " ✗ \(runError!)", color: runError == nil ? 46 : 196) + } catch { + isWriting = false + runError = error.localizedDescription + } + } + } + + func select(_ task: CodingTask) { + guard !isRunning, !isBusyCustom, statuses[task.id] != .pending else { return } + // Re-show a finished task's result by writing it again (greedy, so the code is the same). + Task { + await write(task) + } + } +} diff --git a/Sources/CodeWriterDemo/demo.sh b/Sources/CodeWriterDemo/demo.sh new file mode 100755 index 0000000..25741c9 --- /dev/null +++ b/Sources/CodeWriterDemo/demo.sh @@ -0,0 +1,35 @@ +#!/bin/zsh +# Python writer demo: the app, plus one terminal (Ghostty if installed) with macmon (GPU / ANE / power; asitop as a +# fallback) on top and the live model log below. +# Sources/CodeWriterDemo/demo.sh [app options, e.g. --autostart] +# The model folder comes from CODE_WRITER_MODEL_DIR (default ~/Library/Application Support/FluidUse/qwen2.5-coder-0.5b-coreml). +set -e +ROOT=$(cd "$(dirname "$0")/../.." && pwd) +LOG=${TMPDIR:-/tmp}/code-writer-demo.log +# macmon (brew install macmon) needs no sudo and works on M5; asitop 0.0.24 crashes there (KeyError E0-Cluster_active). +if command -v macmon >/dev/null; then + MONITOR="$(command -v macmon)" +else + MONITOR="sudo sh -c 'pkill -x powermetrics; exec $(command -v asitop || echo "$HOME/.local/bin/asitop")'" +fi + +swift build -c release --product CodeWriterDemo --package-path "$ROOT" +: > "$LOG" +# One tmux session, reused across launches: the monitor on top, the log below. +if ! tmux has-session -t code-writer 2>/dev/null; then + # The pane stays open if the monitor exits, so the reason is visible. + tmux new-session -d -s code-writer -x 160 -y 70 \ + "$MONITOR; echo; echo \"monitor exited (\$?) — press Enter to retry\"; read; $MONITOR" + tmux set-option -t code-writer remain-on-exit on + tmux split-window -v -t code-writer "tail -n 300 -F '$LOG'" + if [[ -d /Applications/Ghostty.app ]]; then + open -na Ghostty --args -e "$(command -v tmux)" attach -t code-writer + else + osascript -e 'tell application "Terminal" to do script "tmux attach -t code-writer"' \ + -e 'tell application "Terminal" to activate' >/dev/null + fi +fi +pkill -x CodeWriterDemo 2>/dev/null || true +cd "$ROOT" +"$ROOT/.build/release/CodeWriterDemo" "$@" > "$LOG" 2> "$LOG.stderr" & +echo "demo pid $! · log $LOG · Play (⌘↩) · Replay (⌘R)" diff --git a/Sources/FluidUse/CodeWriter/CodeWriterManager.swift b/Sources/FluidUse/CodeWriter/CodeWriterManager.swift new file mode 100644 index 0000000..65339d2 --- /dev/null +++ b/Sources/FluidUse/CodeWriter/CodeWriterManager.swift @@ -0,0 +1,302 @@ +import CoreML +import Foundation +import os + +/// Writes Python from a plain-English task with Qwen2.5-Coder-0.5B-Instruct on Core ML. +/// +/// Same graph split as ``ShortReplyManager``: one package holds `prefill` (stateless, fixed prompt length, Neural +/// Engine) and `decode` (stateful KV cache, GPU) over one set of weights. The host tokenizes the Qwen2.5 chat prompt, +/// left-pads it to the prefill length, copies the prefill K/V into the decoder state and decodes greedily until +/// `<|im_end|>`, streaming the text as it grows. +@available(macOS 15.0, iOS 18.0, *) +public actor CodeWriterManager { + public struct Timing: Sendable { + public let promptTokens: Int + public let prefillSeconds: Double + public let decodeSeconds: Double + public let generatedTokens: Int + public var totalSeconds: Double { prefillSeconds + decodeSeconds } + public var tokensPerSecond: Double { decodeSeconds > 0 ? Double(generatedTokens) / decodeSeconds : 0 } + } + + public struct Completion: Sendable { + /// Everything the model wrote (markdown, possibly with prose around the code). + public let text: String + /// The first ```python block, or the whole text when there is none. + public let code: String + /// Generated token ids, including the stop token when one was produced. + public let tokens: [Int] + public let timing: Timing + } + + /// `config.json` next to the package. + struct Config: Decodable { + let package: String + let prefillLength: Int + let cacheLength: Int + let layers: Int + let kvHeads: Int + let headDim: Int + let padID: Int + let stopIDs: [Int] + let maxNewTokens: Int + let systemPrompt: String + } + + private let logger = Logger(subsystem: "FluidUse", category: "CodeWriter") + private let config: Config + private let tokenizer: QwenBPETokenizer + private let prefill: MLModel + private let decode: MLModel + private let stopIDs: Set + + /// Loads `/config.json`, the package it names, and `tokenizer.json`. + public static func load( + from directory: URL, prefillUnits: MLComputeUnits = .cpuAndNeuralEngine, + decodeUnits: MLComputeUnits = .cpuAndGPU + ) async throws -> CodeWriterManager { + let config = try JSONDecoder().decode( + Config.self, from: Data(contentsOf: directory.appendingPathComponent("config.json"))) + var package = directory.appendingPathComponent(config.package) + if package.pathExtension == "mlpackage" { + package = try await KevManager.compiled(package) // compile once, keep the .mlmodelc beside the package + } + return try CodeWriterManager( + config: config, compiled: package, tokenizer: directory.appendingPathComponent("tokenizer.json"), + prefillUnits: prefillUnits, decodeUnits: decodeUnits) + } + + /// Models are created here, inside the actor, so the non-Sendable `MLModel`s never cross an isolation boundary. + init( + config: Config, compiled: URL, tokenizer: URL, prefillUnits: MLComputeUnits, decodeUnits: MLComputeUnits + ) throws { + self.config = config + self.tokenizer = try QwenBPETokenizer(tokenizerJsonURL: tokenizer) + let prefillConfiguration = MLModelConfiguration() + prefillConfiguration.computeUnits = prefillUnits + prefillConfiguration.functionName = "prefill" + let decodeConfiguration = MLModelConfiguration() + decodeConfiguration.computeUnits = decodeUnits + decodeConfiguration.functionName = "decode" + self.prefill = try MLModel(contentsOf: compiled, configuration: prefillConfiguration) + self.decode = try MLModel(contentsOf: compiled, configuration: decodeConfiguration) + self.stopIDs = Set(config.stopIDs) + } + + /// One short task through both functions, so the first real call pays no compile cost. + public func warmUp() async throws { + _ = try await write(task: "Write a Python function that returns 1.", maxNewTokens: 8) + } + + /// The Qwen2.5 chat prompt (system turn from `config.json` unless overridden). + func promptTokens(user: String, system: String?) throws -> [Int] { + let text = + "<|im_start|>system\n\(system ?? config.systemPrompt)<|im_end|>\n" + + "<|im_start|>user\n\(user)<|im_end|>\n<|im_start|>assistant\n" + return try tokenizer.encode(text) + } + + /// Greedy completion of `task`. `onText` receives the full text written so far after each token (skipping + /// states that end mid-character). + public func write( + task: String, system: String? = nil, maxNewTokens: Int? = nil, + onText: (@Sendable (String) -> Void)? = nil + ) async throws -> Completion { + let prompt = try promptTokens(user: task, system: system) + let length = config.prefillLength + guard !task.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { throw CodeWriterError.emptyTask } + guard prompt.count <= length else { throw CodeWriterError.taskTooLong(prompt.count, length) } + // Left-pad to the fixed prefill length; padded cache slots stay masked for every later query. + let pad = length - prompt.count + let ids = Array(repeating: config.padID, count: pad) + prompt + + let started = ContinuousClock.now + let state = decode.makeState() + let prefillOutput = try Self.predict(prefill, try prefillInputs(ids: ids, pad: pad), state: nil) + try copyCache(from: prefillOutput, into: state, promptLength: length) + var logits = try lastLogits(prefillOutput) + let prefillSeconds = seconds(since: started) + + let decodeStarted = ContinuousClock.now + let limit = min(maxNewTokens ?? config.maxNewTokens, config.cacheLength - length) + var generated: [Int] = [] + var position = length + let mask = try MLMultiArray(shape: [1, 1, 1, NSNumber(value: config.cacheLength)], dataType: .float32) + mask.withUnsafeMutableBytes { buffer, _ in + let values = buffer.bindMemory(to: Float32.self) + for column in 0.. String { + guard let open = text.range(of: #"```[A-Za-z]*\n"#, options: .regularExpression) else { + return text.trimmingCharacters(in: .whitespacesAndNewlines) + } + let rest = text[open.upperBound...] + let body = rest.range(of: "```").map { rest[..<$0.lowerBound] } ?? rest + return String(body).trimmingCharacters(in: .newlines) + } + + /// Synchronous prediction (an async context would pick Core ML's async overload). Called only from the actor, + /// which serializes it: Core ML's synchronous prediction is not thread-safe across callers. + private static func predict( + _ model: MLModel, _ input: MLFeatureProvider, state: MLState? + ) throws -> MLFeatureProvider { + if let state { return try model.prediction(from: input, using: state) } + return try model.prediction(from: input) + } + + // MARK: - Feature providers + + private func prefillInputs(ids: [Int], pad: Int) throws -> MLDictionaryFeatureProvider { + let length = ids.count + let inputIDs = try MLMultiArray(shape: [1, NSNumber(value: length)], dataType: .int32) + for (index, id) in ids.enumerated() { inputIDs[index] = NSNumber(value: Int32(id)) } + // Causal mask over the prompt; padded columns masked for every row. + let mask = try MLMultiArray( + shape: [1, 1, NSNumber(value: length), NSNumber(value: length)], dataType: .float32) + mask.withUnsafeMutableBytes { buffer, _ in + let values = buffer.bindMemory(to: Float32.self) + for row in 0.. row || column < pad) ? -1e4 : 0 + } + } + } + return try MLDictionaryFeatureProvider(dictionary: ["input_ids": inputIDs, "attention_mask": mask]) + } + + private func decodeInputs(token: Int, position: Int, mask: MLMultiArray) throws -> MLDictionaryFeatureProvider { + let inputIDs = try MLMultiArray(shape: [1, 1], dataType: .int32) + inputIDs[0] = NSNumber(value: Int32(token)) + let cachePosition = try MLMultiArray(shape: [1], dataType: .int32) + cachePosition[0] = NSNumber(value: Int32(position)) + return try MLDictionaryFeatureProvider(dictionary: [ + "input_ids": inputIDs, "attention_mask": mask, "cache_position": cachePosition, + ]) + } + + /// Prefill returns `keys` / `values` as [layers, kvHeads, prompt, headDim] fp16; the decoder state holds + /// [1, kvHeads, cache, headDim] fp16 per layer. + private func copyCache(from output: MLFeatureProvider, into state: MLState, promptLength: Int) throws { + guard let keys = output.featureValue(for: "keys")?.multiArrayValue, + let values = output.featureValue(for: "values")?.multiArrayValue + else { throw CodeWriterError.invalidAsset("prefill output has no keys/values") } + let headDim = config.headDim + let kvHeads = config.kvHeads + let rowElements = promptLength * headDim + for (name, source) in [("k", keys), ("v", values)] { + // Prefill outputs are fp16 and may carry padded strides; copy head by head using the array's own strides. + guard source.dataType == .float16, source.shape.count == 4, + source.shape[1].intValue == kvHeads, source.shape[2].intValue == promptLength, + source.shape[3].intValue == headDim, source.strides[3].intValue == 1, + source.strides[2].intValue == headDim + else { + throw CodeWriterError.invalidAsset( + "unexpected prefill \(name) layout \(source.shape) / \(source.strides)") + } + let layerStride = source.strides[0].intValue + let headStride = source.strides[1].intValue + for layer in 0...size) + } + } + } + } + } + } + } + + private func lastLogits(_ output: MLFeatureProvider) throws -> MLMultiArray { + guard let logits = output.featureValue(for: "logits")?.multiArrayValue else { + throw CodeWriterError.invalidAsset("model output has no logits") + } + return logits + } + + private func argmax(_ logits: MLMultiArray) -> Int { + let count = logits.count + return logits.withUnsafeBytes { buffer in + switch logits.dataType { + case .float32: + let values = buffer.bindMemory(to: Float32.self) + var best = 0 + for index in 1.. values[best] { best = index } + return best + default: + let values = buffer.bindMemory(to: Float16.self) + var best = 0 + for index in 1.. values[best] { best = index } + return best + } + } + } + + private func seconds(since start: ContinuousClock.Instant) -> Double { + let duration = start.duration(to: .now) + return Double(duration.components.seconds) + Double(duration.components.attoseconds) / 1e18 + } +} + +public enum CodeWriterError: Error, LocalizedError, Sendable { + case emptyTask + case taskTooLong(Int, Int) + case invalidAsset(String) + + public var errorDescription: String? { + switch self { + case .emptyTask: return "Describe what the code should do." + case .taskTooLong(let tokens, let limit): return "The task is too long (\(tokens) of \(limit) prompt tokens)." + case .invalidAsset(let detail): return "Code-writer model asset problem: \(detail)" + } + } +} diff --git a/Sources/FluidUse/CodeWriter/CodeWriterModelStore.swift b/Sources/FluidUse/CodeWriter/CodeWriterModelStore.swift new file mode 100644 index 0000000..6f75927 --- /dev/null +++ b/Sources/FluidUse/CodeWriter/CodeWriterModelStore.swift @@ -0,0 +1,82 @@ +import CryptoKit +import Foundation + +/// Downloads the pinned Qwen2.5-Coder-0.5B Core ML snapshot (FluidInference/qwen2.5-coder-0.5b-coreml): the +/// multifunction package, the tokenizer and the host config, laid out as `CodeWriterManager.load(from:)` expects. +public enum CodeWriterModelStore { + public typealias Progress = @Sendable (_ file: String, _ bytes: Int64) -> Void + + struct Asset { + let path: String + let sha256: String + } + + public static let repository = "FluidInference/qwen2.5-coder-0.5b-coreml" + static let revision = "48abc2c5b3ae57f39b0c7b9d884345389b298097" + static let assets: [Asset] = [ + Asset(path: "config.json", sha256: "36b8b8eba9f2285e72c6560ce1b3f660534d2e3d780c910cd04aadf84871b553"), + Asset(path: "tokenizer.json", sha256: "c0382117ea329cdf097041132f6d735924b697924d6f6fc3945713e96ce87539"), + Asset( + path: "qwen2_5_coder_0_5b.mlpackage/Manifest.json", + sha256: "b50d536c73db6ed30cfcaacd3a92bfda794e6b7d501f0e9c119c840b3adcee0e"), + Asset( + path: "qwen2_5_coder_0_5b.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "e70fe1ecdc5fb13aa62d31ffd9c9d4f52c9aafa0846d34efdb25758e0be833b8"), + Asset( + path: "qwen2_5_coder_0_5b.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "3cca098e86bc97aeca18a7593f32869a937e6d09bcae7ec15ac83e5333d0e11b"), + ] + + /// Ensure the snapshot exists in the FluidUse cache and return its directory. Files are checksummed once per + /// pinned revision; later launches only check that they are present. + public static func ensure(cacheDirectory: URL? = nil, progress: Progress? = nil) async throws -> URL { + let root = cacheDirectory ?? LayaModelStore.defaultCacheDirectory() + let directory = root.appendingPathComponent("qwen2.5-coder-0.5b-coreml", isDirectory: true) + let manager = FileManager.default + let verified = directory.appendingPathComponent(".verified-\(revision)") + if manager.fileExists(atPath: verified.path), + assets.allSatisfy({ manager.fileExists(atPath: directory.appendingPathComponent($0.path).path) }) + { + return directory + } + try manager.createDirectory(at: directory, withIntermediateDirectories: true) + for asset in assets { + try Task.checkCancellation() + let destination = directory.appendingPathComponent(asset.path) + if manager.fileExists(atPath: destination.path), try checksum(of: destination) == asset.sha256 { + continue + } + try manager.createDirectory(at: destination.deletingLastPathComponent(), withIntermediateDirectories: true) + progress?(asset.path, 0) + let escaped = asset.path.addingPercentEncoding(withAllowedCharacters: .urlPathAllowed) ?? asset.path + guard let url = URL(string: "https://huggingface.co/\(repository)/resolve/\(revision)/\(escaped)") else { + throw CodeWriterError.invalidAsset("Invalid Hugging Face asset URL for \(asset.path)") + } + let (temporary, response) = try await URLSession.shared.download(from: url) + defer { try? manager.removeItem(at: temporary) } + guard let http = response as? HTTPURLResponse, http.statusCode == 200 else { + throw CodeWriterError.invalidAsset("Download failed for \(asset.path)") + } + let actual = try checksum(of: temporary) + guard actual == asset.sha256 else { + throw CodeWriterError.invalidAsset( + "Checksum mismatch for \(asset.path): expected \(asset.sha256), got \(actual)") + } + let size = (try manager.attributesOfItem(atPath: temporary.path)[.size] as? NSNumber)?.int64Value ?? 0 + try LayaModelStore.installDownloadedFile(temporary, at: destination) + progress?(asset.path, size) + } + try Data().write(to: verified) + return directory + } + + private static func checksum(of file: URL) throws -> String { + let handle = try FileHandle(forReadingFrom: file) + defer { try? handle.close() } + var digest = SHA256() + while let chunk = try handle.read(upToCount: 4_194_304), !chunk.isEmpty { + digest.update(data: chunk) + } + return digest.finalize().map { String(format: "%02x", $0) }.joined() + } +} diff --git a/Tests/FluidUseTests/CodeWriterTests.swift b/Tests/FluidUseTests/CodeWriterTests.swift new file mode 100644 index 0000000..bf7515a --- /dev/null +++ b/Tests/FluidUseTests/CodeWriterTests.swift @@ -0,0 +1,41 @@ +import XCTest + +@testable import FluidUse + +final class CodeWriterTests: XCTestCase { + func testExtractCodeTakesTheFirstFencedBlock() throws { + guard #available(macOS 15.0, iOS 18.0, *) else { throw XCTSkip("CodeWriterManager needs macOS 15") } + let text = + "Here you go:\n\n```python\ndef f(x):\n return x[::-1]\n```\n\nIt reverses.\n```python\nprint(1)\n```" + XCTAssertEqual(CodeWriterManager.extractCode(text), "def f(x):\n return x[::-1]") + } + + func testExtractCodeKeepsAnUnclosedBlockAndPlainText() throws { + guard #available(macOS 15.0, iOS 18.0, *) else { throw XCTSkip("CodeWriterManager needs macOS 15") } + // Mid-stream (or cut off at the token limit): no closing fence yet. + XCTAssertEqual(CodeWriterManager.extractCode("```python\ndef f():\n retu"), "def f():\n retu") + XCTAssertEqual(CodeWriterManager.extractCode("```\nx = 1\n```"), "x = 1") + XCTAssertEqual(CodeWriterManager.extractCode(" def g(): pass\n"), "def g(): pass") + } + + /// The Swift prompt must tokenize exactly as `apply_chat_template(..., add_generation_prompt=True)` in + /// transformers for Qwen/Qwen2.5-Coder-0.5B-Instruct. + func testPromptTokensMatchPythonReference() throws { + guard let directory = ProcessInfo.processInfo.environment["CODE_WRITER_MODEL_DIR"], !directory.isEmpty else { + throw XCTSkip("Set CODE_WRITER_MODEL_DIR (folder with tokenizer.json) to run") + } + let tokenizer = try QwenBPETokenizer( + tokenizerJsonURL: URL(fileURLWithPath: directory).appendingPathComponent("tokenizer.json")) + let text = + "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n" + + "<|im_start|>user\nWrite a function to reverse each string in a list.\nYour code should pass this test:\n" + + "assert rev([\"ab\"]) == [\"ba\"]<|im_end|>\n<|im_start|>assistant\n" + let reference = [ + 151644, 8948, 198, 2610, 525, 1207, 16948, 11, 3465, 553, 54364, 14817, 13, 1446, 525, 264, 10950, 17847, + 13, + 151645, 198, 151644, 872, 198, 7985, 264, 729, 311, 9931, 1817, 914, 304, 264, 1140, 624, 7771, 2038, 1265, + 1494, 419, 1273, 510, 2207, 5772, 19065, 370, 14013, 621, 4383, 4645, 1341, 151645, 198, 151644, 77091, 198, + ] + XCTAssertEqual(try tokenizer.encode(text), reference) + } +}