diff --git a/Sources/EmbeddingServer/EmbeddingServerMain.swift b/Sources/EmbeddingServer/EmbeddingServerMain.swift index 1339457..b3002b8 100644 --- a/Sources/EmbeddingServer/EmbeddingServerMain.swift +++ b/Sources/EmbeddingServer/EmbeddingServerMain.swift @@ -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) + } +}