Expose token timings from the Unified offline batch path#814
Merged
Alex-Wengg merged 2 commits intoJul 23, 2026
Merged
Conversation
`UnifiedAsrManager.transcribe` already decodes per-token encoder frames — the greedy RNNT decoder records one per emission and the overlap merge preserves them — and then discards them on its final line, so callers that need to align a batch transcript back to its audio (seeking, playback highlighting, word/speaker attribution) have no way to get them. Add `transcribeWithTimings(_:)`, the batch counterpart to `StreamingUnifiedAsrManager.consumeTokenTimings()`. It returns the same `[TokenTiming]` shape, so its output feeds `buildWordTimings(from:)` identically, and it uses the same conversion the streaming manager does: frame x frameSamples/sampleRate (80 ms), with each token's end back-filled only when it would overrun its successor, so real pauses survive as gaps. Both batch entry points now share one `decodedTokens` helper, so the text they return cannot drift apart. The conversion is factored out as a pure function and tested without loading the model.
Member
|
notes from code bot Blocking:
Optional / non-blocking:
|
Drop the per-file NOTICE header: no other file in the repo carries one and git history already covers provenance, so upstream shouldn't end up with a file claiming to be a modified copy of itself. Reword the doc comment to match what the code does. Ends are provisional and clamped back only on overrun, so spans are not contiguous and a pause stays visible as a gap; the old wording invited callers to build contiguous spans off it. Clamp the last token's provisional end to the clip duration. The batch path knows the sample count up front, so a seek target never lands past EOF; streaming keeps the unbounded behavior via the new optional parameter. Pass the unwrapped tokenizer into decodedTokens so the notInitialized guard appears once per entry point rather than three times in one call chain.
NylonDiamond
force-pushed
the
unified-offline-token-timings
branch
from
July 22, 2026 22:48
79b6ec1 to
4a26809
Compare
Contributor
Author
|
Thanks, both blocking items are fixed, plus two of the optional ones. Force-pushed as 4a26809.
Left the other two for follow-ups. Having streaming call a shared append-with-back-fill helper is the right fix but it touches
|
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
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.
Problem
UnifiedAsrManager.transcribe(_:)returns only aString, so there is no way to align a batch transcript back to its audio. The information is already there and is thrown away one line before the return:ChunkProcessor.TokenWindow.timestampis the global encoder frame the greedy RNNT decoder records per emission (transcribeWindowpassesglobalFrameOffset: chunkStart / config.frameSamples), and the overlap merge preserves it.The streaming sibling already publishes exactly this, via
StreamingUnifiedAsrManager.consumeTokenTimings()—buildWordTimings(from:)even cites it in its doc comment. So the offline path is the odd one out. Applications that want Unified's WER and punctuation for batch transcription currently have to fall back to a TDT checkpoint purely to keep seek and playback highlighting, which means shipping a second English model.Change
Adds
UnifiedAsrManager.transcribeWithTimings(_:) -> TranscriptionWithTimings, returning the text plus[TokenTiming], so it drops straight intobuildWordTimings(from:).The conversion mirrors
StreamingUnifiedAsrManagerrather than inventing anything:frame x frameSamples / sampleRate(80 ms),▁normalized to a leading space, and each token'sendTimeclamped back only when it would overrun its successor, so genuine pauses survive as gaps between spans.Both batch entry points now go through one
decodedTokenshelper, so the text fromtranscribeandtranscribeWithTimingscannot drift apart.transcribe(_:)'s behavior and signature are unchanged.Tests
The conversion is factored out as a pure function, so
UnifiedTokenTimingTestscovers it without loading the model: frame→seconds, pause preservation, no span overrunning its successor, tokens sharing a frame, the frontier token's provisional one-frame end, unknown ids, and end-to-end regrouping throughbuildWordTimings. 7 tests, andswift build/ the existing suite are unaffected.Caveat
These are emission times, not forced-alignment boundaries: RNNT emits once the decoder has heard enough context, so a token's start can sit slightly after the word's true onset. That is already true of
consumeTokenTimings()and is documented on the new method.