Skip to content

feat(diarizer): embed known spans with OfflineDiarizerManager.embedSpans - #1015

Open
Toasty360 wants to merge 2 commits into
FluidInference:mainfrom
Toasty360:offline-embed-spans
Open

Toasty360 wants to merge 2 commits into
FluidInference:mainfrom
Toasty360:offline-embed-spans

Conversation

@Toasty360

@Toasty360 Toasty360 commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor

Problem

An app that already has speaker turns (from another diarizer such as Nemotron 3, or a cloud transcript) and needs a voice embedding per turn has only one public way to get it: run OfflineDiarizerManager.process(audio:) on each turn's clip, forced to one speaker, and read speakerDatabase. That runs segmentation, FBANK, embedding, PLDA and VBx clustering for every turn. In our app that's 127 runs and 16 of 20 seconds on an iPhone 13 for a 43-minute, 8-speaker meeting.

OfflineEmbeddingExtractor.embedSpan already embeds one span directly (for the zero-vote re-embed), but it's internal.

Change

  • OfflineDiarizerManager.embedSpans(audioSource:spans:) async throws -> [[Float]]: one raw (not normalized) 256-dim WeSpeaker embedding per time span, the same kind of vector as speakerDatabase values.
  • Each span is embedded the way embedSpan does it: repeated to fill the 10 s window, every frame weighted. A span longer than a window is cut into equal pieces of at most one window, and their embeddings are averaged.
  • FBANK and embedding predictions run in batches of embedding.batchSize windows.
  • embedSpan now uses the same batched path. Its output is unchanged: on EN2002a with zeroVoteReembed on, the segments and speaker database are byte-identical to main.
  • Unit tests for how spans are cut into pieces, and a short section in Documentation/Diarization/GettingStarted.md.

Measurements

Apple M6, macOS 27, .cpuAndNeuralEngine, AMI test set (16 meetings, single distant microphone). The turns are built from the AMI reference words the way a transcript groups them into paragraphs. For each reference speaker we take the 30 longest turns of 2 s or more. Old is a forced-1-speaker process(audio:) on each turn (stepRatio 0.5, PLDA batch 8, 10 VBx iterations). New is one embedSpans call per speaker.

Speed, for 1,683 turns:

Old New
Total 65.4 s 7.0 s
Per turn ~35 ms ~4 ms

That's about 9× faster. On a 43.6-minute, 4-speaker meeting (TS3003d), it goes from 3.9 s to 0.4 s.

Agreement. Cosine between old and new, on the 1,589 turns that the old path embeds:

p5 median min
Per turn 0.897 0.977 0.584
Turns 10–30 s 0.940 0.975 0.732
Mean of a speaker's turns 0.989 0.997 0.973

Spans over 10 s need the averaging. Truncating them to the first window drops long turns (over 30 s) from a median cosine of 0.990 to 0.904.

Separation. Turn-vs-turn cosine between the same speaker and different speakers:

Same speaker, mean Different speakers, mean AUC
Old 0.669 0.105 0.988
New 0.685 0.112 0.989

Speaker means still match every speaker's old speakerDatabase centroid at 0.95 on average, and never match another speaker above 0.54.

One difference: the old path returns nothing for about 6% of turns (2–4 s turns with too little detected speech), while embedSpans embeds every span it's given. Those turns' embeddings are noisy, so callers that relied on that implicit speech check should filter turns themselves (by word count, for example).

iPhone 13 (A15, 4 GB, iOS 27)

A real job in an app on AMI EN2002a (36 min): Parakeet transcription, then Nemotron diarization, then embeddings for its turns.

Before (one process(audio:) per turn) After (embedSpans)
Turns embedded 127 (on a 43-min meeting) 120
Time for those turns 16–16.7 s (~0.13 s each) 2.4 s (~0.02 s each)
Whole embedding step, including model load and one embedding per speaker ~20 s 5.7 s

The before and after numbers come from different meetings of similar length, so read the per-turn cost.

Apps that already have speaker turns (another diarizer, a transcript)
had to run the whole offline pipeline per turn to get one embedding:
segmentation, FBANK, embedding, PLDA and VBx clustering, forced to one
speaker. Add OfflineDiarizerManager.embedSpans(audioSource:spans:),
which embeds each span directly like the internal embedSpan (span
tiled to fill the window, every frame weighted) and returns raw
256-dim vectors, the same kind as speakerDatabase values. Spans longer
than a window are cut into equal pieces whose embeddings are averaged.
FBANK and embedding predictions run in batches of embedding.batchSize.

