From e7a94b32035fe89b51f6779bf4d40073f8174a70 Mon Sep 17 00:00:00 2001 From: MarkBase Admin Date: Mon, 6 Jul 2026 09:47:06 +0800 Subject: [PATCH] v2: add EmbeddingServerApp, fix Sendable issues (WIP) --- .../Embedding/EmbeddingGemmaModel.swift | 2 +- .../MarkBaseServer/EmbeddingServerApp.swift | 113 ++++++++++++++++++ Sources/MarkBaseServer/SimpleServerApp.swift | 48 ++++++-- Sources/MarkBaseServer/main.swift | 12 +- 4 files changed, 162 insertions(+), 13 deletions(-) create mode 100644 Sources/MarkBaseServer/EmbeddingServerApp.swift diff --git a/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift b/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift index 3859f42..895aacc 100644 --- a/Sources/MarkBase/Embedding/EmbeddingGemmaModel.swift +++ b/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 diff --git a/Sources/MarkBaseServer/EmbeddingServerApp.swift b/Sources/MarkBaseServer/EmbeddingServerApp.swift new file mode 100644 index 0000000..07324b4 --- /dev/null +++ b/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() + } +} diff --git a/Sources/MarkBaseServer/SimpleServerApp.swift b/Sources/MarkBaseServer/SimpleServerApp.swift index f0475ba..b8a0f71 100644 --- a/Sources/MarkBaseServer/SimpleServerApp.swift +++ b/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) + + // 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) - 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)") + 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, diff --git a/Sources/MarkBaseServer/main.swift b/Sources/MarkBaseServer/main.swift index f3a90fa..9c24d98 100644 --- a/Sources/MarkBaseServer/main.swift +++ b/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 { - try await SimpleServerApp.main() + 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)