v2: add EmbeddingServer binary, EmbeddingGemma model (WIP forward pass)
This commit is contained in:
@@ -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"],
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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..<m {
|
||||
let inputOffset = i * k * 4
|
||||
let outputOffset = i * n * 4
|
||||
enc.setBuffer(input, offset: inputOffset, index: 0)
|
||||
enc.setBuffer(weight, offset: 0, index: 1)
|
||||
enc.setBuffer(output, offset: outputOffset, index: 2)
|
||||
var mm: UInt32 = 1, 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: n, height: 1, depth: 1),
|
||||
threadsPerThreadgroup: MTLSize(width: min(256, n), height: 1, depth: 1))
|
||||
}
|
||||
enc.endEncoding()
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user