v2: add EmbeddingServerApp, fix Sendable issues (WIP)

Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift

@@ -42,7 +42,7 @@ public struct EmbeddingGemmaConfig: Codable {
 }
 
 /// EmbeddingGemma - Google's 300M parameter embedding model
-public final class EmbeddingGemmaModel {
+public final class EmbeddingGemmaModel: @unchecked Sendable {
     public let config: EmbeddingGemmaConfig
     public let engine: MarkBaseEngine
     public let tokenizer: Tokenizer

Sources/MarkBaseServer/EmbeddingServerApp.swift

@@ -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()
+    }
+}

Sources/MarkBaseServer/SimpleServerApp.swift

@@ -20,16 +20,36 @@ struct SimpleServerApp {
         print("")
 
         let engine = try MarkBaseEngine(autoCompile: true)
-        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
-        let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
-        let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
-        let embeddingModel = try TextEmbeddingModel(modelDir: modelPath, engine: engine, config: TextEmbeddingConfig())
         
-        print("✓ E4B loaded (\(model.numHiddenLayers) layers)")
+        // Detect model type
+        let isEmbeddingGemma = modelName.contains("embeddinggemma") || modelName.contains("embedding")
+        let isE4B = !isEmbeddingGemma
+        
+        var e4bModel: E4BModel? = nil
+        var embedModel: EmbeddingGemmaModel? = nil
+        let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
+        var generator: StreamingGenerator? = nil
+        
+        if isE4B {
+            e4bModel = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
+            generator = StreamingGenerator(model: e4bModel!, tokenizer: tokenizer, engine: engine)
+            print("✓ E4B loaded (\(e4bModel!.numHiddenLayers) layers)")
+        } else {
+            embedModel = try EmbeddingGemmaModel(modelDir: modelPath, engine: engine)
+            print("✓ EmbeddingGemma loaded (\(embedModel!.config.numHiddenLayers) layers, hidden=\(embedModel!.config.hiddenSize))")
+        }
+        
+        // Embedding model (use EmbeddingGemma if available, fallback to E4B)
+        let textEmbedModel: TextEmbeddingModel?
+        if isE4B {
+            textEmbedModel = try TextEmbeddingModel(modelDir: modelPath, engine: engine, config: TextEmbeddingConfig())
+        } else {
+            textEmbedModel = nil
+        }
 
         let router = Router()
 
-        let layers = model.numHiddenLayers
+        let layers = isE4B ? e4bModel!.numHiddenLayers : embedModel!.config.numHiddenLayers
 
         @Sendable func helpJSON() -> String {
             return """
@@ -184,7 +204,8 @@ struct SimpleServerApp {
         }
 
         @Sendable func healthResponse() -> String {
-            return "{\"status\":\"healthy\",\"model\":\"e4b\",\"layers\":\(layers)}"
+            let m = isEmbeddingGemma ? "embeddinggemma-300m" : "e4b"
+            return "{\"status\":\"healthy\",\"model\":\"\(m)\",\"layers\":\(layers)}"
         }
 
         router.get("/") { _, _ in
@@ -229,7 +250,7 @@ struct SimpleServerApp {
                 topP: topP
             )
 
-            let generatedTokens = try generator.generateTokens(promptTokens: promptTokens, config: config)
+            let generatedTokens = try generator!.generateTokens(promptTokens: promptTokens, config: config)
 
             let id = UUID().uuidString
             let ts = Int(Date().timeIntervalSince1970)
@@ -313,7 +334,14 @@ struct SimpleServerApp {
             var embeddings: [[String: Any]] = []
             for (i, text) in inputs.enumerated() {
                 let t0 = Date()
-                let embedding = try embeddingModel.embed(text: text)
+                let embedding: [Float]
+                if let eg = embedModelCopy {
+                    embedding = try eg.embed(text: text)
+                } else if let tem = textEmbedModelCopy {
+                    embedding = try tem.embed(text: text)
+                } else {
+                    return "{\"error\":\"no embedding model available\",\"type\":\"server_error\",\"code\":500}"
+                }
                 let duration = Date().timeIntervalSince(t0)
 
                 let embeddingData: [String: Any]
@@ -338,7 +366,7 @@ struct SimpleServerApp {
 
             let id = UUID().uuidString
             let ts = Int(Date().timeIntervalSince1970)
-            let totalTokens = inputs.reduce(0) { $0 + tokenizer.encode(text: $1).count }
+            let totalTokens = inputs.reduce(0) { $0 + $1.components(separatedBy: .whitespaces).count }
 
             let response: [String: Any] = [
                 "id": id,

Sources/MarkBaseServer/main.swift

@@ -1,9 +1,17 @@
 import Foundation
 
-// Entry point — avoids @main conflict with top-level code
+// Entry point — routes to appropriate server based on model name
 Task {
     do {
+        let args = CommandLine.arguments
+        let modelName = args.count > 1 ? args[1] : "E4B-MarkBase"
+        let isEmbedding = modelName.contains("embedding") || modelName.contains("embeddinggemma")
+        
+        if isEmbedding {
+            try await EmbeddingServerApp.main()
+        } else {
             try await SimpleServerApp.main()
+        }
     } catch {
         print("Server error: \(error)")
         exit(1)