v2: add actor-based serialization for EmbeddingServer, fix stability for ASCII text
This commit is contained in:
@@ -23,6 +23,9 @@ struct EmbeddingServerMain {
|
|||||||
|
|
||||||
print("EmbeddingGemma loaded (\(layers) layers, hidden=\(hiddenSize))")
|
print("EmbeddingGemma loaded (\(layers) layers, hidden=\(hiddenSize))")
|
||||||
|
|
||||||
|
// Use actor to serialize access to embedModel
|
||||||
|
let embedder = Embedder(model: embedModel)
|
||||||
|
|
||||||
let router = Router()
|
let router = Router()
|
||||||
|
|
||||||
router.get("/") { _, _ in
|
router.get("/") { _, _ in
|
||||||
@@ -49,10 +52,8 @@ struct EmbeddingServerMain {
|
|||||||
|
|
||||||
var embeddings: [[String: Any]] = []
|
var embeddings: [[String: Any]] = []
|
||||||
for (i, text) in inputs.enumerated() {
|
for (i, text) in inputs.enumerated() {
|
||||||
let t0 = Date()
|
let result = try await embedder.embed(text: text)
|
||||||
let embedding = try embedModel.embed(text: text)
|
embeddings.append(["object": "embedding", "index": i, "embedding": result.embedding, "usage_ms": result.ms])
|
||||||
let ms = Int(Date().timeIntervalSince(t0) * 1000)
|
|
||||||
embeddings.append(["object": "embedding", "index": i, "embedding": embedding, "usage_ms": ms])
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let id = UUID().uuidString
|
let id = UUID().uuidString
|
||||||
@@ -75,3 +76,23 @@ struct EmbeddingServerMain {
|
|||||||
try await app.run()
|
try await app.run()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct EmbeddingResult {
|
||||||
|
let embedding: [Float]
|
||||||
|
let ms: Int
|
||||||
|
}
|
||||||
|
|
||||||
|
actor Embedder {
|
||||||
|
private let model: EmbeddingGemmaModel
|
||||||
|
|
||||||
|
init(model: EmbeddingGemmaModel) {
|
||||||
|
self.model = model
|
||||||
|
}
|
||||||
|
|
||||||
|
func embed(text: String) throws -> EmbeddingResult {
|
||||||
|
let t0 = Date()
|
||||||
|
let embedding = try model.embed(text: text)
|
||||||
|
let ms = Int(Date().timeIntervalSince(t0) * 1000)
|
||||||
|
return EmbeddingResult(embedding: embedding, ms: ms)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user