v2: add EmbeddingServer binary, EmbeddingGemma model (WIP forward pass)

2026-07-06 10:13:39 +08:00
parent e7a94b3203
commit dbec6b20ea
3 changed files with 104 additions and 9 deletions
@@ -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()
+    }
+}