Skip to content

Coordinator: the feed and warm loops are a RecordingEngineSession #80

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dinooo13 wants to merge 1 commit into
base: core/finish-pipeline
Choose a base branch
from
Open

Coordinator: the feed and warm loops are a RecordingEngineSession #80

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 21 additions & 94 deletions Sources/PladderCore/DictationCoordinator.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,25 +99,12 @@ public final class DictationCoordinator {
private var pendingUnloads: [any TranscriptionEngine] = []
private var cycleEngine: (any TranscriptionEngine)?

// Cancelling cannot abort a CoreML call already started, so a release landing in
// a warm pass waits for it; the log shows that as `engine` above `engine-time`.
private var warmupTask: Task<Void, Never>?
private var stream: Stream?

private struct Stream {
let engine: any StreamingTranscriptionEngine
let utterance: Utterance
let feed: Task<Int, Never>
}
private var engineSession: RecordingEngineSession?

private var recordingID = 0
private var cancelCleanup: Task<Void, Never>?
private var muteRestoreTask: Task<Void, Never>?

// Every utterance is padded to the model's full window, so half a second of
// silence costs the same encoder pass as a real call.
private static let warmupSamples = [Float](repeating: 0, count: 8_000)

@ObservationIgnored private(set) var handledHotkeyEvents = 0

public struct CycleTiming: Sendable, Equatable {
Expand Down Expand Up @@ -387,8 +374,16 @@ public final class DictationCoordinator {
onEvent(.recordingStarted)
armOutputMute()
startKeyDownWork()
await startEngineWork(live: live, recordingID: mine)
guard isCurrent(mine) else { return }
if let engine = cycleEngine {
let session = RecordingEngineSession(
engine: engine, capture: capture, clock: clock,
intervals: .init(feed: feedInterval, livePass: livePassInterval, warmup: warmupInterval),
live: live, onPartial: { [weak self] in self?.partialTranscript = $0 })
engineSession = session
await session.begin()
guard isCurrent(mine) else { return }
session.startLoops()
}
levelTask = Task { [weak self] in
for await level in levels {
guard let self, self.state.isRecording else { return }
Expand Down Expand Up @@ -416,62 +411,6 @@ public final class DictationCoordinator {
state.isRecording && recordingID == id
}

// A streaming engine that cannot begin an utterance is used like a batch engine:
// transcribed whole at release, slower but not lost.
private func startEngineWork(live: Bool, recordingID id: Int) async {
if let streaming = cycleEngine as? (any StreamingTranscriptionEngine),
let utterance = try? await streaming.beginUtterance() {
guard isCurrent(id) else {
await streaming.abandonUtterance(utterance)
return
}
startStreamingFeed(streaming, utterance: utterance, live: live)
// The live loop's pass is the warm pass; a second loop would only compete with it.
if !live { startWarmupLoop { await streaming.warmPass() } }
} else if let engine = cycleEngine {
startWarmupLoop { [warmupSamples = Self.warmupSamples] in
_ = try? await engine.transcribe(warmupSamples)
}
}
}

// One loop for feed and live passes, so they never drain against each other. At
// release the loop is cancelled and awaited: a chunk drained just before is still
// fed, so every sample reaches the engine once and in order.
private func startStreamingFeed(_ engine: any StreamingTranscriptionEngine, utterance: Utterance, live: Bool) {
let feed = Task { [weak self, clock, feedInterval, livePassInterval] in
var fed = 0
while !Task.isCancelled {
if !live { try? await clock.sleep(for: feedInterval) }
guard let self, !Task.isCancelled, self.state.isRecording else { return fed }
let chunk = await self.capture.drain()
if !chunk.isEmpty {
fed += chunk.count
await engine.feed(chunk, to: utterance)
}
guard live, !Task.isCancelled else { continue }
let text = await engine.livePass(utterance)
// A pass finishing after the release would put stale text back on screen.
guard !Task.isCancelled, self.state.isRecording else { return fed }
self.partialTranscript = text
try? await clock.sleep(for: livePassInterval)
}
return fed
}
stream = Stream(engine: engine, utterance: utterance, feed: feed)
}

private func startWarmupLoop(_ warm: @escaping @Sendable () async -> Void) {
warmupTask?.cancel()
warmupTask = Task.detached(priority: .utility) { [clock, interval = warmupInterval] in
while !Task.isCancelled {
await warm()
if Task.isCancelled { return }
try? await clock.sleep(for: interval)
}
}
}

private func armOutputMute() {
guard settings.muteOutputWhileDictating, let outputMuter else { return }
let session = recordingID
Expand All @@ -493,7 +432,7 @@ public final class DictationCoordinator {
// Synchronous, and every call only flips a flag, cancels a task or spawns one,
// so on the release path this costs nothing before `recordingStopped`.
@discardableResult
private func endRecording() -> Stream? {
private func endRecording() -> RecordingEngineSession.Ending? {
settleTask?.cancel()
settleTask = nil
gesture.reset()
Expand All @@ -504,20 +443,17 @@ public final class DictationCoordinator {
inputLevel = 0
maxDurationTask?.cancel()
maxDurationTask = nil
let ended = stream
ended?.feed.cancel()
stream = nil
// Before the state flip, so no further pass is queued ahead of the real call.
warmupTask?.cancel()
warmupTask = nil
let ending = engineSession?.end()
engineSession = nil
partialTranscript = nil
restoreOutputDevice()
return ended
return ending
}

public func hotkeyReleased(submit: Bool = false) {
guard state.isRecording else { return }
let stream = endRecording()
let ending = endRecording()
state = .transcribing
// Snapshots anything copied while speaking now, beside the engine pass, not in
// the paste.
Expand All @@ -534,19 +470,14 @@ public final class DictationCoordinator {
// The feed loop finishes handing over a chunk drained before the release. What
// it waits for is engine work `endUtterance` would wait for anyway: engine time.
let handover = ContinuousClock.now
let fed = await stream?.feed.value
let fed = await ending?.feed?.value
let stopped = ContinuousClock.now
let feedWait = stopped - handover
let audio = await self.capture.stop()
let captureStop = ContinuousClock.now - stopped
let source: DictationFinish.Source = if let stream {
.utterance(stream.engine, stream.utterance)
} else {
.whole(engine)
}
let outcome = await finish.run(
DictationFinish.Input(
audio: audio, fedSamples: fed ?? 0, source: source, submit: submit,
audio: audio, fedSamples: fed ?? 0, source: ending?.source ?? .whole(engine), submit: submit,
captureStop: captureStop, feedWait: feedWait),
onStage: { self.enter($0) })
self.finished(outcome)
Expand Down Expand Up @@ -594,19 +525,15 @@ public final class DictationCoordinator {
// Synchronous, so the state is idle on return; the task stops the microphone.
private func beginCancel() -> Task<Void, Never>? {
guard state.isRecording else { return nil }
let stream = endRecording()
let ending = endRecording()
willPolish = false
becomeIdle()
cycleEngine = nil
// The mic goes off first; the engine drops the utterance only once the feed has
// exited, so no chunk lands after it. Then it waits for the cancel before it.
// The mic goes off first, then the utterance; then it waits for the cancel before it.
let previous = cancelCleanup
let cleanup = Task { [capture] in
_ = await capture.stop()
if let stream {
_ = await stream.feed.value
await stream.engine.abandonUtterance(stream.utterance)
}
await ending?.abandon()
await previous?.value
}
cancelCleanup = cleanup
Expand Down
134 changes: 134 additions & 0 deletions Sources/PladderCore/RecordingEngineSession.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
import Foundation

// See docs/ARCHITECTURE.md, "Dictation flow".
@MainActor
final class RecordingEngineSession {
struct Intervals: Sendable {
var feed: Duration
var livePass: Duration
var warmup: Duration
}

struct Ending: Sendable {
let source: DictationFinish.Source
// Awaited at release: a chunk drained just before is still fed, so every sample
// reaches the engine once and in order.
let feed: Task<Int, Never>?

// The engine drops the utterance only once the feed has exited, so no chunk lands
// after it.
func abandon() async {
guard case .utterance(let engine, let utterance) = source else { return }
_ = await feed?.value
await engine.abandonUtterance(utterance)
}
}

// Every utterance is padded to the model's full window, so half a second of
// silence costs the same encoder pass as a real call.
static let warmupSamples = [Float](repeating: 0, count: 8_000)

private let engine: any TranscriptionEngine
private let capture: any AudioCapture
private let clock: any Clock<Duration>
private let intervals: Intervals
private let live: Bool
private let onPartial: @MainActor (String?) -> Void
private var streaming: (engine: any StreamingTranscriptionEngine, utterance: Utterance)?
private var feed: Task<Int, Never>?
// Cancelling cannot abort a CoreML call already started, so a release landing in
// a warm pass waits for it; the log shows that as `engine` above `engine-time`.
private var warmup: Task<Void, Never>?
private(set) var isEnded = false

init(
engine: any TranscriptionEngine,
capture: any AudioCapture,
clock: any Clock<Duration>,
intervals: Intervals,
live: Bool,
onPartial: @escaping @MainActor (String?) -> Void
) {
self.engine = engine
self.capture = capture
self.clock = clock
self.intervals = intervals
self.live = live
self.onPartial = onPartial
}

// A streaming engine that cannot begin an utterance is used like a batch engine:
// transcribed whole at release, slower but not lost.
func begin() async {
guard let engine = engine as? (any StreamingTranscriptionEngine),
let utterance = try? await engine.beginUtterance()
else { return }
guard !isEnded else {
await engine.abandonUtterance(utterance)
return
}
streaming = (engine, utterance)
}

func startLoops() {
guard !isEnded else { return }
if let streaming {
startFeed(streaming.engine, utterance: streaming.utterance)
// The live loop's pass is the warm pass; a second loop would only compete with it.
if !live { startWarmupLoop { [engine = streaming.engine] in await engine.warmPass() } }
} else {
startWarmupLoop { [engine, warmupSamples = Self.warmupSamples] in
_ = try? await engine.transcribe(warmupSamples)
}
}
}

// Synchronous: it only cancels, so the release path waits on nothing here.
@discardableResult
func end() -> Ending {
isEnded = true
feed?.cancel()
warmup?.cancel()
warmup = nil
let source: DictationFinish.Source = if let streaming {
.utterance(streaming.engine, streaming.utterance)
} else {
.whole(engine)
}
return Ending(source: source, feed: feed)
}

// One loop for feed and live passes, so they never drain against each other.
private func startFeed(_ engine: any StreamingTranscriptionEngine, utterance: Utterance) {
feed = Task { [weak self, capture, clock, intervals, live] in
var fed = 0
while !Task.isCancelled {
if !live { try? await clock.sleep(for: intervals.feed) }
guard let self, !Task.isCancelled, !self.isEnded else { return fed }
let chunk = await capture.drain()
if !chunk.isEmpty {
fed += chunk.count
await engine.feed(chunk, to: utterance)
}
guard live, !Task.isCancelled else { continue }
let text = await engine.livePass(utterance)
// A pass finishing after the release would put stale text back on screen.
guard !Task.isCancelled, !self.isEnded else { return fed }
self.onPartial(text)
try? await clock.sleep(for: intervals.livePass)
}
return fed
}
}

private func startWarmupLoop(_ warm: @escaping @Sendable () async -> Void) {
warmup?.cancel()
warmup = Task.detached(priority: .utility) { [clock, interval = intervals.warmup] in
while !Task.isCancelled {
await warm()
if Task.isCancelled { return }
try? await clock.sleep(for: interval)
}
}
}
}
Loading
Loading