Repository navigation
Conversation
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).
Alex-Wengg
left a comment
There was a problem hiding this comment.
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( |
There was a problem hiding this comment.
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())) |
There was a problem hiding this comment.
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 { |
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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( |
There was a problem hiding this comment.
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( |
There was a problem hiding this comment.
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( |
There was a problem hiding this comment.
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() { |
There was a problem hiding this comment.
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).
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 readspeakerDatabase. 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.embedSpanalready 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 asspeakerDatabasevalues.embedSpandoes 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.embedding.batchSizewindows.embedSpannow uses the same batched path. Its output is unchanged: on EN2002a withzeroVoteReembedon, the segments and speaker database are byte-identical tomain.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-speakerprocess(audio:)on each turn (stepRatio 0.5, PLDA batch 8, 10 VBx iterations). New is oneembedSpanscall per speaker.Speed, for 1,683 turns:
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:
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:
Speaker means still match every speaker's old
speakerDatabasecentroid 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
embedSpansembeds 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.
process(audio:)per turn)embedSpans)The before and after numbers come from different meetings of similar length, so read the per-turn cost.