Skip to content

Expose token timings from the Unified offline batch path#814

Merged
Alex-Wengg merged 2 commits into
FluidInference:mainfrom
NylonDiamond:unified-offline-token-timings
Jul 23, 2026
Merged

Expose token timings from the Unified offline batch path#814
Alex-Wengg merged 2 commits into
FluidInference:mainfrom
NylonDiamond:unified-offline-token-timings

Conversation

@NylonDiamond

Copy link
Copy Markdown
Contributor

Problem

UnifiedAsrManager.transcribe(_:) returns only a String, 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:

merged.sort { $0.timestamp < $1.timestamp }
merged = merger.collapseSeamWordDuplicates(merged, vocabulary: tokenizer.vocabulary)
return tokenizer.decode(ids: merged.map(\.token))   // timestamps dropped here

ChunkProcessor.TokenWindow.timestamp is the global encoder frame the greedy RNNT decoder records per emission (transcribeWindow passes globalFrameOffset: 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 into buildWordTimings(from:).

The conversion mirrors StreamingUnifiedAsrManager rather than inventing anything: frame x frameSamples / sampleRate (80 ms), normalized to a leading space, and each token's endTime clamped back only when it would overrun its successor, so genuine pauses survive as gaps between spans.

Both batch entry points now go through one decodedTokens helper, so the text from transcribe and transcribeWithTimings cannot drift apart. transcribe(_:)'s behavior and signature are unchanged.

Tests

The conversion is factored out as a pure function, so UnifiedTokenTimingTests covers 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 through buildWordTimings. 7 tests, and swift 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.

`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.
@Alex-Wengg

Copy link
Copy Markdown
Member

notes from code bot

Blocking:

  1. Drop the NOTICE header (lines 1–8). It's a downstream-fork provenance notice — once merged, upstream's own file would claim to be "modified from the original FluidAudio source" and point at this PR as pending. No other file in the repo carries a per-file header; git history covers provenance.
  2. Doc comment on transcribeWithTimings is inaccurate. It says "each token's endTime is the next token's start" — but the implementation (correctly) gives every token a provisional one-frame end and only clamps back on overrun, so pauses leave gaps. Your own silenceSurvivesAsAGapBetweenTokens test asserts a 0.4 s hole. The PR description words it right ("clamped back only when it would overrun its successor") — please use that phrasing in the doc so nobody builds contiguous spans off it.

Optional / non-blocking:

  • The conversion is a near-verbatim copy of the inline loop in StreamingUnifiedAsrManager.processWindow (same clamp, same ▁ replacement, same comments). Fine for this PR, but a follow-up that has streaming call a shared append-with-back-fill helper would prevent the two paths drifting — the exact drift decodedTokens exists to prevent for text.
  • Since the batch path has samples.count, the final token's provisional end could be clamped to the clip duration — streaming can't know the total, but offline can, and it avoids endTime past EOF for seek use cases.
  • The tokenizer guard now appears three times in one call chain (transcribe, transcribeWithTimings, decodedTokens); passing the unwrapped tokenizer into decodedTokens would collapse it to one per entry point.
  • If the method pair sticks around, the AVAudioPCMBuffer overload probably wants a WithTimings twin eventually — one-liner workaround exists via AudioConverter, so no rush.

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
NylonDiamond force-pushed the unified-offline-token-timings branch from 79b6ec1 to 4a26809 Compare July 22, 2026 22:48
@NylonDiamond

Copy link
Copy Markdown
Contributor Author

Thanks, both blocking items are fixed, plus two of the optional ones. Force-pushed as 4a26809.

  1. NOTICE header dropped, along with the commit that added it. Fair point that the file would end up claiming to be a modified copy of itself once merged.
  2. Doc comment now uses the PR description's phrasing: every token gets a provisional one-frame end that is clamped back only when it would overrun its successor, so the spans are explicitly not contiguous and a pause stays visible as a gap.
  3. Took the clip-duration suggestion. tokenTimings grew a clipDuration: Double? (nil leaves it unbounded, which is what streaming would pass) and the batch path passes samples.count / sampleRate, so the last token's end can no longer sit past EOF. Two tests cover the clamp and the no-op case where the clip runs longer.
  4. Collapsed the tokenizer guard: decodedTokens now takes the unwrapped tokenizer, so notInitialized is checked once per entry point instead of three times in one chain.

Left the other two for follow-ups. Having streaming call a shared append-with-back-fill helper is the right fix but it touches StreamingUnifiedAsrManager beyond this PR's scope, and the AVAudioPCMBuffer twin is easy to add whenever the method pair settles. Happy to do either in this PR if you would rather they land together.

swift build clean, swift format lint --strict clean, and the timing tests are at 9 passing.

@Alex-Wengg
Alex-Wengg merged commit f852921 into FluidInference:main Jul 23, 2026
11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants