MarkBase Admin
78 lines
2.9 KiB
Swift
78 lines
2.9 KiB
Swift
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()
|
|
}
|
|
}
|