v2: add embedding tests, multilingual embedding support
This commit is contained in:
@@ -23,6 +23,7 @@ struct SimpleServerApp {
|
||||
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)")
|
||||
|
||||
@@ -168,14 +169,15 @@ struct SimpleServerApp {
|
||||
"deployment": "docs/DEPLOYMENT.md",
|
||||
"performance": "docs/PERFORMANCE.md"
|
||||
},
|
||||
"notes": [
|
||||
"notes": [
|
||||
"All responses are in JSON format",
|
||||
"Text generation only (multimodal not yet supported via API)",
|
||||
"E4B model with 42 layers, ~4B parameters",
|
||||
"For multimodal (vision/audio) support, use the MarkBase Swift library directly",
|
||||
"Streaming support is planned but not yet implemented",
|
||||
"Function calling uses native Gemma 4 special tokens",
|
||||
"Messages can include tool_calls (assistant) and tool responses (tool role) for multi-turn function calling"
|
||||
"Messages can include tool_calls (assistant) and tool responses (tool role) for multi-turn function calling",
|
||||
"Text embeddings available via /v1/embeddings endpoint (OpenAI-compatible)"
|
||||
]
|
||||
}
|
||||
"""
|
||||
@@ -287,6 +289,73 @@ struct SimpleServerApp {
|
||||
}
|
||||
}
|
||||
|
||||
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) ?? "e4b"
|
||||
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 embeddingModel.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 totalTokens = inputs.reduce(0) { $0 + tokenizer.encode(text: $1).count }
|
||||
|
||||
let response: [String: Any] = [
|
||||
"id": id,
|
||||
"object": "list",
|
||||
"created": ts,
|
||||
"model": modelId,
|
||||
"data": embeddings,
|
||||
"usage": [
|
||||
"prompt_tokens": totalTokens,
|
||||
"total_tokens": totalTokens
|
||||
]
|
||||
]
|
||||
|
||||
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))
|
||||
@@ -299,6 +368,7 @@ struct SimpleServerApp {
|
||||
print(" GET /health - Health check")
|
||||
print(" GET /v1/models - Model list")
|
||||
print(" POST /v1/chat/completions - Chat completion")
|
||||
print(" POST /v1/embeddings - Text embeddings")
|
||||
print("")
|
||||
print("Model: \(modelName)")
|
||||
if modelName.contains("E4B") {
|
||||
|
||||
Reference in New Issue
Block a user