spb/prisme Public MIT

1//2//  SentenceEmbedder.swift3//  Prisme4//5//  Author: Simon-Pierre Boucher <contact@spboucher.ai>6//78import Foundation9import NaturalLanguage1011/// On-device sentence embeddings (NaturalLanguage framework — no Foundation12/// Models involved, works on every device, offline). Vectors from different13/// language models live in different spaces: the store tags each vector14/// with its language and only compares like with like.15@MainActor16final class SentenceEmbedder {17    private var cache: [NLLanguage: NLEmbedding] = [:]1819    func languageCode(for text: String) -> String? {20        let recognizer = NLLanguageRecognizer()21        recognizer.processString(String(text.prefix(400)))22        return recognizer.dominantLanguage?.rawValue23    }2425    func embed(_ text: String, languageCode: String?) -> [Float]? {26        guard let languageCode else { return nil }27        let language = NLLanguage(rawValue: languageCode)28        let embedding: NLEmbedding?29        if let cached = cache[language] {30            embedding = cached31        } else {32            embedding = NLEmbedding.sentenceEmbedding(for: language)33            if let embedding { cache[language] = embedding }34        }35        guard let vector = embedding?.vector(for: String(text.prefix(600))) else {36            return nil37        }38        return vector.map(Float.init)39    }4041    static func cosine(_ a: [Float], _ b: [Float]) -> Float {42        guard a.count == b.count, !a.isEmpty else { return 0 }43        var dot: Float = 0, normA: Float = 0, normB: Float = 044        for i in a.indices {45            dot += a[i] * b[i]46            normA += a[i] * a[i]47            normB += b[i] * b[i]48        }49        let denominator = (normA.squareRoot() * normB.squareRoot())50        return denominator > 0 ? dot / denominator : 051    }52}53