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))")
|
||||
|
||||
// Use actor to serialize access to embedModel
|
||||
let embedder = Embedder(model: embedModel)
|
||||
|
||||
let router = Router()
|
||||
|
||||
router.get("/") { _, _ in
|
||||
@@ -49,10 +52,8 @@ struct EmbeddingServerMain {
|
||||
|
||||
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 result = try await embedder.embed(text: text)
|
||||
embeddings.append(["object": "embedding", "index": i, "embedding": result.embedding, "usage_ms": result.ms])
|
||||
}
|
||||
|
||||
let id = UUID().uuidString
|
||||
@@ -75,3 +76,23 @@ struct EmbeddingServerMain {
|
||||
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