Something went wrong. Try again.
Routing embedding stack for Qwen3-Embedding-0.6B on Apple Silicon: orchestrator + MLX + llama.cpp backends
Something went wrong. Try again.
2.8 kB · 80 lines
Swift
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081import Foundationimport MLXimport MLXEmbeddersimport MLXLMCommon
/// Everything that runs inside `ModelContainer.perform { ... }`./// Pure function: given model/tokenizer/pool + texts, returns `[[Float]]`.enum Batching { static func embed( model: EmbeddingModel, tokenizer: any Tokenizer, pooling: Pooling, texts: [String], kind: InputKind, maxLen: Int ) -> [[Float]] { precondition(!texts.isEmpty, "embed called with empty batch")
// 1. Prompt wrap + tokenize let wrapped = texts.map { Prompt.wrap($0, kind) } var rows: [[Int32]] = wrapped.map { text in let ids = tokenizer.encode(text: text, addSpecialTokens: true) // Truncate to maxLen hard cap (rare; buckets usually govern this). let capped = ids.count > maxLen ? Array(ids.prefix(maxLen)) : ids return capped.map { Int32($0) } }
// 2. Pick length bucket from the longest row. let maxRow = rows.map { $0.count }.max() ?? 1 let length = Buckets.length(for: max(maxRow, 1), maxLen: maxLen)
// 3. Pad rows to length; build attention mask. let padId: Int32 = { if let eos = tokenizer.eosTokenId { return Int32(eos) } return 0 }()
let batchSize = rows.count var idsFlat = [Int32](repeating: padId, count: batchSize * length) var maskFlat = [Int32](repeating: 0, count: batchSize * length) for (i, row) in rows.enumerated() { let n = min(row.count, length) let base = i * length for j in 0..<n { idsFlat[base + j] = row[j] maskFlat[base + j] = 1 } } rows.removeAll(keepingCapacity: false)
let inputIds = MLXArray(idsFlat, [batchSize, length]) let attentionMask = MLXArray(maskFlat, [batchSize, length])
// 4. Forward. let out = model( inputIds, positionIds: nil, tokenTypeIds: nil, attentionMask: attentionMask )
// 5. Last-token pool + L2 normalize. Pooling mask consumed as-is; it // casts to int32 internally (see Pooling.swift: `_mask.asType(.int32)`). let pooled = pooling(out, mask: attentionMask, normalize: true)
// 6. Evaluate before leaving the perform closure (MLXArray is not Sendable). eval(pooled)
// 7. Copy to host memory as `[[Float]]`. let dim = pooled.dim(1) let host: [Float] = pooled.asArray(Float.self) var result = [[Float]]() result.reserveCapacity(batchSize) for i in 0..<batchSize { let base = i * dim result.append(Array(host[base..<(base + dim)])) } return result }}