feat: streaming transcribe API for Parakeet ONNX engine#7
Open
fiorelorenzo wants to merge 5 commits into
Open
Conversation
Adds `Model::transcribe_stream()` returning a `Box<dyn StreamSession>` that lets callers push audio incrementally and receive partial transcripts. v1 implements re-encode-growing-buffer for the Parakeet ONNX engine: each push recomputes mel + runs the encoder on the full buffer-so-far, then resumes the greedy TDT decode from the previous time-step. Quadratic in audio length; acceptable for short utterances (typical dictation <30s on M-series M4 encodes in tens of ms per pass), to be optimized in a follow-up via chunked-attention encoder state. `Engine` trait grows a `transcribe_stream` method with a default `Err(NotSupported)` impl, so Whisper / qwen3-asr / parakeet-mlx engines compile unchanged. Streaming support for those is a follow-up. Added `tests/parakeet_streaming.rs` (ignored by default, needs cached weights) asserting the chunked transcript matches the one-shot transcript on the same audio. Consumer: the Lirevo dictation app (https://github.com/fiorelorenzo/lirevo) needs this for its live-overlay UX.
The v1 streaming API committed tokens off whatever encoder context was available at the time of each push, and never revisited those tokens when later pushes grew the encoder's input. Since Parakeet's Conformer encoder is non-causal, the encoder output for the same early frames is not stable across pushes, and on utterances longer than ~20s the cumulative streamed text could drift from the one-shot transcribe() result on the same audio. v1.1 always re-decodes from t=0 every push and computes the consumer- facing `delta` as the diff vs the previously-emitted text. The cumulative text in `PartialTranscript.text` is now byte-identical to what `transcribe()` would return on the same audio at every push. The UI should treat `text` as authoritative and use `delta` as a hint; early tokens may be corrected retroactively when later pushes give the encoder more context. Removed the per-session TDT state carry-across (`state_h`, `state_c`, `last_token`, `emitted`) - no longer needed since each push decodes from scratch. Updated the integration test to assert strict text equality between streamed cumulative text and one-shot text. Build, clippy, and fmt gates green.
Mirrors the Parakeet ONNX streaming impl from the previous commits on this branch onto the MLX engine. Same re-decode-growing-buffer architecture, same PartialTranscript contract, same byte-equivalence to one-shot `transcribe()` by construction. Refactors ParakeetMlxEngine to hold `conformer`, `predict`, `joint` behind `Arc<Mutex<_>>` so the streaming session can hold cloned handles, paralleling the ONNX engine's `Arc<Mutex<ort::Session>>` pattern. Vocab, durations, mel_config and time_ratio are cloned into the session by value. The mel + encoder + TDT decode pipeline is pulled out of the engine's `transcribe` into `run_full_pipeline` so the streaming session can reuse the same code path. `mlx_synchronize` is now called via a small `gpu_synchronize` helper at the end of `transcribe`, `push`, and `finish` to keep the GPU-buffer-lifetime contract intact. Send on `ParakeetMlxStreamSession` is asserted unsafe with the same single-thread-at-a-time invariant the engine documents: the trait surface holds the session behind `&mut self` and `Box<dyn StreamSession + Send>`. Extends `tests/parakeet_streaming.rs` with a parakeet-mlx variant gated on the `parakeet-mlx` feature. The file's outer cfg switched from `feature = "parakeet"` to `any(parakeet, parakeet-mlx)` so each engine's test can be enabled independently.
Author
|
Adds a third commit extending the streaming API to the Parakeet MLX engine using the same re-decode-growing-buffer architecture. Same PartialTranscript contract, same byte-equivalence to one-shot transcribe(). gated on the |
rc.10 pinned smallvec to a 2.0.0-alpha.10 prerelease; rc.12 uses smallvec ^1.15 which resolves a downstream resolver collision for projects that pull mistral.rs (whose transitive serde-saphyr requires smallvec ^2.0.0-alpha.12, a different prerelease that is mutually exclusive with rc.10's pin). rc.12 also moved to ndarray 0.17 (audiopipe was on 0.16) and made `ort::Error` generic on a recovery payload — added a `From<ort::Error<SessionBuilder>>` bridge in `error.rs` so `?` keeps working at session-builder call sites. No behaviour change in audiopipe's parakeet / whisper paths; build is clean against `parakeet,parakeet-mlx,qwen3-asr-ggml,whisper,metal,coreml`.
Author
|
Added a small follow-up commit on this branch bumping ort =2.0.0-rc.10 to =2.0.0-rc.12. Reason: rc.10 pinned a smallvec 2.0.0-alpha prerelease that collides with another consumer downstream. rc.12 uses a stable smallvec; rc.12 also moved to ndarray 0.17 and made ort::Error generic on a recovery payload, so the bump includes an ndarray dep bump and a small From<Error> bridge in error.rs. No API surface change in audiopipe's parakeet / whisper paths. Happy to split this into its own PR if you would prefer. |
Author
|
Added a follow-up commit (d4ef5d7 bump ort to 2.0.0-rc.12) bumping ort =2.0.0-rc.10 to =2.0.0-rc.12. Reason: rc.10 pinned a smallvec 2.0.0-alpha prerelease that collides with a downstream consumer (mistral.rs, whose transitive serde-saphyr requires smallvec ^2.0.0-alpha.12, a different incompatible prerelease). rc.12 uses stable smallvec ^1.15; no API surface change inside audiopipe's parakeet/whisper paths. Happy to split into its own PR if you'd prefer. |
Streaming re-decodes the whole (growing) audio buffer on every push, and MLX caches Metal buffers keyed by size. Dictations of increasing length keep minting fresh buffer sizes the cache never reuses, so the MLX cache and the process IOAccelerator footprint climb unbounded across a session (~11 GB after a dozen dictations observed in the Lirevo host). Add the MLX memory FFI (clear_cache / set_cache_limit / get_*_memory) and: - set a 512 MiB cache ceiling once at engine init (AUDIOPIPE_MLX_CACHE_LIMIT_MB to tune, =0 to disable); - clear_cache after finish() and batch transcribe() (AUDIOPIPE_NO_MLX_CACHE_CLEAR to opt out). Validated with examples/mlx_stream_leak_probe over increasing-length dictations: mlx_cache 58 -> 968 -> ... -> 2959 MB (unbounded) before, flat at 0 MB after; phys_footprint 6.0 GB -> 3.0 GB (just the model).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
Model::transcribe_stream()returning aBox<dyn StreamSession>. Consumers push audio incrementally and receive aPartialTranscript { text, delta, segments, is_final }per chunk. Cumulativetextis byte-identical toModel::transcribe()on the same audio by construction.Implementation is re-decode-growing-buffer: each
pushrecomputes the mel on the buffer-so-far, runs the encoder on the full buffer, and re-decodes the TDT fromt=0. The new tokens beyond the previously-emitted text are exposed asdelta; consumers should treattextas authoritative and usedeltaas a hint (a later push may shorten or correct earlier text once the encoder sees more context).Cost is quadratic in audio length. On Apple Silicon M-series this is fine for typical dictation utterances under 30 seconds (encoder pass is tens of milliseconds). A follow-up could add chunked-attention state to keep the per-push cost constant, at the price of an encoder-side ONNX export change.
Enginetrait grows atranscribe_streammethod with a defaultErr(Error::NotSupported), sowhisper.rs,qwen3_asr_*.rs, andparakeet_mlxcompile unchanged. Streaming support for those engines is a separate concern, not in this PR.Ownership:
ParakeetEnginenow holdsencoderanddecoderasArc<Mutex<ort::Session>>so the streaming session can carry handles without lifetime-through-trait gymnastics.ort::SessionisSend + Syncupstream; contention is a non-issue because callers go through&mut selfonModel.tests/parakeet_streaming.rsasserts strict equality between the chunked cumulative text and the one-shot text on a fixture WAV. The test is#[ignore]d because it needs cached Parakeet weights and a local fixture file (run withcargo test --features parakeet -- --ignored).Consumer: the Lirevo dictation app (https://github.com/fiorelorenzo/lirevo) wired its Cargo dep to this branch to drive a live-overlay UX during push-to-talk.
Two commits on the branch keep the history honest: the first lands the streaming session machinery with an incremental decode that could drift on long utterances; the second refactors the decode to re-run from scratch each push so the cumulative text matches one-shot exactly.