diff --git a/Package.swift b/Package.swift index 1c34907..e677e77 100644 --- a/Package.swift +++ b/Package.swift @@ -7,6 +7,7 @@ let package = Package( products: [ .library(name: "MarkBase", targets: ["MarkBase"]), .executable(name: "MarkBaseServer", targets: ["MarkBaseServer"]), + .executable(name: "EmbeddingServer", targets: ["EmbeddingServer"]), .executable(name: "CLITest", targets: ["CLITest"]), ], dependencies: [ @@ -34,6 +35,17 @@ let package = Package( .linkedFramework("Foundation"), ] ), + .executableTarget( + name: "EmbeddingServer", + dependencies: [ + "MarkBase", + .product(name: "Hummingbird", package: "hummingbird"), + ], + linkerSettings: [ + .linkedFramework("Metal"), + .linkedFramework("Foundation"), + ] + ), .executableTarget( name: "CLITest", dependencies: ["MarkBase"], diff --git a/Sources/EmbeddingServer/EmbeddingServerMain.swift b/Sources/EmbeddingServer/EmbeddingServerMain.swift new file mode 100644 index 0000000..1339457 --- /dev/null +++ b/Sources/EmbeddingServer/EmbeddingServerMain.swift @@ -0,0 +1,77 @@ +import Foundation +import MarkBase +import Hummingbird + +@main +struct EmbeddingServerMain { + static func main() async throws { + let args = CommandLine.arguments + let modelName = args.count > 1 ? args[1] : "embeddinggemma-300m" + let port = args.count > 2 ? Int(args[2]) ?? 8084 : 8084 + + let modelPath = NSString(string: "~/MarkBaseEngine/models/\(modelName)").expandingTildeInPath + + print("MarkBaseEngine Embedding Server") + print(" Model: \(modelName)") + print(" Port: \(port)") + print(" Path: \(modelPath)") + + let engine = try MarkBaseEngine(autoCompile: true) + let embedModel = try EmbeddingGemmaModel(modelDir: modelPath, engine: engine) + let layers = embedModel.config.numHiddenLayers + let hiddenSize = embedModel.config.hiddenSize + + print("EmbeddingGemma loaded (\(layers) layers, hidden=\(hiddenSize))") + + let router = Router() + + router.get("/") { _, _ in + return "{\"server\":\"MarkBaseEngine Embedding\",\"model\":\"embeddinggemma-300m\",\"layers\":\(layers),\"hidden_size\":\(hiddenSize)}" + } + + router.get("/health") { _, _ in + return "{\"status\":\"healthy\",\"model\":\"embeddinggemma-300m\",\"layers\":\(layers)}" + } + + router.post("/v1/embeddings") { request, _ in + let buffer = try await request.body.collect(upTo: .max) + let data = Data(buffer: buffer) + + guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any], + let input = json["input"] else { + return "{\"error\":\"missing 'input' field\"}" + } + + let inputs: [String] + if let str = input as? String { inputs = [str] } + else if let arr = input as? [String] { inputs = arr } + else { return "{\"error\":\"'input' must be string or array\"}" } + + var embeddings: [[String: Any]] = [] + for (i, text) in inputs.enumerated() { + let t0 = Date() + let embedding = try embedModel.embed(text: text) + let ms = Int(Date().timeIntervalSince(t0) * 1000) + embeddings.append(["object": "embedding", "index": i, "embedding": embedding, "usage_ms": ms]) + } + + let id = UUID().uuidString + let ts = Int(Date().timeIntervalSince1970) + let response: [String: Any] = [ + "id": id, "object": "list", "created": ts, "model": "embeddinggemma-300m", + "data": embeddings + ] + + let jsonData = try JSONSerialization.data(withJSONObject: response) + return String(data: jsonData, encoding: .utf8) ?? "{}" + } + + let app = Application( + router: router, + configuration: .init(address: .hostname("0.0.0.0", port: port)) + ) + + print("Server starting on port \(port)...") + try await app.run() + } +} diff --git a/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift b/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift index 895aacc..9e55d4f 100644 --- a/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift +++ b/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift @@ -284,15 +284,21 @@ public final class EmbeddingGemmaModel: @unchecked Sendable { let enc = cmdBuf.makeComputeCommandEncoder()! let pso = try engine.pipeline(named: "matmul_f32") enc.setComputePipelineState(pso) - enc.setBuffer(input, offset: 0, index: 0) - enc.setBuffer(weight, offset: 0, index: 1) - enc.setBuffer(output, offset: 0, index: 2) - var mm = UInt32(m), kk = UInt32(k), nn = UInt32(n) - enc.setBytes(&mm, length: 4, index: 3) - enc.setBytes(&kk, length: 4, index: 4) - enc.setBytes(&nn, length: 4, index: 5) - enc.dispatchThreads(MTLSize(width: m * n, height: 1, depth: 1), - threadsPerThreadgroup: MTLSize(width: min(256, m * n), height: 1, depth: 1)) + // Note: matmul_f32 only handles M=1, so we loop over positions + // For production: implement a proper batch matmul kernel + for i in 0..