embedSpan (zero-vote re-embedding) now shares the batched path; its
output is unchanged (EN2002a with zeroVoteReembed on: identical
segments and speaker database).
@Toasty360
Toasty360 marked this pull request as draft October 12, 2026 22:58
@Toasty360
Toasty360 marked this pull request as ready for review October 13, 2026 02:14

@Alex-Wengg Alex-Wengg left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review of embedSpans. The first four are correctness issues on the new public path (shared-model concurrency, trap on non-finite spans, no config.validate(), blocking the caller's executor); the rest are cleanup and test gaps.

/// - audioSource: Mono audio at the model's sample rate.
/// - spans: Time spans in seconds; each must cover at least one sample of the audio.
/// - Returns: One embedding per span.
public func embedSpans(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

embedSpans runs sync MLModel.predictions on the shared fbank/embedding models with no serialization. Two overlapping calls (the doc's one-call-per-speaker pattern fanned out in a task group, or a call during process(audio:)) hit the same non-thread-safe CoreML path. prepare() avoids this by owning a single extractor task; this should serialize or the doc should say it must not run concurrently.

endSeconds: Double
) throws -> Range<Int> {
let sampleRate = Double(config.sampleRate)
let startSample = max(0, Int((startSeconds * sampleRate).rounded()))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Int(x.rounded()) traps on NaN, ±inf, or out-of-range doubles, and spans now come from a public API. A malformed transcript timestamp kills the process instead of throwing OfflineDiarizationError.

audioSource: AudioSampleSource,
spans: [Range<TimeInterval>]
) async throws -> [[Float]] {
if models == nil {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

prepare() calls config.validate(); this path does not. With windowDurationSeconds == 0 (or batchSize 0) spanPieces returns [] for every span and the call returns empty vectors with no error. Same root cause regresses embedSpan: a zero-length span used to throw, now truncated is empty and tiledWindow yields a zero window.

pldaTransform: PLDATransform(pldaRhoModel: models.pldaRhoModel, psi: models.pldaPsi),
config: config
)
return try extractor.embedSpans(audioSource: audioSource, spans: spans)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All CoreML work here runs synchronously inside the async function on the caller's executor, so a 36-minute meeting blocks a cooperative-pool thread for seconds and cancellation is only seen between batches. prepare() offloads via group.addTask(priority: .userInitiated); this should do the same or be documented as blocking.

}

/// Runs the embedding model once per FBANK feature array, with the same frame weights for all.
private func runEmbeddingBatch(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

runEmbeddingBatch duplicates the provider/output-extraction logic in runEmbeddingModel and drops its prefetchToNeuralEngine calls, so the zero-vote re-embed path loses the ANE prefetch. Suggest runEmbeddingModel wraps runEmbeddingBatch the way runFbankModel wraps runFbankBatch.

guard let models else {
throw OfflineDiarizationError.modelNotLoaded("offline-diarizer")
}
let extractor = OfflineEmbeddingExtractor(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Each call builds a fresh OfflineEmbeddingExtractor (model-description parsing, shape resolution, an ANEMemoryOptimizer) plus a PLDATransform that embedding never uses. Models are immutable after initialize, so the extractor can be cached on the manager, or the extractor could take an optional PLDA.

let spanLength = min(endSample - startSample, config.samplesPerWindow)
guard spanLength > 0 else {
guard endSample > startSample else {
throw OfflineDiarizationError.processingFailed(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A span that starts past the end of the audio is reported as empty span after endSample is clamped, so a bad timestamp from a longer transcript timeline reads like a degenerate span. Worth a distinct out-of-range error.

XCTAssertEqual(OfflineEmbeddingExtractor.spanPieces(100..<260, windowLength: 160), [100..<260])
}

func testLongSpanIsCutIntoEqualPiecesOfAtMostOneWindow() {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing the two boundary cases where the ceil arithmetic is most likely to go off by one: a span that is an exact multiple of windowLength (0..<20 with 10 must give exactly two pieces) and a span of exactly one window (one piece).

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