v2: add EmbeddingServerApp, fix Sendable issues (WIP)

2026-07-06 09:47:06 +08:00
parent 88aeff7935
commit e7a94b3203
4 changed files with 162 additions and 13 deletions
@@ -0,0 +1,113 @@
+import Foundation
+import MarkBase
+import Hummingbird
+
+struct EmbeddingServerApp {
+    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("═══════════════════════════════════════════════════════════════════")
+        print("  MarkBaseEngine Embedding Server")
+        print("═══════════════════════════════════════════════════════════════════")
+        print("  Model: \(modelName)")
+        print("  Port: \(port)")
+        print("  Path: \(modelPath)")
+        print("")
+
+        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": {
+                "name": "MarkBaseEngine Embedding",
+                "model": "embeddinggemma-300m",
+                "layers": \(layers),
+                "hidden_size": \(hiddenSize),
+                "output_dim": \(hiddenSize),
+                "framework": "Hummingbird 2.x + Metal GPU"
+              },
+              "endpoints": [
+                {"method": "POST", "path": "/v1/embeddings", "summary": "Generate text embeddings"}
+              ]
+            }
+            """
+        }
+
+        router.get("/health") { _, _ in
+            return "{\"status\":\"healthy\",\"model\":\"embeddinggemma-300m\",\"layers\":\(layers),\"hidden_size\":\(hiddenSize)}"
+        }
+
+        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\":\"invalid request\",\"type\":\"invalid_request_error\",\"code\":400,\"message\":\"missing 'input' field\"}"
+            }
+
+            let modelId = (json["model"] as? String) ?? "embeddinggemma-300m"
+            let encodingFormat = (json["encoding_format"] as? String) ?? "float"
+
+            let inputs: [String]
+            if let str = input as? String { inputs = [str] }
+            else if let arr = input as? [String] { inputs = arr }
+            else {
+                return "{\"error\":\"invalid request\",\"type\":\"invalid_request_error\",\"code\":400,\"message\":\"'input' must be string or array of strings\"}"
+            }
+
+            var embeddings: [[String: Any]] = []
+            for (i, text) in inputs.enumerated() {
+                let t0 = Date()
+                let embedding = try embedModel.embed(text: text)
+                let duration = Date().timeIntervalSince(t0)
+
+                let embeddingData: [String: Any]
+                if encodingFormat == "base64" {
+                    let base64 = embedding.withUnsafeBytes { Data($0).base64EncodedString() }
+                    embeddingData = ["object": "embedding", "index": i, "embedding": base64, "usage_ms": Int(duration * 1000)]
+                } else {
+                    embeddingData = ["object": "embedding", "index": i, "embedding": embedding, "usage_ms": Int(duration * 1000)]
+                }
+                embeddings.append(embeddingData)
+            }
+
+            let id = UUID().uuidString
+            let ts = Int(Date().timeIntervalSince1970)
+            let response: [String: Any] = [
+                "id": id, "object": "list", "created": ts, "model": modelId,
+                "data": embeddings,
+                "usage": ["prompt_tokens": inputs.reduce(0) { $0 + $1.components(separatedBy: .whitespaces).count }, "total_tokens": inputs.reduce(0) { $0 + $1.components(separatedBy: .whitespaces).count }]
+            ]
+
+            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)...")
+        print("Endpoints:")
+        print("  GET  /             - Info")
+        print("  GET  /health       - Health check")
+        print("  POST /v1/embeddings - Text embeddings")
+        print("")
+
+        try await app.run()
+    }
+}