warren/markbaseengine
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

v2: add EmbeddingServerApp, fix Sendable issues (WIP)

Browse Source
This commit is contained in:
MarkBase Admin
2026-07-06 09:47:06 +08:00
parent 88aeff7935
commit e7a94b3203
4 changed files with 162 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sources/MarkBaseServer/EmbeddingServerApp.swift
+113
View File
@@ -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
+38 -10
View File
@@ -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,
Sources/MarkBaseServer/main.swift
+10 -2
View File
@@ -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)