@glassly/bluetooth-sdk 0.1.0-dev.26

package/ios/Source/speakerid/SpeakerIdEngine.swift

@@ -0,0 +1,573 @@
+//
+//  SpeakerIdEngine.swift
+//
+//  On-device speaker identification over the canonical 16 kHz mono PCM stream.
+//
+//  Taps DeviceManager.handlePcm (the single audio choke point), windows voiced
+//  audio with a local energy gate, computes a speaker embedding per
+//  utterance-ish window (sherpa-onnx CAM++ model), matches against enrolled
+//  profile centroids by cosine similarity, and emits `speaker_detected`
+//  bridge events. Raw audio never leaves the phone for speaker ID; enrollment
+//  audio is discarded as soon as its embeddings are computed.
+//
+//  Profiles are pushed in from JS (`setSpeakerProfiles`) — MMKV in the engine
+//  layer is the source of truth; native holds them in memory only.
+//
+
+import Foundation
+
+#if !SWIFT_PACKAGE || GLASSLY_FEATURE_LOCAL_STT
+
+final class SpeakerIdEngine {
+    static let shared = SpeakerIdEngine()
+
+    struct Profile {
+        let id: String
+        let name: String
+        /// L2-normalized centroid.
+        let embedding: [Float]
+    }
+
+    struct Match {
+        let profileId: String
+        let name: String
+        let similarity: Float
+        /// Epoch ms when the window producing this match ended.
+        let endMs: Int64
+    }
+
+    // MARK: - Tuning (see notes/speaker-id-plan.md §2)
+
+    /// Cosine threshold for accepting a match (normalized embeddings).
+    /// Benchmarked cross-speaker ceiling for this model is ~0.46 (the MAX
+    /// over clean-audio pairs; typical impostors score far lower), while
+    /// on-device same-speaker scores run below the 0.67+ clean bench.
+    /// Owner-only voice recognition favors recall: a borderline accept just
+    /// labels a line "You", and the JS segment resolver's overlap voting +
+    /// conflict damping still filter isolated stray windows — so sit at the
+    /// edge of the impostor range rather than above it. (With one enrolled
+    /// profile the margin rule below is vacuous.)
+    private static let matchThreshold: Float = 0.45
+    /// Reject when the runner-up profile is within this margin of the best.
+    private static let matchMargin: Float = 0.08
+    /// Flush a window once this much voiced audio accumulated.
+    private static let windowVoicedTargetMs = 3000.0
+    /// End-of-speech hangover: flush after this much trailing silence.
+    private static let endOfSpeechSilenceMs = 500.0
+    /// Windows with less voiced audio than this are skipped (unreliable).
+    private static let minVoicedMs = 1000.0
+    /// Hard cap on buffered window samples (10 s) to bound memory.
+    private static let maxWindowSamples = 16000 * 10
+    /// RMS energy gate on [-1, 1] floats.
+    private static let voicedRmsThreshold: Float = 0.01
+    /// Enrollment chunks: pairwise cosine consistency floor.
+    private static let enrollConsistencyThreshold: Float = 0.6
+    private static let enrollChunkCount = 3
+
+    private static let sampleRate = 16000
+    private static let modelPathKey = "SpeakerIdModelPath"
+
+    // MARK: - State
+
+    private let queue = DispatchQueue(label: "com.glassly.speakerid.processor", qos: .utility)
+
+    /// Cheap gate read on the audio thread before hopping queues.
+    private let activeFlag = NSLock()
+    private var _identifyEnabled = false
+    private var _enrolling = false
+    private var isActive: Bool {
+        activeFlag.lock()
+        defer { activeFlag.unlock() }
+        return _identifyEnabled || _enrolling
+    }
+
+    // Everything below is touched on `queue` only.
+    private var extractor: SherpaOnnxSpeakerEmbeddingExtractorWrapper?
+    private var extractorLoadAttempted = false
+    private var profiles: [Profile] = []
+
+    private var windowSamples: [Float] = []
+    private var windowStartMs: Int64 = 0
+    private var windowVoicedMs = 0.0
+    private var trailingSilenceMs = 0.0
+
+    private struct EnrollmentSession {
+        var targetVoicedMs: Double
+        var source: String
+        var samples: [Float] = []
+        var voicedMs = 0.0
+        var lastProgressAt = Date(timeIntervalSince1970: 0)
+        var lastLevel: Float = 0
+    }
+
+    private var enrollment: EnrollmentSession?
+
+    private let matchLock = NSLock()
+    private var _lastMatch: Match?
+
+    // MARK: - Public API (any thread unless noted)
+
+    func setIdentifyEnabled(_ enabled: Bool) {
+        activeFlag.lock()
+        _identifyEnabled = enabled
+        activeFlag.unlock()
+        Bridge.log("SpeakerID: identify \(enabled ? "enabled" : "disabled")")
+        queue.async { [weak self] in
+            guard let self else { return }
+            if enabled {
+                self.ensureExtractor()
+            } else {
+                self.resetWindow()
+            }
+        }
+    }
+
+    func setProfiles(_ profiles: [Profile]) {
+        queue.async { [weak self] in
+            self?.profiles = profiles
+            Bridge.log(
+                "SpeakerID: profiles set — \(profiles.count): "
+                    + profiles.map { "\($0.name)(\($0.embedding.count)d)" }.joined(separator: ", ")
+            )
+        }
+    }
+
+    static func setModelPath(_ path: String) {
+        UserDefaults.standard.set(path, forKey: modelPathKey)
+        Bridge.log("SpeakerID: model path set → \(path)")
+        SpeakerIdEngine.shared.queue.async {
+            // Force a reload on next use — the model may have just arrived.
+            SpeakerIdEngine.shared.extractor = nil
+            SpeakerIdEngine.shared.extractorLoadAttempted = false
+        }
+    }
+
+    static func isModelAvailable() -> Bool {
+        guard let path = resolvedModelPath() else { return false }
+        return FileManager.default.fileExists(atPath: path)
+    }
+
+    /// `onEnded` fires (on an arbitrary queue) whenever the enrollment session
+    /// finishes or fails, so the caller can release the mic requirement.
+    func startEnrollment(targetSpeechSeconds: Double, source: String, onEnded: @escaping () -> Void) {
+        activeFlag.lock()
+        _enrolling = true
+        activeFlag.unlock()
+        queue.async { [weak self] in
+            guard let self else { return }
+            self.ensureExtractor()
+            guard self.extractor != nil else {
+                Bridge.log("SpeakerID: enrollment aborted — model unavailable")
+                self.finishEnrollment(ok: false, reason: "model_unavailable", embeddings: [], source: source)
+                onEnded()
+                return
+            }
+            Bridge.log("SpeakerID: enrollment started — target \(targetSpeechSeconds)s voiced, mic=\(source)")
+            self.enrollmentEnded = onEnded
+            self.enrollment = EnrollmentSession(
+                targetVoicedMs: max(5.0, targetSpeechSeconds) * 1000.0,
+                source: source
+            )
+        }
+    }
+
+    private var enrollmentEnded: (() -> Void)?
+
+    func cancelEnrollment() {
+        activeFlag.lock()
+        _enrolling = false
+        activeFlag.unlock()
+        queue.async { [weak self] in
+            guard let self else { return }
+            self.enrollment = nil
+            let onEnded = self.enrollmentEnded
+            self.enrollmentEnded = nil
+            onEnded?()
+        }
+    }
+
+    /// Latest accepted identification, if it ended within `withinMs` of now.
+    /// Used by the local STT path to stamp transcripts (native-side join).
+    func recentMatch(withinMs: Int64) -> Match? {
+        matchLock.lock()
+        defer { matchLock.unlock() }
+        guard let match = _lastMatch else { return nil }
+        let now = Int64(Date().timeIntervalSince1970 * 1000)
+        return now - match.endMs <= withinMs ? match : nil
+    }
+
+    /// Called from DeviceManager.handlePcm with 16 kHz mono s16le PCM.
+    func acceptAudio(pcm16le: Data) {
+        guard isActive, pcm16le.count >= 2 else { return }
+        queue.async { [weak self] in
+            self?.process(pcm16le)
+        }
+    }
+
+    // MARK: - Processing (queue only)
+
+    private func process(_ pcmData: Data) {
+        let samples = Self.toFloatArray(pcmData)
+        guard !samples.isEmpty else { return }
+        let chunkMs = Double(samples.count) / Double(Self.sampleRate) * 1000.0
+        let rms = Self.rms(samples)
+        let voiced = rms > Self.voicedRmsThreshold
+
+        if enrollment != nil {
+            processEnrollment(samples: samples, chunkMs: chunkMs, rms: rms, voiced: voiced)
+            return
+        }
+
+        activeFlag.lock()
+        let identify = _identifyEnabled
+        activeFlag.unlock()
+        guard identify else { return }
+        processIdentification(samples: samples, chunkMs: chunkMs, voiced: voiced)
+    }
+
+    private func processIdentification(samples: [Float], chunkMs: Double, voiced: Bool) {
+        let nowMs = Int64(Date().timeIntervalSince1970 * 1000)
+
+        if windowSamples.isEmpty {
+            guard voiced else { return }
+            windowStartMs = nowMs - Int64(chunkMs)
+            windowVoicedMs = 0
+            trailingSilenceMs = 0
+        }
+
+        // Accumulate VOICED samples only — enrollment computes centroids from
+        // voiced-only audio, so identification must embed like-for-like.
+        // Windows padded with inter-phrase silence score far below the same
+        // speaker's centroid (observed 0.1–0.48 vs the 0.67+ bench).
+        if voiced {
+            if windowSamples.count < Self.maxWindowSamples {
+                windowSamples.append(contentsOf: samples)
+            }
+            windowVoicedMs += chunkMs
+            trailingSilenceMs = 0
+        } else {
+            trailingSilenceMs += chunkMs
+        }
+
+        let endOfSpeech = trailingSilenceMs >= Self.endOfSpeechSilenceMs
+        let windowFull = windowVoicedMs >= Self.windowVoicedTargetMs
+        guard endOfSpeech || windowFull else { return }
+
+        let startMs = windowStartMs
+        let voicedMs = windowVoicedMs
+        let window = windowSamples
+        resetWindow()
+
+        guard voicedMs >= Self.minVoicedMs else {
+            Bridge.log("SpeakerID: window skipped — only \(Int(voicedMs))ms voiced (<\(Int(Self.minVoicedMs))ms)")
+            return
+        }
+        identify(window: window, startMs: startMs, endMs: nowMs, voicedMs: voicedMs)
+    }
+
+    private func identify(window: [Float], startMs: Int64, endMs: Int64, voicedMs: Double) {
+        ensureExtractor()
+        guard let extractor else {
+            Bridge.log("SpeakerID: window dropped — extractor not loaded (model missing?)")
+            return
+        }
+        guard !profiles.isEmpty else {
+            Bridge.log("SpeakerID: window dropped — no enrolled profiles pushed")
+            return
+        }
+
+        let stream = extractor.createStream()
+        stream.acceptWaveform(samples: window, sampleRate: Self.sampleRate)
+        stream.inputFinished()
+        guard extractor.isReady(stream: stream) else {
+            Bridge.log("SpeakerID: window dropped — extractor not ready (\(window.count) samples)")
+            return
+        }
+        let embedding = Self.l2Normalize(extractor.compute(stream: stream))
+        guard !embedding.isEmpty else {
+            Bridge.log("SpeakerID: window dropped — empty embedding")
+            return
+        }
+
+        var best: (profile: Profile, similarity: Float)?
+        var second: Float = -1
+        for profile in profiles where profile.embedding.count == embedding.count {
+            let cos = Self.dot(profile.embedding, embedding)
+            if best == nil || cos > best!.similarity {
+                second = best?.similarity ?? -1
+                best = (profile, cos)
+            } else if cos > second {
+                second = cos
+            }
+        }
+
+        var matched: Profile?
+        var similarity: Float = 0
+        if let best, best.similarity >= Self.matchThreshold {
+            // Margin rule: reject ambiguous windows where the runner-up is
+            // close (overlapping talkers produce garbage embeddings).
+            if second < 0 || best.similarity - second >= Self.matchMargin {
+                matched = best.profile
+                similarity = best.similarity
+            }
+        }
+
+        if let matched {
+            matchLock.lock()
+            _lastMatch = Match(profileId: matched.id, name: matched.name, similarity: similarity, endMs: endMs)
+            matchLock.unlock()
+        }
+
+        let durationMs = Int(endMs - startMs)
+        let speech = "speech=\(Int(voicedMs))ms window=\(durationMs)ms"
+        if let matched {
+            Bridge.log(
+                "SpeakerID: matched \(matched.name) cos=\(String(format: "%.3f", similarity)) "
+                    + "second=\(String(format: "%.3f", second)) \(speech)"
+            )
+        } else if let best {
+            let reason = best.similarity < Self.matchThreshold
+                ? "below threshold \(Self.matchThreshold)"
+                : "margin to runner-up < \(Self.matchMargin)"
+            Bridge.log(
+                "SpeakerID: no match — best \(best.profile.name) cos=\(String(format: "%.3f", best.similarity)) "
+                    + "second=\(String(format: "%.3f", second)) (\(reason)) \(speech)"
+            )
+        }
+
+        Bridge.sendSpeakerDetected(
+            speakerId: matched?.id,
+            name: matched?.name,
+            similarity: Double(matched != nil ? similarity : (best?.similarity ?? 0)),
+            startMs: startMs,
+            endMs: endMs
+        )
+    }
+
+    // MARK: - Enrollment (queue only)
+
+    private func processEnrollment(samples: [Float], chunkMs: Double, rms: Float, voiced: Bool) {
+        guard var session = enrollment else { return }
+        if voiced {
+            session.samples.append(contentsOf: samples)
+            session.voicedMs += chunkMs
+        }
+        session.lastLevel = rms
+
+        let now = Date()
+        if now.timeIntervalSince(session.lastProgressAt) >= 0.25 {
+            session.lastProgressAt = now
+            Bridge.sendSpeakerEnrollProgress(
+                voicedSeconds: session.voicedMs / 1000.0,
+                targetSeconds: session.targetVoicedMs / 1000.0,
+                level: Double(rms)
+            )
+        }
+
+        if session.voicedMs >= session.targetVoicedMs {
+            enrollment = nil
+            activeFlag.lock()
+            _enrolling = false
+            activeFlag.unlock()
+            completeEnrollment(session)
+        } else {
+            enrollment = session
+        }
+    }
+
+    private func completeEnrollment(_ session: EnrollmentSession) {
+        guard let extractor else {
+            finishEnrollment(ok: false, reason: "model_unavailable", embeddings: [], source: session.source)
+            return
+        }
+
+        // One embedding per prompt-sized chunk; internal consistency check
+        // catches wind/noise/multiple talkers.
+        let chunkCount = Self.enrollChunkCount
+        let chunkSize = session.samples.count / chunkCount
+        guard chunkSize >= Self.sampleRate else {
+            finishEnrollment(ok: false, reason: "too_short", embeddings: [], source: session.source)
+            return
+        }
+
+        var embeddings: [[Float]] = []
+        for i in 0 ..< chunkCount {
+            let start = i * chunkSize
+            let end = i == chunkCount - 1 ? session.samples.count : start + chunkSize
+            let chunk = Array(session.samples[start ..< end])
+            let stream = extractor.createStream()
+            stream.acceptWaveform(samples: chunk, sampleRate: Self.sampleRate)
+            stream.inputFinished()
+            guard extractor.isReady(stream: stream) else { continue }
+            let embedding = Self.l2Normalize(extractor.compute(stream: stream))
+            if !embedding.isEmpty {
+                embeddings.append(embedding)
+            }
+        }
+
+        guard embeddings.count >= 2 else {
+            finishEnrollment(ok: false, reason: "too_short", embeddings: [], source: session.source)
+            return
+        }
+
+        for i in 0 ..< embeddings.count {
+            for j in (i + 1) ..< embeddings.count {
+                let cos = Self.dot(embeddings[i], embeddings[j])
+                if cos < Self.enrollConsistencyThreshold {
+                    Bridge.log(
+                        "SpeakerID: enrollment chunks \(i)/\(j) inconsistent — cos=\(String(format: "%.3f", cos)) "
+                            + "< \(Self.enrollConsistencyThreshold) (noise / multiple talkers?)"
+                    )
+                    finishEnrollment(ok: false, reason: "inconsistent", embeddings: [], source: session.source)
+                    return
+                }
+            }
+        }
+
+        let encoded = embeddings.map { Self.base64(from: $0) }
+        finishEnrollment(ok: true, reason: nil, embeddings: encoded, source: session.source, dim: embeddings[0].count)
+    }
+
+    private func finishEnrollment(ok: Bool, reason: String?, embeddings: [String], source: String, dim: Int = 0) {
+        activeFlag.lock()
+        _enrolling = false
+        activeFlag.unlock()
+        Bridge.log(
+            ok
+                ? "SpeakerID: enrollment ok — \(embeddings.count) embeddings (dim=\(dim), mic=\(source))"
+                : "SpeakerID: enrollment failed — \(reason ?? "unknown") (mic=\(source))"
+        )
+        Bridge.sendSpeakerEnrollResult(ok: ok, reason: reason, embeddings: embeddings, dim: dim, source: source)
+        let onEnded = enrollmentEnded
+        enrollmentEnded = nil
+        onEnded?()
+    }
+
+    // MARK: - Model loading (queue only)
+
+    private func ensureExtractor() {
+        guard extractor == nil, !extractorLoadAttempted else { return }
+        extractorLoadAttempted = true
+        guard let modelPath = Self.resolvedModelPath(),
+              FileManager.default.fileExists(atPath: modelPath)
+        else {
+            Bridge.log("SpeakerID: model not downloaded yet — engine idle")
+            return
+        }
+        // Config and construction share one C-string scope; the config's
+        // `const char*` fields point at buffers owned by it (see toCPointer in
+        // SherpaOnnx.swift) and must not outlive it.
+        extractor = withSherpaOnnxCStrings {
+            var config = sherpaOnnxSpeakerEmbeddingExtractorConfig(model: modelPath, numThreads: 2)
+            return SherpaOnnxSpeakerEmbeddingExtractorWrapper(config: &config)
+        }
+        Bridge.log("SpeakerID: embedding extractor initialized (dim=\(extractor?.dim ?? -1))")
+    }
+
+    /// Same Documents-relative reconstruction as the STT model path (the app
+    /// container id changes between dev builds).
+    private static func resolvedModelPath() -> String? {
+        guard let storedPath = UserDefaults.standard.string(forKey: modelPathKey) else { return nil }
+        let documentsURL = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first!
+        if let range = storedPath.range(of: "/Documents/") {
+            let relativePath = String(storedPath[range.upperBound...])
+            return documentsURL.appendingPathComponent(relativePath).path
+        }
+        return storedPath
+    }
+
+    // MARK: - Helpers
+
+    private func resetWindow() {
+        windowSamples = []
+        windowStartMs = 0
+        windowVoicedMs = 0
+        trailingSilenceMs = 0
+    }
+
+    private static func toFloatArray(_ pcmData: Data) -> [Float] {
+        let count = pcmData.count / 2
+        var samples = [Float](repeating: 0, count: count)
+        pcmData.withUnsafeBytes { (bufferPointer: UnsafeRawBufferPointer) in
+            guard let address = bufferPointer.baseAddress else { return }
+            let int16Pointer = address.bindMemory(to: Int16.self, capacity: count)
+            for i in 0 ..< count {
+                samples[i] = Float(Int16(littleEndian: int16Pointer[i])) / 32768.0
+            }
+        }
+        return samples
+    }
+
+    private static func rms(_ samples: [Float]) -> Float {
+        guard !samples.isEmpty else { return 0 }
+        var sum: Float = 0
+        for s in samples { sum += s * s }
+        return (sum / Float(samples.count)).squareRoot()
+    }
+
+    private static func dot(_ a: [Float], _ b: [Float]) -> Float {
+        guard a.count == b.count else { return 0 }
+        var sum: Float = 0
+        for i in 0 ..< a.count { sum += a[i] * b[i] }
+        return sum
+    }
+
+    private static func l2Normalize(_ v: [Float]) -> [Float] {
+        var sum: Float = 0
+        for x in v { sum += x * x }
+        let norm = sum.squareRoot()
+        guard norm > 0 else { return v }
+        return v.map { $0 / norm }
+    }
+
+    private static func base64(from embedding: [Float]) -> String {
+        embedding.withUnsafeBufferPointer { buffer in
+            Data(buffer: buffer).base64EncodedString()
+        }
+    }
+
+    static func embedding(fromBase64 base64: String) -> [Float]? {
+        guard let data = Data(base64Encoded: base64), data.count % 4 == 0, !data.isEmpty else { return nil }
+        let count = data.count / 4
+        var floats = [Float](repeating: 0, count: count)
+        _ = floats.withUnsafeMutableBytes { data.copyBytes(to: $0) }
+        return floats
+    }
+}
+
+#else
+
+/// Stub for SWIFT_PACKAGE builds without GLASSLY_FEATURE_LOCAL_STT — the
+/// sherpa-onnx runtime isn't linked, so speaker ID is unavailable.
+final class SpeakerIdEngine {
+    static let shared = SpeakerIdEngine()
+
+    struct Profile {
+        let id: String
+        let name: String
+        let embedding: [Float]
+    }
+
+    struct Match {
+        let profileId: String
+        let name: String
+        let similarity: Float
+        let endMs: Int64
+    }
+
+    func setIdentifyEnabled(_: Bool) {}
+    func setProfiles(_: [Profile]) {}
+    static func setModelPath(_: String) {}
+    static func isModelAvailable() -> Bool { false }
+    func startEnrollment(targetSpeechSeconds _: Double, source: String, onEnded: @escaping () -> Void) {
+        Bridge.sendSpeakerEnrollResult(ok: false, reason: "unsupported", embeddings: [], dim: 0, source: source)
+        onEnded()
+    }
+
+    func cancelEnrollment() {}
+    func recentMatch(withinMs _: Int64) -> Match? { nil }
+    func acceptAudio(pcm16le _: Data) {}
+    static func embedding(fromBase64 _: String) -> [Float]? { nil }
+}
+
+#endif