diff --git a/Package.swift b/Package.swift index 13b14ea..49c9fe6 100644 --- a/Package.swift +++ b/Package.swift @@ -67,6 +67,8 @@ let package = Package( .executableTarget(name: "ImageSortDemo", dependencies: ["ImageSort"], exclude: ["README.md"]), .executableTarget(name: "KevCheck", dependencies: ["FluidUse"]), .executableTarget(name: "ClefVisionCheck", dependencies: ["FluidUse"]), + .executableTarget(name: "UnlimitedOcrCheck", dependencies: ["FluidUse"]), + .executableTarget(name: "UnlimitedOcrDemo", dependencies: ["FluidUse"], exclude: ["README.md"]), .executableTarget( name: "ClefFlashDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh", "bars.py"]), .executableTarget(name: "ClefTextCheck", dependencies: ["FluidUse"]), diff --git a/README.md b/README.md index bb34ddb..9404c21 100644 --- a/README.md +++ b/README.md @@ -332,6 +332,27 @@ 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)). +## Document OCR (multi-page) + +`UnlimitedOcrManager` runs Baidu's [Unlimited-OCR](https://huggingface.co/FluidInference/unlimited-ocr-coreml) +(DeepSeek-OCR architecture, 3B MoE / ~0.5B active, MIT) on the GPU through Core ML: up to 8 document pages go in as +one prompt and the whole document comes out of a single generation as markdown with layout boxes. Its R-SWA cache +(full prefix + a ring of the last 128 generated tokens) is Core ML state, so speed stays flat over a long read. The +pinned, checksummed bundle (~3.8 GB) downloads from the Hub on first use (macOS 15 / iOS 18). + +```swift +let ocr = try await UnlimitedOcrManager.load(from: try await UnlimitedOcrModelStore.ensure()) +let pages = try (0..` sections of `<|det|>label [x1, y1, x2, y2]<|/det|>content` blocks. Boxes are in +/// upstream's 0...999 page coordinates (top-left origin). Parsing is tolerant of a partially streamed tail. +public struct UnlimitedOcrDocument: Sendable { + public struct Block: Sendable, Identifiable { + public let id: Int + public let label: String + /// Normalized 0...1 rectangles (top-left origin); usually one. + public let boxes: [CGRect] + public let content: String + } + + public struct Page: Sendable, Identifiable { + public let id: Int + public let blocks: [Block] + } + + public let pages: [Page] + + public init(parsing text: String) { + var pages: [Page] = [] + var sections = text.components(separatedBy: "") + if sections.first?.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty == true { sections.removeFirst() } + var blockID = 0 + for (pageIndex, section) in sections.enumerated() { + var blocks: [Block] = [] + let parts = section.components(separatedBy: "<|det|>") + for part in parts.dropFirst() { + let pieces = part.components(separatedBy: "<|/det|>") + let header = pieces[0] + let content = pieces.count > 1 ? pieces[1...].joined(separator: "") : "" + let label = + header.split(separator: "[", maxSplits: 1).first.map { + $0.trimmingCharacters(in: .whitespaces) + } ?? "" + blocks.append( + Block( + id: blockID, label: label, boxes: Self.boxes(in: header), + content: content.trimmingCharacters(in: .whitespacesAndNewlines))) + blockID += 1 + } + if blocks.isEmpty, !section.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty { + blocks.append( + Block( + id: blockID, label: "text", boxes: [], + content: section.trimmingCharacters(in: .whitespacesAndNewlines))) + blockID += 1 + } + pages.append(Page(id: pageIndex, blocks: blocks)) + } + self.pages = pages + } + + /// Every complete `[x1, y1, x2, y2]` group in a det header, normalized. + static func boxes(in header: String) -> [CGRect] { + var rects: [CGRect] = [] + var numbers: [Double] = [] + var current = "" + var inside = false + for character in header { + switch character { + case "[": + inside = true + numbers = [] + current = "" + case "]": + if let value = Double(current.trimmingCharacters(in: .whitespaces)) { numbers.append(value) } + if inside, numbers.count == 4 { + rects.append( + CGRect( + x: numbers[0] / 999, y: numbers[1] / 999, width: (numbers[2] - numbers[0]) / 999, + height: (numbers[3] - numbers[1]) / 999)) + } + inside = false + current = "" + case ",": + if let value = Double(current.trimmingCharacters(in: .whitespaces)) { numbers.append(value) } + current = "" + default: + if inside { current.append(character) } + } + } + return rects + } + + /// Plain markdown: layout markup removed, pages separated by a rule. + public var markdown: String { + pages.map { page in page.blocks.map(\.content).filter { !$0.isEmpty }.joined(separator: "\n\n") } + .joined(separator: "\n\n---\n\n") + } +} diff --git a/Sources/FluidUse/UnlimitedOcr/UnlimitedOcrManager.swift b/Sources/FluidUse/UnlimitedOcr/UnlimitedOcrManager.swift new file mode 100644 index 0000000..23c62d3 --- /dev/null +++ b/Sources/FluidUse/UnlimitedOcr/UnlimitedOcrManager.swift @@ -0,0 +1,430 @@ +@preconcurrency import CoreML +import Foundation +import os + +/// Baidu Unlimited-OCR (3B MoE, ~0.5B active) on Core ML: several document pages in, structured markdown out, in one +/// generation (upstream's "Multi page parsing." mode). +/// +/// - Page encoder (SAM-B + CLIP-L + projector, fp16, GPU): one 1024 x 1024 page -> 273 embeddings. +/// - Decoder bundle (8-bit, GPU): per prefix-capacity bucket P a `prefill_P` (128 tokens per call) and a `decode_P` +/// (1 token) function sharing weights. The K/V cache is Core ML state: P prefix slots + a 128-slot ring for +/// generated tokens (upstream's R-SWA keeps the whole prefix plus the last 128 generated tokens). +/// - Host: embeddings, RoPE, masks, ring slots and upstream's 35-gram no-repeat ban. +/// +/// An actor: Core ML's synchronous `prediction` is not thread-safe. +@available(macOS 15.0, iOS 18.0, *) +public actor UnlimitedOcrManager { + private struct Bucket { + let capacity: Int + let prefill: MLModel + let decode: MLModel + } + + let config: UnlimitedOcrConfig + private let directory: URL + private let computeUnits: MLComputeUnits + private let encoder: MLModel + private let embeddings: Data // fp16 [vocab, hidden] + private let tokenDecoder: UnlimitedOcrTokenDecoder + private var buckets: [Int: Bucket] = [:] + /// Read by the generation loop every token; a lock rather than actor state because the loop holds the actor + /// without suspending, so an actor-isolated `pause()` would only run after the read ends. + private nonisolated let pauseFlag = OSAllocatedUnfairLock(initialState: false) + + /// Largest document, in pages, the bundle's biggest capacity bucket holds. + public nonisolated let maxPages: Int + + public static func load( + from directory: URL, computeUnits: MLComputeUnits = .cpuAndGPU + ) async throws + -> UnlimitedOcrManager + { + let config = try JSONDecoder().decode( + UnlimitedOcrConfig.self, from: Data(contentsOf: directory.appendingPathComponent("config.json"))) + let embeddings = try Data( + contentsOf: directory.appendingPathComponent("embeddings.f16"), options: .alwaysMapped) + guard embeddings.count == config.vocabSize * config.hiddenSize * 2 else { + throw UnlimitedOcrError.invalidAsset("embeddings.f16 size") + } + let tokenDecoder = try UnlimitedOcrTokenDecoder( + tokenizerJsonURL: directory.appendingPathComponent("tokenizer.json"), + skipping: [config.bosId, config.eosId]) + let encoder = try await compiledModel( + directory: directory, package: config.encoder, functionName: nil, computeUnits: computeUnits) + return UnlimitedOcrManager( + config: config, directory: directory, computeUnits: computeUnits, encoder: encoder, + embeddings: embeddings, tokenDecoder: tokenDecoder) + } + + private init( + config: UnlimitedOcrConfig, directory: URL, computeUnits: MLComputeUnits, encoder: MLModel, embeddings: Data, + tokenDecoder: UnlimitedOcrTokenDecoder + ) { + self.config = config + self.directory = directory + self.computeUnits = computeUnits + self.encoder = encoder + self.embeddings = embeddings + self.tokenDecoder = tokenDecoder + maxPages = config.maxPages + } + + /// Compiles a package once into `compiled/` and loads it (optionally one function of a multifunction package). + private static func compiledModel( + directory: URL, package: String, functionName: String?, computeUnits: MLComputeUnits + ) async throws -> MLModel { + let source = directory.appendingPathComponent(package) + let compiledDir = directory.appendingPathComponent("compiled") + try FileManager.default.createDirectory(at: compiledDir, withIntermediateDirectories: true) + let target = compiledDir.appendingPathComponent(source.deletingPathExtension().lastPathComponent + ".mlmodelc") + if !FileManager.default.fileExists(atPath: target.path) { + let temporary = try await MLModel.compileModel(at: source) + try FileManager.default.moveItem(at: temporary, to: target) + } + let configuration = MLModelConfiguration() + configuration.computeUnits = computeUnits + configuration.functionName = functionName + return try await MLModel.load(contentsOf: target, configuration: configuration) + } + + /// Loads the decoder functions for the buckets these page counts need (default: all) and runs one prediction + /// through each model, so the first real read does not pay Core ML's lazy GPU setup. + public func warm(pages: [Int]? = nil) async throws { + let counts = pages ?? Array(1...max(1, maxPages)) // [] warms the page encoder only + let blank = UnlimitedOcrPage( + pixels: [Float]( + repeating: 0, count: 3 * UnlimitedOcrPagePreprocessor.size * UnlimitedOcrPagePreprocessor.size)) + _ = try encode(blank) + var warmed = Set() + for count in counts { + guard let capacity = config.capacity(forPrefix: config.prefixLength(pages: count, prompt: "multi")), + warmed.insert(capacity).inserted + else { continue } + let bucket = try await bucket(capacity) + let slots = capacity + config.ring + for (model, c) in [(bucket.prefill, config.chunk), (bucket.decode, 1)] { + let x = try MLMultiArray(shape: [1, c as NSNumber, config.hiddenSize as NSNumber], dataType: .float16) + let mask = try MLMultiArray(shape: [1, 1, c as NSNumber, (slots + c) as NSNumber], dataType: .float16) + let writeIdx = try MLMultiArray(shape: [c as NSNumber], dataType: .int32) + let outIndex = try MLMultiArray(shape: [1], dataType: .int32) + // MLMultiArray memory is not initialized. + x.dataPointer.initializeMemory(as: UInt8.self, repeating: 0, count: x.count * 2) + mask.dataPointer.initializeMemory(as: UInt8.self, repeating: 0, count: mask.count * 2) + for i in 0.. Bucket { + if let bucket = buckets[capacity] { return bucket } + let prefill = try await Self.compiledModel( + directory: directory, package: config.decoder, functionName: "prefill_P\(capacity)", + computeUnits: computeUnits) + let decode = try await Self.compiledModel( + directory: directory, package: config.decoder, functionName: "decode_P\(capacity)", + computeUnits: computeUnits) + let bucket = Bucket(capacity: capacity, prefill: prefill, decode: decode) + buckets[capacity] = bucket + return bucket + } + + /// Reads `pages` (from `UnlimitedOcrPagePreprocessor`) as one document. With one page the "single" prompt + /// (`document parsing.`) can be used; the default multi-page prompt marks each page with ``. + public nonisolated func read( + pages: [UnlimitedOcrPage], prompt: String = "multi", maxNewTokens: Int? = nil + ) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + let task = Task { + do { + try await self.run( + pages: pages, prompt: prompt, maxNewTokens: maxNewTokens, + emit: { continuation.yield($0) }) + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + continuation.onTermination = { _ in task.cancel() } + } + } + + /// Holds a running read between tokens until `resume()`; the K/V state and position are kept. + public nonisolated func pause() { pauseFlag.withLock { $0 = true } } + + public nonisolated func resume() { pauseFlag.withLock { $0 = false } } + + public nonisolated var isPaused: Bool { pauseFlag.withLock { $0 } } + + // MARK: - Generation + + private func run( + pages: [UnlimitedOcrPage], prompt: String, maxNewTokens: Int?, emit: @Sendable (UnlimitedOcrEvent) -> Void + ) async throws { + guard !pages.isEmpty else { throw UnlimitedOcrError.invalidInput("no pages") } + resume() + guard let suffix = config.promptSuffixIds[prompt] else { + throw UnlimitedOcrError.invalidInput("unknown prompt \(prompt)") + } + let hidden = config.hiddenSize + let n = config.prefixLength(pages: pages.count, prompt: prompt) + guard let capacity = config.capacity(forPrefix: n) else { + throw UnlimitedOcrError.invalidInput( + "\(pages.count) pages need \(n) prefix slots; bundle holds \(maxPages) pages") + } + let bucket = try await bucket(capacity) + + // Prefix: BOS, the page embeddings, the prompt text. + var promptIDs = [config.bosId] + var prefix = [Float16](repeating: 0, count: n * hidden) + copyEmbedding(config.bosId, into: &prefix, row: 0) + let encodeStart = Date() + for (index, page) in pages.enumerated() { + let pageStart = Date() + let embeds = try encode(page) + try copyRows(embeds, rows: config.pageTokens, into: &prefix, startRow: 1 + index * config.pageTokens) + promptIDs += Array(repeating: config.imageTokenId, count: config.pageTokens) + emit(.pageEncoded(index: index, milliseconds: Date().timeIntervalSince(pageStart) * 1000)) + try Task.checkCancellation() + } + let encodeMs = Date().timeIntervalSince(encodeStart) * 1000 + for (i, id) in suffix.enumerated() { copyEmbedding(id, into: &prefix, row: n - suffix.count + i) } + promptIDs += suffix + + // Prefill in chunks, then hand the cache to the decode function's state. + let slots = capacity + config.ring + let chunk = config.chunk + let prefillStart = Date() + let prefillState = bucket.prefill.makeState() + var top: (values: [Float], ids: [Int]) = ([], []) + let x = try MLMultiArray(shape: [1, chunk as NSNumber, hidden as NSNumber], dataType: .float16) + let mask = try MLMultiArray(shape: [1, 1, chunk as NSNumber, (slots + chunk) as NSNumber], dataType: .float16) + let writeIdx = try MLMultiArray(shape: [chunk as NSNumber], dataType: .int32) + let outIndex = try MLMultiArray(shape: [1], dataType: .int32) + for start in stride(from: 0, to: n, by: chunk) { + let rows = min(chunk, n - start) + let xp = x.dataPointer.assumingMemoryBound(to: Float16.self) + prefix.withUnsafeBufferPointer { src in + xp.update(from: src.baseAddress! + start * hidden, count: rows * hidden) + } + (xp + rows * hidden).update(repeating: 0, count: (chunk - rows) * hidden) + let mp = mask.dataPointer.assumingMemoryBound(to: Float16.self) + let width = slots + chunk + for i in 0..