diff --git a/Package.swift b/Package.swift index 4f21087..424ceb9 100644 --- a/Package.swift +++ b/Package.swift @@ -53,6 +53,8 @@ let package = Package( .target(name: "BookmarkSort", dependencies: ["FluidUse"]), .executableTarget(name: "BookmarkSortCheck", dependencies: ["BookmarkSort", "FluidUse"]), .executableTarget(name: "BookmarkSortDemo", dependencies: ["BookmarkSort"]), + .executableTarget(name: "AudioSearchCheck", dependencies: ["FluidUse"]), + .executableTarget(name: "AudioSearchDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh"]), .executableTarget( name: "TopicSortDemo", dependencies: ["BookmarkSort", "FluidUse"], exclude: ["README.md", "demo.sh"], resources: [.copy("Resources")]), diff --git a/Sources/AudioSearchCheck/main.swift b/Sources/AudioSearchCheck/main.swift new file mode 100644 index 0000000..b303ae6 --- /dev/null +++ b/Sources/AudioSearchCheck/main.swift @@ -0,0 +1,74 @@ +import Foundation +import FluidUse + +/// Indexes audio files with EmbeddingGemma 2 (10 s windows) and searches them with text. +/// +/// swift run -c release AudioSearchCheck file.wav [more files…] [--query="…"]… [--dump=vectors.jsonl] +@main +struct AudioSearchCheck { + static func main() async throws { + let arguments = CommandLine.arguments.dropFirst() + // Folders expand to the audio files inside them, sorted. + let audioExtensions: Set = ["wav", "flac", "mp3", "m4a", "aiff", "caf"] + let files = arguments.filter { !$0.hasPrefix("--") }.flatMap { path -> [URL] in + let url = URL(fileURLWithPath: path) + var isDirectory: ObjCBool = false + guard FileManager.default.fileExists(atPath: path, isDirectory: &isDirectory), isDirectory.boolValue else { + return [url] + } + let names = (try? FileManager.default.contentsOfDirectory(atPath: path)) ?? [] + return names.filter { audioExtensions.contains(($0 as NSString).pathExtension.lowercased()) }.sorted() + .map { url.appendingPathComponent($0) } + } + let queries = arguments.filter { $0.hasPrefix("--query=") }.map { String($0.dropFirst(8)) } + let dump = arguments.first { $0.hasPrefix("--dump=") }.map { String($0.dropFirst(7)) } + guard !files.isEmpty else { + print("usage: AudioSearchCheck file.wav [more files…] [--query=\"…\"]… [--dump=vectors.jsonl]") + exit(2) + } + var start = DispatchTime.now().uptimeNanoseconds + let text = try await EmbeddingGemma2Manager.loadDefault() + let audio = try await EmbeddingGemma2Audio.load(text: text) + _ = try await audio.embed(window: [Float](repeating: 0, count: 16_000)) + print(String(format: "models loaded in %.1f s", seconds(since: start))) + start = DispatchTime.now().uptimeNanoseconds + let recordings = try files.map { try EmbeddingGemma2Audio.samples(contentsOf: $0) } + let decodeSeconds = seconds(since: start) + start = DispatchTime.now().uptimeNanoseconds + let embedded = try await audio.embed(recordings: recordings) + let embedSeconds = seconds(since: start) + let audioSeconds = recordings.reduce(0.0) { $0 + Double($1.count) } / Double(EmbeddingGemma2Audio.sampleRate) + let index: [(file: String, window: EmbeddingGemma2Audio.Window)] = zip(files, embedded).flatMap { + file, windows in + windows.map { (file: file.lastPathComponent, window: $0) } + } + print( + String( + format: "%.1f min of audio, %d windows: decode %.1f s, embed %.1f s = %.0fx real time (embed only)", + audioSeconds / 60, index.count, decodeSeconds, embedSeconds, audioSeconds / max(embedSeconds, 1e-6))) + if let dump { + let lines = index.map { item -> String in + let object: [String: Any] = ["file": item.file, "start": item.window.start, "v": item.window.embedding] + return String(decoding: try! JSONSerialization.data(withJSONObject: object), as: UTF8.self) + } + try lines.joined(separator: "\n").write(toFile: dump, atomically: true, encoding: .utf8) + } + for query in queries { + let vector = try await text.embed(query, prompt: .searchQuery) + let ranked = index.map { item in + (item, zip(vector, item.window.embedding).reduce(Float(0)) { $0 + $1.0 * $1.1 }) + } + .sorted { $0.1 > $1.1 }.prefix(3) + print("\(query):") + for (item, score) in ranked { + print(String(format: " %.3f %@ @ %@", score, item.file, clock(item.window.start))) + } + } + } + + static func seconds(since start: UInt64) -> Double { Double(DispatchTime.now().uptimeNanoseconds - start) / 1e9 } + + static func clock(_ seconds: TimeInterval) -> String { + String(format: "%d:%02d", Int(seconds) / 60, Int(seconds) % 60) + } +} diff --git a/Sources/AudioSearchDemo/AudioSearchDemoApp.swift b/Sources/AudioSearchDemo/AudioSearchDemoApp.swift new file mode 100644 index 0000000..6e904d7 --- /dev/null +++ b/Sources/AudioSearchDemo/AudioSearchDemoApp.swift @@ -0,0 +1,24 @@ +import AppKit +import SwiftUI + +@main +struct AudioSearchDemoApp: App { + @StateObject private var model = AudioSearchModel() + + init() { + setvbuf(stdout, nil, _IOLBF, 0) + // Bare SwiftPM executables start as background processes; make this one a regular windowed app. + NSApplication.shared.setActivationPolicy(.regular) + NSApplication.shared.activate(ignoringOtherApps: true) + } + + var body: some Scene { + WindowGroup("Search audio — on-device") { + ContentView() + .environmentObject(model) + .frame(minWidth: 900, minHeight: 640) + .task { await model.prepare() } + } + .defaultSize(width: 1100, height: 900) + } +} diff --git a/Sources/AudioSearchDemo/AudioSearchModel.swift b/Sources/AudioSearchDemo/AudioSearchModel.swift new file mode 100644 index 0000000..baec102 --- /dev/null +++ b/Sources/AudioSearchDemo/AudioSearchModel.swift @@ -0,0 +1,655 @@ +import Accelerate +import AVFoundation +import FluidUse +import Foundation +import SwiftUI +import os + +/// Indexes audio collections with EmbeddingGemma 2 (10 s windows: audio model on the GPU, text model on the Neural +/// Engine) and searches every window with a text query. Play runs the show once: index everything, listen to the +/// top three of a few queries, then run searches back to back as fast as they go. Pause holds it anywhere (timers +/// included); Replay starts over from an empty index. +@MainActor +final class AudioSearchModel: ObservableObject { + enum Step: Equatable { + case loading(String) + case idle + case indexing + case listening + case speed + case finished + case failed(String) + } + + struct Collection: Identifiable { + let id: Int + let name: String + let color: Color + let files: [URL] + let suggestions: [String] + } + + struct Entry: Sendable { + let collection: Int + let file: URL + let start: TimeInterval + let duration: TimeInterval + let embedding: [Float] + } + + struct Result: Identifiable { + let entry: Entry + let score: Float + var id: String { "\(entry.file.path)#\(entry.start)" } + } + + @Published private(set) var step: Step = .loading("Starting…") + @Published private(set) var paused = false + @Published private(set) var collections: [Collection] = [] + @Published private(set) var windowsDone = 0 + @Published private(set) var windowsTotal = 0 + @Published private(set) var audioSeconds: Double = 0 + @Published private(set) var indexSeconds: Double = 0 + @Published private(set) var indexedWindows = 0 + @Published private(set) var results: [Result] = [] + @Published private(set) var queryMilliseconds: Double = 0 + @Published private(set) var playing: String? + @Published private(set) var stepRemaining = 0 + @Published private(set) var burstQueries = 0 + @Published private(set) var burstPerSecond: Double = 0 + @Published private(set) var burstMilliseconds: Double = 0 + @Published var query = "" + + private var entries: [Entry] = [] + /// Every window's embedding, row after row, so a query is one matrix-vector product. + private var matrix: [Float] = [] + private var text: EmbeddingGemma2Manager? + private var audio: EmbeddingGemma2Audio? + private var player: AVAudioPlayer? + private var searchTask: Task? + private var showTask: Task? + /// The last query the show put in the field; the field's change callback echoing it is not typing. + private var showQuery = "" + /// Paused time so far, so timers and speeds count only running time. + private var pausedNanoseconds: UInt64 = 0 + private var pauseBegan: UInt64? + private let gate = PauseGate() + /// Seconds of each result to play (`--clip=`) and length of the listen and speed steps (`--segment=`). + private let clipSeconds = max( + 1, CommandLine.arguments.first { $0.hasPrefix("--clip=") }.flatMap { Double($0.dropFirst(7)) } ?? 3) + private let segmentSeconds = max( + 5, CommandLine.arguments.first { $0.hasPrefix("--segment=") }.flatMap { Double($0.dropFirst(10)) } ?? 30) + + var realTimeFactor: Double { indexSeconds > 0 ? audioSeconds / indexSeconds : 0 } + var allSuggestions: [String] { collections.flatMap(\.suggestions) } + var isRunning: Bool { [.indexing, .listening, .speed].contains(step) } + var canPlay: Bool { step == .idle || (isRunning && paused) } + var canReplay: Bool { isRunning || step == .finished } + + func prepare() async { + do { + collections = Self.defaultCollections() + guard !collections.isEmpty else { + throw EmbeddingGemma2Error.invalidAsset("No audio found: pass --audio=") + } + DemoLog.event( + "Search audio · \(collections.map { "\($0.name) (\($0.files.count))" }.joined(separator: " · "))") + let loadStart = DispatchTime.now().uptimeNanoseconds + let ticker = Task { @MainActor [weak self] in + var seconds = 0 + while !Task.isCancelled { + try? await Task.sleep(for: .seconds(1)) + seconds += 1 + guard let self, case .loading = self.step else { continue } + self.step = .loading( + "Loading EmbeddingGemma 2… \(seconds) s" + + (seconds >= 8 ? " (first launch downloads and compiles the models once)" : "")) + } + } + defer { ticker.cancel() } + let text = try await EmbeddingGemma2Manager.loadDefault() + let audio = try await EmbeddingGemma2Audio.load(text: text) + _ = try await audio.embed(window: [Float](repeating: 0, count: 16_000)) + _ = try await text.embed("warm up", prompt: .searchQuery) + // The speed step embeds queries eight per call (pack_256): load that path now too. + _ = try await text.embed(Array(repeating: "warm up", count: 64), prompt: .searchQuery) + self.text = text + self.audio = audio + DemoLog.model( + String( + format: "models loaded in %.1f s · audio encoder on the GPU, text model on the Neural Engine", + Double(DispatchTime.now().uptimeNanoseconds - loadStart) / 1e9)) + step = .idle + if CommandLine.arguments.contains("--autostart") { play() } + } catch { + step = .failed(error.localizedDescription) + DemoLog.line("failed: \(error.localizedDescription)", color: 196, bold: true) + } + } + + // MARK: Play / Pause / Replay + + /// Starts the show, or resumes it after Pause. + func play() { + if isRunning, paused { + if let began = pauseBegan { pausedNanoseconds += DispatchTime.now().uptimeNanoseconds - began } + pauseBegan = nil + paused = false + gate.set(paused: false) + player?.play() + DemoLog.event("▶ resume") + return + } + guard step == .idle else { return } + showTask = Task { await show() } + } + + func pause() { + guard isRunning, !paused else { return } + paused = true + pauseBegan = DispatchTime.now().uptimeNanoseconds + gate.set(paused: true) + player?.pause() + DemoLog.event("❚❚ pause") + } + + /// Clears the index and runs the show again from the start. + func replay() { + guard canReplay else { return } + showTask?.cancel() + showTask = nil + stop() + paused = false + pauseBegan = nil + gate.set(paused: false) + entries = [] + matrix = [] + indexedWindows = 0 + results = [] + query = "" + showQuery = "" + windowsDone = 0 + windowsTotal = 0 + audioSeconds = 0 + indexSeconds = 0 + burstQueries = 0 + burstPerSecond = 0 + burstMilliseconds = 0 + step = .idle + DemoLog.event("↺ replay") + showTask = Task { await show() } + } + + /// Nanoseconds of running (unpaused) time on a monotonic clock. + private var activeNow: UInt64 { + let now = DispatchTime.now().uptimeNanoseconds + return now - pausedNanoseconds - (pauseBegan.map { now - $0 } ?? 0) + } + + private func waitWhilePaused() async { + while paused, !Task.isCancelled { try? await Task.sleep(for: .milliseconds(50)) } + } + + /// Index, listen, speed, done. + private func show() async { + DemoLog.event( + "▶ play: index everything, top 3 of \(Self.showcase.count) example searches, \(Int(segmentSeconds)) s of searches as fast as they go" + ) + await index() + guard !Task.isCancelled, indexedWindows > 0 else { return } + await listen() + guard !Task.isCancelled else { return } + stop() + await speed() + guard !Task.isCancelled else { return } + step = .finished + DemoLog.event( + String( + format: "■ done: %@ of audio indexed in %.1f s (%.0fx real time); %d searches at %.0f/s", + Self.clock(audioSeconds), indexSeconds, realTimeFactor, burstQueries, burstPerSecond)) + } + + // MARK: Index + + private func index() async { + guard let audio else { return } + step = .indexing + // Round-robin across collections. + var queues = collections.map { collection in collection.files.map { (collection.id, $0) } } + var ordered: [(Int, URL)] = [] + while queues.contains(where: { !$0.isEmpty }) { + for index in queues.indices where !queues[index].isEmpty { ordered.append(queues[index].removeFirst()) } + } + let files = ordered + DemoLog.event("📥 indexing \(files.count) files") + do { + let decodeStart = DispatchTime.now().uptimeNanoseconds + let recordings = try await Task.detached(priority: .userInitiated) { + try files.map { try EmbeddingGemma2Audio.samples(contentsOf: $0.1) } + }.value + DemoLog.line( + String( + format: "decoded %d files to 16 kHz in %.1f s", files.count, + Double(DispatchTime.now().uptimeNanoseconds - decodeStart) / 1e9), color: 245) + await waitWhilePaused() + let total = recordings.reduce(0.0) { $0 + Double($1.count) } / Double(EmbeddingGemma2Audio.sampleRate) + let start = activeNow + let progress = IndexProgress() + let inbox = WindowInbox() + let gate = gate + // The embedding tasks report into `progress` and `inbox`; the UI drains them ten times a second. + let watcher = Task { @MainActor [weak self] in + while !Task.isCancelled { + try? await Task.sleep(for: .milliseconds(100)) + guard let self else { return } + self.refresh(progress: progress, total: total, start: start) + self.add(inbox.drain()) + } + } + defer { watcher.cancel() } + _ = try await audio.embed( + recordings: recordings, progress: { done, count in progress.update(done: done, total: count) }, + onWindow: { recording, window in + inbox.append( + Entry( + collection: files[recording].0, file: files[recording].1, start: window.start, + duration: window.duration, embedding: window.embedding)) + }, + beforeWindow: { await gate.wait() }) + try Task.checkCancellation() + refresh(progress: progress, total: total, start: start) + add(inbox.drain()) + DemoLog.event( + String( + format: "📥 indexed %@ of audio (%d windows) in %.1f s = %.0fx real time", Self.clock(total), + indexedWindows, indexSeconds, realTimeFactor)) + } catch is CancellationError { + } catch { + step = .failed(error.localizedDescription) + DemoLog.line("failed: \(error.localizedDescription)", color: 196, bold: true) + } + } + + private func add(_ fresh: [Entry]) { + guard !fresh.isEmpty else { return } + entries += fresh + matrix += fresh.flatMap(\.embedding) + indexedWindows = entries.count + } + + private func refresh(progress: IndexProgress, total: Double, start: UInt64) { + let (done, count) = progress.snapshot + guard count > 0 else { return } + windowsDone = done + windowsTotal = count + audioSeconds = total * Double(done) / Double(count) + indexSeconds = Double(activeNow - start) / 1e9 + } + + // MARK: Listen + + /// The listening step's examples: each is typed, then its top three play. Two phrases from the earnings call, + /// six everyday sounds; every one returns three right hits. + static let showcase = [ + "forward-looking statements disclaimer", "a baby crying", "church bells ringing", + "the operator opens the line for questions", "a siren", "someone laughing", "a cow mooing", "birds chirping", + ] + + static let earningsExamples: Set = [ + "forward-looking statements disclaimer", "the operator opens the line for questions", + ] + + /// Types each showcase query (those whose collection is present), then plays its top three. + private func listen() async { + step = .listening + // Earnings-call examples need that collection; sound examples need the sound clips. + let names = Set(collections.map(\.name)) + var queries = Self.showcase.filter { query in + Self.earningsExamples.contains(query) ? names.contains("Earnings calls") : names.contains("Sounds (ESC-50)") + } + if queries.isEmpty { queries = Array(allSuggestions.prefix(8)) } + guard !queries.isEmpty else { return } + for (number, query) in queries.enumerated() { + guard !Task.isCancelled else { return } + stepRemaining = queries.count - number + showQuery = "" + self.query = "" + for character in query { + await waitWhilePaused() + showQuery.append(character) + self.query.append(character) + try? await Task.sleep(for: .milliseconds(30)) + } + await waitWhilePaused() + guard !Task.isCancelled, let ranked = await rank(query) else { return } + results = ranked + for (rank, result) in ranked.prefix(3).enumerated() { + guard !Task.isCancelled else { return } + DemoLog.line( + String( + format: " ▶ #%d %@ %@ @ %@", rank + 1, collections[result.entry.collection].name, + result.entry.file.lastPathComponent, Self.clock(result.entry.start)), color: 120) + await playClip(result, seconds: min(clipSeconds, result.entry.duration)) + } + await waitWhilePaused() + try? await Task.sleep(for: .milliseconds(700)) + } + } + + /// Plays `seconds` of a result's window, holding while paused; returns when it ends or is stopped. + private func playClip(_ result: Result, seconds: TimeInterval) async { + stop() + guard let player = try? AVAudioPlayer(contentsOf: result.entry.file) else { return } + player.currentTime = result.entry.start + player.play() + self.player = player + playing = result.id + var played: UInt64 = 0 + var last = DispatchTime.now().uptimeNanoseconds + while !Task.isCancelled, playing == result.id, Double(played) / 1e9 < seconds { + try? await Task.sleep(for: .milliseconds(50)) + let now = DispatchTime.now().uptimeNanoseconds + if !paused { played += now - last } + last = now + } + if playing == result.id { stop() } + } + + // MARK: Speed + + /// Queries back to back for `segmentSeconds` of running time: 64 at a time, embedded eight per Neural Engine + /// call, ranked against every window in one matrix multiply. The field shows the latest query of each batch. + private func speed() async { + guard let text, !entries.isEmpty else { return } + step = .speed + let pool = queryPool + let dimension = EmbeddingGemma2Manager.dimension + let windows = entries.count + let batchSize = 64 + burstQueries = 0 + burstPerSecond = 0 + burstMilliseconds = 0 + DemoLog.event("⚡ search speed: \(pool.count) different queries, back to back, against \(windows) windows") + let start = activeNow + let length = UInt64(segmentSeconds * 1e9) + let ticker = tick { [weak self] in + guard let self else { return } + self.stepRemaining = max(0, Int((Double(length) - Double(self.activeNow - start)) / 1e9)) + } + defer { ticker.cancel() } + var embedNanoseconds: UInt64 = 0 + var rankNanoseconds: UInt64 = 0 + var next = 0 + var lastLog = start + var recent: [(time: UInt64, count: Int)] = [(start, 0)] + while !Task.isCancelled, activeNow - start < length { + if paused { + await waitWhilePaused() + recent = [(activeNow, 0)] + continue + } + let batch = (0.. 1_000_000_000 } + let span = Double(now - (recent.first?.time ?? start)) / 1e9 + let recentQueries = recent.dropFirst().reduce(0) { $0 + $1.count } + if span > 0.2, recentQueries > 0 { + burstPerSecond = Double(recentQueries) / span + burstMilliseconds = 1000 / burstPerSecond + } + showQuery = batch[batchSize - 1] + query = showQuery + results = shown + queryMilliseconds = burstMilliseconds + if now - lastLog > 1_000_000_000 { + lastLog = now + DemoLog.line( + String( + format: + "⚡ %d searches · %.0f/s · %.2f ms each (embed %.2f + rank %.2f) · %.1f M window scores/s", + burstQueries, burstPerSecond, burstMilliseconds, + Double(embedNanoseconds) / 1e6 / Double(burstQueries), + Double(rankNanoseconds) / 1e6 / Double(burstQueries), + burstPerSecond * Double(windows) / 1e6), color: 81) + } + } + // The tiles keep the whole step's average. + let elapsed = Double(activeNow - start) / 1e9 + if elapsed > 0, burstQueries > 0 { + burstPerSecond = Double(burstQueries) / elapsed + burstMilliseconds = 1000 / burstPerSecond + } + } + + private func tick(_ update: @escaping @MainActor () -> Void) -> Task { + Task { @MainActor in + while !Task.isCancelled { + update() + try? await Task.sleep(for: .milliseconds(250)) + } + } + } + + // MARK: Search by hand (when the show is not running) + + /// Ranks every indexed window against the typed query (debounced). + func search() { + if query == showQuery || isRunning { return } + searchTask?.cancel() + let query = query.trimmingCharacters(in: .whitespacesAndNewlines) + guard !query.isEmpty, !entries.isEmpty else { + results = [] + return + } + searchTask = Task { + try? await Task.sleep(for: .milliseconds(120)) + guard !Task.isCancelled, let ranked = await rank(query), !Task.isCancelled else { return } + results = ranked + } + } + + /// Top 25 windows for `query`: one text embedding, then one matrix-vector product over every window. + private func rank(_ query: String) async -> [Result]? { + guard let text else { return nil } + let begin = DispatchTime.now().uptimeNanoseconds + guard let vector = try? await text.embed(query, prompt: .searchQuery) else { return nil } + var scores = [Float](repeating: 0, count: entries.count) + cblas_sgemv( + CblasRowMajor, CblasNoTrans, Int32(entries.count), Int32(vector.count), 1, matrix, Int32(vector.count), + vector, 1, 0, &scores, 1) + let ranked = scores.indices.sorted { scores[$0] > scores[$1] }.prefix(25).map { + Result(entry: entries[$0], score: scores[$0]) + } + queryMilliseconds = Double(DispatchTime.now().uptimeNanoseconds - begin) / 1e6 + DemoLog.line( + String(format: "🔎 “%@” → %d windows ranked in %.0f ms", query, entries.count, queryMilliseconds), color: 81) + return ranked + } + + /// Row play button (when the show is not running): plays or stops one window. + func play(_ result: Result) { + guard !isRunning else { return } + if playing == result.id { + stop() + return + } + Task { await playClip(result, seconds: result.entry.duration) } + } + + func stop() { + player?.stop() + player = nil + playing = nil + } + + static func clock(_ seconds: TimeInterval) -> String { + let total = Int(seconds.rounded()) + return total >= 3600 + ? String(format: "%d:%02d:%02d", total / 3600, total / 60 % 60, total % 60) + : String(format: "%d:%02d", total / 60, total % 60) + } + + /// Indices of the `count` highest scores in one row of a row-major score matrix, best first. + static func topIndices(_ scores: [Float], row: Int, width: Int, count: Int) -> [Int] { + var best: [(index: Int, score: Float)] = [] + best.reserveCapacity(count + 1) + for column in 0.. best[best.count - 1].score { + let position = best.firstIndex { score > $0.score } ?? best.count + best.insert((column, score), at: position) + if best.count > count { best.removeLast() } + } + } + return best.map(\.index) + } + + /// Everything the search-speed segment cycles through: the suggestions plus a spread of topics and sounds. + var queryPool: [String] { + var pool = Self.showcase + allSuggestions + Self.extraQueries + // Sound clips have random file names; their classes come from labels.json next to the folder. + var sounds = Set() + for collection in collections { + guard let folder = collection.files.first?.deletingLastPathComponent(), + let data = try? Data( + contentsOf: folder.deletingLastPathComponent().appendingPathComponent("labels.json")), + let labels = try? JSONSerialization.jsonObject(with: data) as? [String: [String: String]] + else { continue } + for label in labels.values { + if let category = label["category"] { + sounds.insert("the sound of " + category.replacingOccurrences(of: "_", with: " ")) + } + } + } + pool += sounds.sorted() + var seen = Set() + return pool.filter { seen.insert($0).inserted } + } + + static let extraQueries = [ + "revenue growth this quarter", "guidance for next year", "operating margins improved", + "cash flow and dividends", + "share buyback program", "supply chain problems", "demand from customers", "pricing pressure", + "new product launch", "hiring and headcount", "interest rates and inflation", "currency headwinds", + "an acquisition of a company", "questions from analysts", "thank you for joining the call", + "the chief executive officer speaks", "regulation and government policy", "research and development spending", + "the weather and climate", "a scientific discovery", "a sports team won the game", "a famous city in Europe", + "the history of a war", "a recipe for dinner", "travelling by train", "music and concerts", + "a disease and its treatment", "planets and space", "an election and voters", "children at school", + "a river and a mountain", "computers and the internet", "money and banks", "an old church", + "farming and crops", "the ocean and fish", "a museum exhibition", "an earthquake", "a new law was passed", + "a festival with fireworks", + ] + + /// `--audio=` (repeatable), else the public datasets already on this Mac. + static func defaultCollections() -> [Collection] { + let palette: [Color] = [.blue, .orange, .green, .pink, .purple, .teal] + let given = CommandLine.arguments.filter { $0.hasPrefix("--audio=") }.map { String($0.dropFirst(8)) } + if !given.isEmpty { + return given.enumerated().compactMap { index, path in + let files = audioFiles(at: URL(fileURLWithPath: path)) + guard !files.isEmpty else { return nil } + return Collection( + id: index, name: URL(fileURLWithPath: path).lastPathComponent, + color: palette[index % palette.count], files: files, suggestions: []) + } + } + let fluidAudio = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0] + .appendingPathComponent("FluidAudio") + let fluidUse = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0] + .appendingPathComponent("FluidUse/Datasets") + let candidates: [(String, URL, [String])] = [ + ( + "Earnings calls", fluidAudio.appendingPathComponent("earnings22-1h"), + ["forward-looking statements disclaimer", "the operator opens the line for questions"] + ), + ( + "FLEURS English", fluidAudio.appendingPathComponent("FLEURS/en_us"), + ["the history of an ancient empire", "animals in the wild"] + ), + ("FLEURS French", fluidAudio.appendingPathComponent("FLEURS/fr_fr"), ["a story about a famous painter"]), + ( + "Sounds (ESC-50)", fluidUse.appendingPathComponent("esc50/wav"), + ["a baby crying", "someone laughing", "a dog barking", "church bells"] + ), + ] + var collections: [Collection] = [] + for candidate in candidates { + let files = audioFiles(at: candidate.1) + guard !files.isEmpty else { continue } + let index = collections.count + collections.append( + Collection( + id: index, name: candidate.0, color: palette[index % palette.count], files: files, + suggestions: candidate.2)) + } + return collections + } + + static func audioFiles(at url: URL) -> [URL] { + let extensions: Set = ["wav", "flac", "mp3", "m4a", "aiff", "caf"] + var isDirectory: ObjCBool = false + guard FileManager.default.fileExists(atPath: url.path, isDirectory: &isDirectory) else { return [] } + guard isDirectory.boolValue else { return [url] } + let names = (try? FileManager.default.contentsOfDirectory(atPath: url.path)) ?? [] + return names.filter { extensions.contains(($0 as NSString).pathExtension.lowercased()) }.sorted() + .map { url.appendingPathComponent($0) } + } +} + +/// Windows finished by the embedding tasks, waiting for the UI to add them to the index. +final class WindowInbox: Sendable { + private let pending = OSAllocatedUnfairLock(initialState: [AudioSearchModel.Entry]()) + + func append(_ entry: AudioSearchModel.Entry) { pending.withLock { $0.append(entry) } } + + func drain() -> [AudioSearchModel.Entry] { + pending.withLock { entries in + defer { entries = [] } + return entries + } + } +} + +/// Holds embedding tasks while the show is paused. +final class PauseGate: Sendable { + private let state = OSAllocatedUnfairLock(initialState: false) + + func set(paused: Bool) { state.withLock { $0 = paused } } + + func wait() async { + while state.withLock({ $0 }), !Task.isCancelled { try? await Task.sleep(for: .milliseconds(50)) } + } +} + +/// Window counts from the embedding tasks, read by the UI on a timer. +final class IndexProgress: Sendable { + private let counts = OSAllocatedUnfairLock(initialState: (done: 0, total: 0)) + + func update(done: Int, total: Int) { counts.withLock { $0 = (done, total) } } + + var snapshot: (Int, Int) { counts.withLock { ($0.done, $0.total) } } +} diff --git a/Sources/AudioSearchDemo/ContentView.swift b/Sources/AudioSearchDemo/ContentView.swift new file mode 100644 index 0000000..b2cb8ee --- /dev/null +++ b/Sources/AudioSearchDemo/ContentView.swift @@ -0,0 +1,246 @@ +import SwiftUI + +struct ContentView: View { + @EnvironmentObject private var model: AudioSearchModel + + var body: some View { + VStack(spacing: 0) { + Header() + Divider() + SearchBar() + Divider() + Results() + } + .background(Color(nsColor: .windowBackgroundColor)) + } +} + +struct Header: View { + @EnvironmentObject private var model: AudioSearchModel + + var body: some View { + VStack(alignment: .leading, spacing: 10) { + HStack(alignment: .center, spacing: 10) { + Text("Search audio").font(.title2.bold()) + Spacer() + Text(status).font(.callout).foregroundStyle(.secondary) + Button { + model.play() + } label: { + Label(model.paused ? "Resume" : "Play", systemImage: "play.fill") + } + .keyboardShortcut(.return, modifiers: []) + .disabled(!model.canPlay) + Button { + model.pause() + } label: { + Label("Pause", systemImage: "pause.fill") + } + .keyboardShortcut(.space, modifiers: []) + .disabled(!model.isRunning || model.paused) + Button { + model.replay() + } label: { + Label("Replay", systemImage: "arrow.counterclockwise") + } + .keyboardShortcut("r", modifiers: [.command]) + .disabled(!model.canReplay) + } + GeometryReader { proxy in + Capsule().fill(Color.primary.opacity(0.08)) + .overlay(alignment: .leading) { + Capsule().fill(Color.accentColor) + .frame( + width: proxy.size.width * Double(model.windowsDone) + / Double(max(model.windowsTotal, 1))) + } + } + .frame(height: 6) + HStack(spacing: 10) { + Tile(value: AudioSearchModel.clock(model.audioSeconds), caption: "audio indexed") + Tile(value: String(format: "%.1f s", model.indexSeconds), caption: "indexing time") + Tile( + value: model.realTimeFactor == 0 ? "–" : String(format: "%.0f×", model.realTimeFactor), + caption: "faster than real time") + Tile(value: "\(model.indexedWindows)", caption: "windows searchable") + } + HStack(spacing: 10) { + Tile(value: model.burstQueries == 0 ? "–" : "\(model.burstQueries)", caption: "searches run") + Tile( + value: model.burstPerSecond == 0 ? "–" : String(format: "%.0f", model.burstPerSecond), + caption: "searches / second") + Tile( + value: model.burstMilliseconds == 0 ? "–" : String(format: "%.2f ms", model.burstMilliseconds), + caption: "per search") + Tile( + value: model.burstPerSecond == 0 + ? "–" : String(format: "%.1f M", model.burstPerSecond * Double(model.indexedWindows) / 1e6), + caption: "audio windows ranked / s") + } + .opacity(model.step == .speed || model.step == .finished ? 1 : 0.55) + HStack(spacing: 8) { + ForEach(model.collections) { collection in + HStack(spacing: 5) { + Circle().fill(collection.color).frame(width: 8, height: 8) + Text("\(collection.name) · \(collection.files.count)") + } + .font(.caption) + .padding(.horizontal, 8).padding(.vertical, 3) + .background(Capsule().fill(collection.color.opacity(0.12))) + } + Spacer() + Text("Sound clips have random file names: the model only hears the audio") + .font(.caption.bold()).foregroundStyle(.secondary) + } + HStack { + Spacer() + Text("EmbeddingGemma 2 · audio on the GPU, text on the Neural Engine · nothing leaves this Mac") + .font(.caption).foregroundStyle(.secondary) + } + } + .padding(16) + } + + private var status: String { + let pause = model.paused ? " · paused" : "" + switch model.step { + case .loading(let message): return message + case .idle: return "Ready · press Play" + case .indexing: return "📥 Indexing \(model.windowsDone)/\(model.windowsTotal) windows" + pause + case .listening: + return + "🔎 Top 3 for each example search · \(AudioSearchModel.showcase.count - model.stepRemaining + 1) of \(AudioSearchModel.showcase.count)" + + pause + case .speed: return "⚡ Searching as fast as it can · \(model.stepRemaining) s left" + pause + case .finished: return "Done · Replay to run it again" + case .failed(let message): return message + } + } +} + +struct SearchBar: View { + @EnvironmentObject private var model: AudioSearchModel + + var body: some View { + VStack(alignment: .leading, spacing: 8) { + HStack { + Image(systemName: "magnifyingglass").foregroundStyle(.secondary) + TextField("Describe what you're looking for: a phrase, a topic, a sound…", text: $model.query) + .textFieldStyle(.plain).font(.title3) + .onSubmit { model.search() } + .onChange(of: model.query) { model.search() } + .disabled(model.isRunning) + if model.queryMilliseconds > 0, !model.results.isEmpty { + Text(String(format: "%.0f ms", model.queryMilliseconds)).font(.caption.monospacedDigit()) + .foregroundStyle(.secondary) + } + } + ScrollView(.horizontal, showsIndicators: false) { + HStack(spacing: 6) { + ForEach(model.allSuggestions, id: \.self) { suggestion in + Button(suggestion) { + model.query = suggestion + model.search() + } + .buttonStyle(.bordered).controlSize(.small) + .disabled(model.isRunning) + } + } + } + } + .padding(.horizontal, 16).padding(.vertical, 12) + } +} + +struct Results: View { + @EnvironmentObject private var model: AudioSearchModel + + var body: some View { + ScrollView { + LazyVStack(spacing: 0) { + ForEach(Array(model.results.enumerated()), id: \.element.id) { rank, result in + ResultRow(rank: rank + 1, result: result) + Divider() + } + } + } + .overlay { + if model.results.isEmpty { + Text(model.step == .finished ? "Type a query or pick a suggestion" : "Press Play") + .foregroundStyle(.secondary) + } + } + } +} + +struct ResultRow: View { + @EnvironmentObject private var model: AudioSearchModel + let rank: Int + let result: AudioSearchModel.Result + + var body: some View { + let collection = model.collections[result.entry.collection] + HStack(spacing: 12) { + Text("\(rank)").font(.callout.monospacedDigit()).foregroundStyle(.secondary).frame(width: 24) + Button { + model.play(result) + } label: { + Image(systemName: model.playing == result.id ? "stop.circle.fill" : "play.circle.fill") + .font(.system(size: 26)).foregroundStyle(collection.color) + } + .buttonStyle(.plain) + .disabled(model.isRunning) + VStack(alignment: .leading, spacing: 3) { + HStack(spacing: 6) { + Text(collection.name).font(.caption.bold()).foregroundStyle(collection.color) + Text(result.entry.file.lastPathComponent).font(.callout).lineLimit(1) + } + Text( + "\(AudioSearchModel.clock(result.entry.start)) – " + + AudioSearchModel.clock(result.entry.start + result.entry.duration) + ) + .font(.caption.monospacedDigit()).foregroundStyle(.secondary) + } + Spacer() + ScoreBar(score: result.score, color: collection.color) + } + .padding(.horizontal, 16).padding(.vertical, 8) + .background(model.playing == result.id ? collection.color.opacity(0.08) : .clear) + } +} + +struct ScoreBar: View { + let score: Float + let color: Color + + var body: some View { + HStack(spacing: 6) { + GeometryReader { proxy in + Capsule().fill(Color.primary.opacity(0.08)) + .overlay(alignment: .leading) { + // Cosine scores here sit around 0.5–0.8; stretch that range across the bar. + Capsule().fill(color) + .frame(width: proxy.size.width * CGFloat(min(max((score - 0.5) / 0.35, 0.03), 1))) + } + } + .frame(width: 120, height: 6) + Text(String(format: "%.3f", score)).font(.caption.monospacedDigit()).foregroundStyle(.secondary) + } + } +} + +struct Tile: View { + let value: String + let caption: String + + var body: some View { + VStack(alignment: .leading, spacing: 2) { + Text(value).font(.system(size: 24, weight: .semibold, design: .rounded)).monospacedDigit() + .contentTransition(.numericText()) + Text(caption).font(.caption).foregroundStyle(.secondary) + } + .frame(maxWidth: .infinity, alignment: .leading) + .padding(10) + .background(RoundedRectangle(cornerRadius: 10).fill(Color.primary.opacity(0.05))) + } +} diff --git a/Sources/AudioSearchDemo/DemoLog.swift b/Sources/AudioSearchDemo/DemoLog.swift new file mode 100644 index 0000000..7e5ddcd --- /dev/null +++ b/Sources/AudioSearchDemo/DemoLog.swift @@ -0,0 +1,29 @@ +import Foundation + +/// Colour-coded console lines for the side terminal (`demo.sh` tails stdout): model load, indexing, searches. +enum DemoLog { + /// 256-colour codes matching the collection colours. + private static let topicColors = [33, 208, 34, 205, 135, 37, 196, 61, 130, 121] + + static func line(_ text: String, color: Int? = nil, bold: Bool = false) { + let time = timeFormatter.string(from: Date()) + var styled = text + if let color { styled = "\u{1B}[38;5;\(color)m\(styled)\u{1B}[0m" } + if bold { styled = "\u{1B}[1m\(styled)\u{1B}[0m" } + print("\u{1B}[2m\(time)\u{1B}[0m \(styled)") + } + + static func colored(_ text: String, slot: Int) -> String { + "\u{1B}[38;5;\(topicColors[slot % topicColors.count])m\(text)\u{1B}[0m" + } + + static func event(_ text: String) { line(text, color: 220, bold: true) } + + static func model(_ text: String) { line("[model] " + text, color: 81) } + + private static let timeFormatter: DateFormatter = { + let formatter = DateFormatter() + formatter.dateFormat = "HH:mm:ss.SSS" + return formatter + }() +} diff --git a/Sources/AudioSearchDemo/README.md b/Sources/AudioSearchDemo/README.md new file mode 100644 index 0000000..cfafc31 --- /dev/null +++ b/Sources/AudioSearchDemo/README.md @@ -0,0 +1,42 @@ +# Search audio + +Indexes hours of audio in seconds, then finds moments in it from a text query: a phrase someone said, a topic, a +sound. Everything runs on the Mac with EmbeddingGemma 2: its audio encoder turns each 10 s window into tokens on the +GPU, and its text model (Neural Engine) embeds those tokens into the same space as text queries. + +```bash +swift run -c release AudioSearchDemo # the public datasets already on this Mac (see below) +swift run -c release AudioSearchDemo --audio=~/Podcasts --audio=meeting.m4a +Sources/AudioSearchDemo/demo.sh # same app, plus one terminal: macmon above the live log +``` + +The first run downloads the models (about 1.1 GB, checksum-verified) from +[FluidInference/embeddinggemma-2-coreml](https://huggingface.co/FluidInference/embeddinggemma-2-coreml) and compiles +them once (a few minutes). Needs macOS 15+. + +- **Play** (↩) runs the show once, in order: **📥 index** every file (16 kHz mono, 10 s windows across all files with + four in flight, so the GPU and the Neural Engine work at the same time), then **🔎 listen** to eight example searches (each + typed, then its top three windows play, 3 s each), then **⚡ search speed** for 30 s (~100 different queries + back to back, 64 at a time, embedded eight per Neural Engine call, ranked against every window with one matrix + multiply). **Pause** (Space) holds it anywhere, timers included; **Replay** (⌘R) starts over from an empty index. +- When it is done, type a query (or pick a suggestion); ▶ plays a window. `--segment=` and `--clip=` change the step + and clip lengths; `--autostart` presses Play on launch + +Default audio, when present: the 1-hour Earnings-22 sample and FLEURS English/French used by FluidAudio's +benchmarks (`~/Library/Application Support/FluidAudio`), and ESC-50 sound clips +(`~/Library/Application Support/FluidUse/Datasets/esc50/wav`), renamed to random ids (`clip_3f9a1c.wav`) so the file +name cannot give the answer away; the classes live in `esc50/labels.json`. Nothing is bundled. + +Numbers (M5 Pro, macOS 27): + +| | | +|---|---| +| 1 hour earnings call, 360 windows | 4.6–5.7 s, 626–777× real time | +| 4 h 36 min (2,301 files incl. 2,000 five-second clips) | 34.7 s, 478× real time | +| Search speed over those 2,786 windows | ~700 searches/s, 1.4 ms each (query embedding 1.4 ms on the ANE, ranking 0.02 ms) | +| Swift vs sentence-transformers (fp32) | cosine 0.9995 on a LibriSpeech clip | + +Quality: speech search works well when the query names something said ("forward-looking statements" finds the +safe-harbor sentence; "operator opens the line" finds the call-in prompt). Sound events are weaker: zero-shot ESC-50 +accuracy is 24% over 50 classes (PyTorch fp32 gets the same, so that is the model, not the conversion); distinctive +sounds such as crying or laughing work, similar ones (rain vs waves vs helicopter) get mixed up. diff --git a/Sources/AudioSearchDemo/demo.sh b/Sources/AudioSearchDemo/demo.sh new file mode 100755 index 0000000..1b8ec43 --- /dev/null +++ b/Sources/AudioSearchDemo/demo.sh @@ -0,0 +1,34 @@ +#!/bin/zsh +# Search audio demo: the app, plus one terminal (Ghostty if installed) with macmon (GPU / ANE / power; asitop as a +# fallback) on top and the live Neural Engine log below. +# Sources/AudioSearchDemo/demo.sh [app options, e.g. --autostart --auto-split --rate=200] +set -e +ROOT=$(cd "$(dirname "$0")/../.." && pwd) +LOG=${TMPDIR:-/tmp}/audio-search-demo.log +# macmon (brew install macmon) needs no sudo and works on M5; asitop 0.0.24 crashes there (KeyError E0-Cluster_active). +if command -v macmon >/dev/null; then + MONITOR="$(command -v macmon)" +else + MONITOR="sudo sh -c 'pkill -x powermetrics; exec $(command -v asitop || echo "$HOME/.local/bin/asitop")'" +fi + +swift build -c release --product AudioSearchDemo --package-path "$ROOT" +: > "$LOG" +# One tmux session, reused across launches: the monitor on top, the log below. +if ! tmux has-session -t audio-search 2>/dev/null; then + # The pane stays open if the monitor exits, so the reason is visible. + tmux new-session -d -s audio-search -x 160 -y 70 \ + "$MONITOR; echo; echo \"monitor exited (\$?) — press Enter to retry\"; read; $MONITOR" + tmux set-option -t audio-search remain-on-exit on + tmux split-window -v -t audio-search "tail -n 300 -F '$LOG'" + if [[ -d /Applications/Ghostty.app ]]; then + open -na Ghostty --args -e "$(command -v tmux)" attach -t audio-search + else + osascript -e 'tell application "Terminal" to do script "tmux attach -t audio-search"' \ + -e 'tell application "Terminal" to activate' >/dev/null + fi +fi +pkill -f "release/AudioSearchDemo" 2>/dev/null || true +cd "$ROOT" +"$ROOT/.build/release/AudioSearchDemo" "$@" > "$LOG" 2> "$LOG.stderr" & +echo "demo pid $! · log $LOG · press Play (↩) · Pause (Space) · Replay (⌘R)" diff --git a/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Audio.swift b/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Audio.swift new file mode 100644 index 0000000..33c5fe1 --- /dev/null +++ b/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Audio.swift @@ -0,0 +1,190 @@ +@preconcurrency import AVFoundation +@preconcurrency import CoreML +import Foundation + +/// EmbeddingGemma 2 audio embeddings: 10 s windows of 16 kHz mono audio go through `EmbeddingGemma2Audio` (log-mel + +/// USM conformer, on the GPU, where it runs 3x faster than on the Neural Engine) and then the text model (Neural +/// Engine), so audio lands in the same space as text and can be searched with a text query. +public final class EmbeddingGemma2Audio: Sendable { + public static let sampleRate = 16_000 + public static let windowSamples = 160_000 + static let leftPad = 160 + static let frames = 1_000 + public static let maxInFlight = 4 + + /// One embedded window of a recording. + public struct Window: Sendable { + public let start: TimeInterval + public let duration: TimeInterval + public let embedding: [Float] + } + + public let text: EmbeddingGemma2Manager + private let model: MLModel + + init(text: EmbeddingGemma2Manager, model: MLModel) { + self.text = text + self.model = model + } + + /// Downloads (once, checksum-verified) and loads the audio model next to an already loaded text model. + public static func load( + text: EmbeddingGemma2Manager, computeUnits: MLComputeUnits = .cpuAndGPU, + progress: EmbeddingGemma2ModelStore.Progress? = nil + ) async throws -> EmbeddingGemma2Audio { + let directory: URL + if let path = ProcessInfo.processInfo.environment["EMBEDDINGGEMMA2_MODEL_DIR"], !path.isEmpty { + directory = URL(fileURLWithPath: path) + } else { + directory = try await EmbeddingGemma2ModelStore.ensureAudio(progress: progress) + } + let compiled = directory.appendingPathComponent("EmbeddingGemma2Audio.mlmodelc") + let package = directory.appendingPathComponent("EmbeddingGemma2Audio.mlpackage") + let url: URL + if FileManager.default.fileExists(atPath: compiled.path) { + url = compiled + } else if FileManager.default.fileExists(atPath: package.path) { + url = try await MLModel.compileModel(at: package) + } else { + throw EmbeddingGemma2Error.invalidAsset("Missing EmbeddingGemma2Audio in \(directory.path)") + } + let configuration = MLModelConfiguration() + configuration.computeUnits = computeUnits + return EmbeddingGemma2Audio( + text: text, model: try await MLModel.load(contentsOf: url, configuration: configuration)) + } + + /// Decodes any file AVFoundation reads to 16 kHz mono Float samples. + public static func samples(contentsOf url: URL) throws -> [Float] { + let file = try AVAudioFile(forReading: url) + guard + let target = AVAudioFormat( + commonFormat: .pcmFormatFloat32, sampleRate: Double(sampleRate), channels: 1, interleaved: false), + let converter = AVAudioConverter(from: file.processingFormat, to: target) + else { throw EmbeddingGemma2Error.invalidAsset("Cannot convert \(url.lastPathComponent) to 16 kHz mono") } + let chunk: AVAudioFrameCount = 65_536 + var result: [Float] = [] + result.reserveCapacity(Int(Double(file.length) * Double(sampleRate) / file.processingFormat.sampleRate) + 1) + while true { + let capacity = + AVAudioFrameCount(Double(chunk) * Double(sampleRate) / file.processingFormat.sampleRate) + 1024 + guard let output = AVAudioPCMBuffer(pcmFormat: target, frameCapacity: capacity) else { break } + var error: NSError? + // Once the file is exhausted every read returns no frames, which ends the stream; no state to share. + let status = converter.convert(to: output, error: &error) { _, outStatus in + guard let input = AVAudioPCMBuffer(pcmFormat: file.processingFormat, frameCapacity: chunk), + (try? file.read(into: input, frameCount: chunk)) != nil, input.frameLength > 0 + else { + outStatus.pointee = .endOfStream + return nil + } + outStatus.pointee = .haveData + return input + } + if let error { throw error } + if let channel = output.floatChannelData?[0], output.frameLength > 0 { + result.append(contentsOf: UnsafeBufferPointer(start: channel, count: Int(output.frameLength))) + } + if status == .endOfStream || status == .error { break } + } + return result + } + + /// Embeddings of consecutive 10 s windows of `samples` (16 kHz mono), in order. `progress` reports windows done. + public func embed( + samples: [Float], progress: (@Sendable (_ done: Int, _ total: Int) -> Void)? = nil + ) async throws -> [Window] { + try await embed(recordings: [samples], progress: progress).first ?? [] + } + + /// Windows of several recordings, embedded together so short files still keep the GPU and the Neural Engine + /// busy (one file at a time leaves a single window in flight). Result: one window list per recording, in order. + /// `onWindow` receives each window as soon as it is embedded (completion order), for indexes that grow live; + /// `beforeWindow` is awaited before each window starts (a pause gate, for example). + public func embed( + recordings: [[Float]], progress: (@Sendable (_ done: Int, _ total: Int) -> Void)? = nil, + onWindow: (@Sendable (_ recording: Int, _ window: Window) -> Void)? = nil, + beforeWindow: (@Sendable () async -> Void)? = nil + ) async throws -> [[Window]] { + var jobs: [(recording: Int, start: Int)] = [] + for (recording, samples) in recordings.enumerated() { + var starts = Array(stride(from: 0, to: samples.count, by: Self.windowSamples)) + // A tail under 0.5 s carries almost nothing and may be too short to produce a token. + if starts.count > 1, let last = starts.last, samples.count - last < Self.sampleRate / 2 { + starts.removeLast() + } + jobs += starts.map { (recording, $0) } + } + return try await withThrowingTaskGroup(of: (Int, Window).self) { group in + var windows = [Window?](repeating: nil, count: jobs.count) + var next = 0 + var done = 0 + func addJob() { + guard next < jobs.count else { return } + let index = next + next += 1 + let job = jobs[index] + let samples = recordings[job.recording] + let slice = Array(samples[job.start.. Int { + let real = leftPad + min(samples, windowSamples) + return real < 321 ? 0 : min(frames, (real - 321) / 160 + 1) + } + + /// Audio tokens for `validFrames` frames: two stride-2 convolutions, so token k is valid when frame 4k is. + static func tokenCount(validFrames: Int) -> Int { (validFrames + 3) / 4 } + + /// Embedding of up to 10 s of 16 kHz mono audio. + public func embed(window: [Float]) async throws -> [Float] { + let count = min(window.count, Self.windowSamples) + let waveform = try MLMultiArray( + shape: [1, NSNumber(value: Self.windowSamples + Self.leftPad)], dataType: .float32) + let mask = try MLMultiArray(shape: [1, NSNumber(value: Self.frames)], dataType: .float16) + let wave = waveform.dataPointer.assumingMemoryBound(to: Float.self) + wave.update(repeating: 0, count: Self.windowSamples + Self.leftPad) + window.withUnsafeBufferPointer { (wave + Self.leftPad).update(from: $0.baseAddress!, count: count) } + let maskPointer = mask.dataPointer.assumingMemoryBound(to: Float16.self) + let validFrames = Self.validFrames(samples: count) + for frame in 0.. 0 else { throw EmbeddingGemma2Error.predictionFailed("window too short to embed") } + return try await text.embed(audioTokens: tokens, count: tokenCount) + } +} diff --git a/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Manager.swift b/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Manager.swift index 4d6b3b2..c245b1f 100644 --- a/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Manager.swift +++ b/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Manager.swift @@ -206,15 +206,55 @@ public final class EmbeddingGemma2Manager: Sendable { } private func embed(ids: [Int32]) async throws -> [Float] { + let padded = ids + [Int32](repeating: tokenizer.padId, count: max(Self.lengths.last!, ids.count) - ids.count) let length = Self.lengths.first { $0 >= ids.count } ?? Self.lengths.last! + return try await run(length: length, count: ids.count) { destination in + writeEmbeddings(Array(padded.prefix(length)), to: destination) + } + } + + /// Embedding of one audio window: ` <|audio> tokens ` through the text model, where `tokens` + /// are the first `count` rows of an `EmbeddingGemma2Audio` output ([1, 250, 512], fp16 or fp32). + public func embed(audioTokens tokens: MLMultiArray, count: Int) async throws -> [Float] { + let total = count + 4 + guard let length = Self.lengths.first(where: { $0 >= total }) else { + throw EmbeddingGemma2Error.predictionFailed("\(count) audio tokens do not fit") + } + let tokenStride = tokens.strides[tokens.strides.count - 2].intValue + return try await run(length: length, count: total) { destination in + writeEmbeddings([bosId, Self.audioStartId], to: destination) + let target = destination + 2 * Self.hiddenSize + for row in 0.. + static let audioEndId: Int32 = 258_883 // + private var bosId: Int32 { tokenizer.bosId } + private var eosId: Int32 { tokenizer.eosId } + + /// One `embed_` call: `fill` writes `length` rows of fp16 inputs; the first `count` are attended. + private func run( + length: Int, count: Int, fill: (UnsafeMutablePointer) throws -> Void + ) async throws -> [Float] { guard let model = models[length] else { throw EmbeddingGemma2Error.predictionFailed("no model for \(length)") } let embeds = try MLMultiArray( shape: [1, NSNumber(value: length), NSNumber(value: Self.hiddenSize)], dataType: .float16) let mask = try MLMultiArray(shape: [1, NSNumber(value: length)], dataType: .float16) let maskPointer = mask.dataPointer.assumingMemoryBound(to: Float16.self) - let padded = ids + [Int32](repeating: tokenizer.padId, count: length - ids.count) - writeEmbeddings(padded, to: embeds.dataPointer.assumingMemoryBound(to: Float16.self)) - for position in 0...mlmodelc` beside it). Text and audio +/// are separate bundles, each pinned to the revision that added it, so adding audio never recompiles the text model. public enum EmbeddingGemma2ModelStore { public typealias Progress = @Sendable (_ file: String, _ bytes: Int64) -> Void @@ -16,6 +17,20 @@ public enum EmbeddingGemma2ModelStore { let sha256: String } + static let audioRevision = "c0fefd0ab4c1140ab35fb582ddbe51b4f93205e1" + + private static let audioAssets = [ + Asset( + path: "EmbeddingGemma2Audio.mlpackage/Manifest.json", + sha256: "65c8e9811516e7916a3d742261c7d5a02d4e43c113b437330ce86f98b3da3514"), + Asset( + path: "EmbeddingGemma2Audio.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "d40896e5c88d62b02d15cc975f27e254c28a5774a1702e09bf5483071623548d"), + Asset( + path: "EmbeddingGemma2Audio.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "41d7b552caabe376e377dcbe1c8c5f9879853bc4f384a161bb1f886f19e84b56"), + ] + private static let assets = [ Asset(path: "config.json", sha256: "4ff891a28a5d55562b55ff3aa4ef362968d378629f9453f337cb70a9e1b45f61"), Asset(path: "tokenizer.json", sha256: "4d777ef5bdc1aa36227abdfb77c3e49e7b9c892d16e1b6bda41c393504828be4"), @@ -31,15 +46,33 @@ public enum EmbeddingGemma2ModelStore { sha256: "75b97ab036ca87eae60560c4034a63d604201498dc4b6f6971dd0b1e4558d61f"), ] - /// Ensures every file exists and matches its checksum, and that the compiled model exists; returns the directory. + /// Ensures the text model, token table and tokenizer exist, match their checksums, and are compiled; returns the + /// directory. public static func ensure(cacheDirectory: URL? = nil, progress: Progress? = nil) async throws -> URL { + try await ensure( + assets: assets, revision: revision, package: "EmbeddingGemma2Text", verifiedName: ".verified-\(revision)", + cacheDirectory: cacheDirectory, progress: progress) + } + + /// Ensures the audio model (`EmbeddingGemma2Audio`) exists, matches its checksums, and is compiled; returns the + /// directory. The text bundle is needed too (`ensure`): audio tokens are embedded by the text model. + public static func ensureAudio(cacheDirectory: URL? = nil, progress: Progress? = nil) async throws -> URL { + try await ensure( + assets: audioAssets, revision: audioRevision, package: "EmbeddingGemma2Audio", + verifiedName: ".verified-audio-\(audioRevision)", cacheDirectory: cacheDirectory, progress: progress) + } + + private static func ensure( + assets: [Asset], revision: String, package: String, verifiedName: String, cacheDirectory: URL?, + progress: Progress? + ) async throws -> URL { let root = cacheDirectory ?? LayaModelStore.defaultCacheDirectory() let directory = root.appendingPathComponent(directoryName) let manager = FileManager.default try manager.createDirectory(at: directory, withIntermediateDirectories: true) - // Hashing ~550 MB takes a moment; skip it once this revision has been verified and nothing is missing. - let verified = directory.appendingPathComponent(".verified-\(revision)") - let compiled = directory.appendingPathComponent("EmbeddingGemma2Text.mlmodelc") + // Hashing hundreds of MB takes a moment; skip it once this revision has been verified and nothing is missing. + let verified = directory.appendingPathComponent(verifiedName) + let compiled = directory.appendingPathComponent("\(package).mlmodelc") let stamp = compiled.appendingPathComponent("fluiduse-revision") let allPresent = assets.allSatisfy { manager.fileExists(atPath: directory.appendingPathComponent($0.path).path) @@ -78,14 +111,13 @@ public enum EmbeddingGemma2ModelStore { // Recompile unless the compiled model carries this revision's stamp (a revision bump or an interrupted // compile leaves an old or missing stamp). Staged rename, so a reader never sees a half-written bundle. if (try? String(contentsOf: stamp, encoding: .utf8)) != revision { - progress?("EmbeddingGemma2Text.mlmodelc", 0) - let temporary = try await MLModel.compileModel( - at: directory.appendingPathComponent("EmbeddingGemma2Text.mlpackage")) + progress?("\(package).mlmodelc", 0) + let temporary = try await MLModel.compileModel(at: directory.appendingPathComponent("\(package).mlpackage")) try revision.write( to: temporary.appendingPathComponent("fluiduse-revision"), atomically: true, encoding: .utf8) - let staged = directory.appendingPathComponent("EmbeddingGemma2Text.\(UUID().uuidString).mlmodelc") + let staged = directory.appendingPathComponent("\(package).\(UUID().uuidString).mlmodelc") try manager.moveItem(at: temporary, to: staged) - let retired = directory.appendingPathComponent("EmbeddingGemma2Text.\(UUID().uuidString).old") + let retired = directory.appendingPathComponent("\(package).\(UUID().uuidString).old") if manager.fileExists(atPath: compiled.path) { try manager.moveItem(at: compiled, to: retired) } try manager.moveItem(at: staged, to: compiled) try? manager.removeItem(at: retired) diff --git a/Tests/FluidUseTests/EmbeddingGemma2Tests.swift b/Tests/FluidUseTests/EmbeddingGemma2Tests.swift index bd90dd0..b0746d4 100644 --- a/Tests/FluidUseTests/EmbeddingGemma2Tests.swift +++ b/Tests/FluidUseTests/EmbeddingGemma2Tests.swift @@ -88,3 +88,24 @@ final class EmbeddingGemma2Tests: XCTestCase { } } } + +/// Window arithmetic for the audio path, against the Hugging Face feature extractor's counts. +final class EmbeddingGemma2AudioWindowTests: XCTestCase { + func testFullWindowMatchesFeatureExtractor() { + // 10 s at 16 kHz: the extractor yields 999 frames and the encoder 250 tokens. + let frames = EmbeddingGemma2Audio.validFrames(samples: 160_000) + XCTAssertEqual(frames, 999) + XCTAssertEqual(EmbeddingGemma2Audio.tokenCount(validFrames: frames), 250) + } + + func testShortClipMatchesFeatureExtractor() { + // 55,680 samples (3.48 s): 347 frames, 87 audio tokens. + let frames = EmbeddingGemma2Audio.validFrames(samples: 55_680) + XCTAssertEqual(frames, 347) + XCTAssertEqual(EmbeddingGemma2Audio.tokenCount(validFrames: frames), 87) + } + + func testTooShortHasNoFrames() { + XCTAssertEqual(EmbeddingGemma2Audio.validFrames(samples: 100), 0) + } +}