From 1b0e5445dbd18911a1a62cc2f0ebcea5ecd325fa Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Thu, 16 Jul 2026 13:16:24 -0400 Subject: [PATCH] =?UTF-8?q?fix(offline-diarizer):=20pyannote-parity=20clus?= =?UTF-8?q?tering=20=E2=80=94=20threshold=20semantics,=20constraint=20coun?= =?UTF-8?q?t,=20constrained=20assignment=20(#801)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three porting bugs against pyannote community-1, all surfaced by #801 where auto clustering and numSpeakers=2 returned the same speaker count (2) but materially different partitions (69% vs 92% correctly-attributed speech): 1. AHC threshold: the config value was documented as a Euclidean cut distance (pyannote applies it directly via fcluster(..., "distance")), but the implementation reinterpreted it as a cosine similarity and cut at sqrt(2 - 2t). This inverted the knob's direction (raising the value split more instead of merging more) and shifted the default 0.6 to an effective cut of 0.894. The value is now applied directly, clamped to [0, 2]; config validation and CLI help updated to match. Old configs map to new ones via sqrt(2 - 2*old): old 0.6 ≈ new 0.894, old 0.7 ≈ new 0.775. 2. Speaker-count constraints compared against the AHC warm-start cluster count instead of the number of clusters VBx actually kept (pi > 1e-7, pyannote's auto_num_clusters). Since the user-visible detected count is the post-collapse survivor count, numSpeakers=N triggered K-Means re-clustering even when VBx had already collapsed to exactly N, producing a different partition at the same count. The check now uses VBxOutput.activeClusterCount. 3. Embedding-to-centroid assignment was a plain per-embedding argmax, which lets two local speakers sharing a segmentation chunk snap to the same centroid — silently absorbing one speaker's turns into another's (the mechanism behind the bad 82s-boundary partition in #801). Ported pyannote's constrained_argmax: per chunk, local speakers are matched to distinct clusters via Hungarian assignment maximizing total similarity. Controlled by Clustering.constrainedAssignment (default true) and automatically disabled when the count is forced via K-Means, matching pyannote. The Hungarian solver moved from DiarizationDER.swift into an internal HungarianAssignment enum with a rectangular max-score wrapper. Docs: Benchmarks.md threshold guidance rewritten for the new semantics (the old --threshold 0.7 advice corresponds to 0.775 today); min/max/ numSpeakers doc comments now state that the floor alone cannot correct an under-partition and that numSpeakers is a target, not a guarantee. --- Documentation/Benchmarks.md | 2 +- .../FluidAudio/Diarizer/DiarizationDER.swift | 63 +----------- .../Diarizer/HungarianAssignment.swift | 98 +++++++++++++++++++ .../Offline/Clustering/AHCClustering.swift | 22 +++-- .../ConstrainedClusterAssignment.swift | 43 ++++++++ .../Offline/Clustering/VBxClustering.swift | 9 +- .../Offline/Core/OfflineDiarizerManager.swift | 55 ++++++++--- .../Offline/Core/OfflineDiarizerTypes.swift | 45 ++++++++- .../Commands/DiarizationBenchmark.swift | 2 +- .../Commands/ProcessCommand.swift | 2 +- .../Diarizer/HungarianAssignmentTests.swift | 75 ++++++++++++++ .../Diarizer/Offline/AHCClusteringTests.swift | 47 +++++++-- .../ConstrainedClusterAssignmentTests.swift | 77 +++++++++++++++ .../Diarizer/Offline/OfflineConfigTests.swift | 1 + .../Diarizer/Offline/OfflineModuleTests.swift | 2 +- .../Diarizer/Offline/VBxConstraintTests.swift | 40 ++++++++ 16 files changed, 482 insertions(+), 101 deletions(-) create mode 100644 Sources/FluidAudio/Diarizer/HungarianAssignment.swift create mode 100644 Sources/FluidAudio/Diarizer/Offline/Clustering/ConstrainedClusterAssignment.swift create mode 100644 Tests/FluidAudioTests/Diarizer/HungarianAssignmentTests.swift create mode 100644 Tests/FluidAudioTests/Diarizer/Offline/ConstrainedClusterAssignmentTests.swift diff --git a/Documentation/Benchmarks.md b/Documentation/Benchmarks.md index 74be2866d..07463064e 100644 --- a/Documentation/Benchmarks.md +++ b/Documentation/Benchmarks.md @@ -634,7 +634,7 @@ AVERAGE 10.6 17.4 5.4 2.0 3.3 - 323.2 12/16 meetings detect the correct speaker count. Average DER 10.62% matches published pyannote-community-1 offline numbers on this split (~11-12%). Results are fully deterministic (two consecutive runs produce bit-identical metrics). -**Clustering threshold matters on this split.** The table above uses `--threshold 0.7`. Since #616 the CLI default is the community-1 preset (0.6), which merges clusters more aggressively and undercounts speakers on 4 meetings (EN2002a 2/4 → 41.9% DER, ES2004d 3/4 → 34.8%, EN2002b 19.0%, IS1009d 16.4%), degrading the average to 15.5% DER. Pass `--threshold 0.7` (or set `clusteringThreshold: 0.7` in `OfflineDiarizerConfig`) for AMI-SDM-like meeting audio. +**Clustering threshold matters on this split.** The table above predates #801 and uses `--threshold 0.7` under the old (inverted) threshold semantics. Since #801 the threshold is a Euclidean cut distance applied directly to the AHC dendrogram (pyannote parity): **larger values merge more aggressively and yield fewer speakers**. Old values map to new ones via `sqrt(2 − 2·old)` — the old default 0.6 behaved like a cut at 0.894, and the old `--threshold 0.7` from this table behaves like `--threshold 0.775` today. The CLI default remains the community-1 preset value (0.6), now interpreted as pyannote does. ### Streaming/online Diarization diff --git a/Sources/FluidAudio/Diarizer/DiarizationDER.swift b/Sources/FluidAudio/Diarizer/DiarizationDER.swift index 89c471081..f5879c485 100644 --- a/Sources/FluidAudio/Diarizer/DiarizationDER.swift +++ b/Sources/FluidAudio/Diarizer/DiarizationDER.swift @@ -121,7 +121,7 @@ public enum DiarizationDER { cost[h * n + r] = maxO - overlap[h * R + r] } } - let assign = Hungarian.solve(costSquare: cost, n: n) + let assign = HungarianAssignment.solve(costSquare: cost, n: n) for h in 0.. 0 { @@ -229,64 +229,3 @@ public enum DiarizationDER { return mask } } - -// MARK: - Hungarian (O(n^3) min-cost assignment on a square cost matrix) - -fileprivate enum Hungarian { - /// Kuhn-Munkres with potentials. `cost` is row-major `n × n`, - /// non-negative integers. Returns `assign[row] = col`. - static func solve(costSquare cost: [Int], n: Int) -> [Int] { - if n == 0 { return [] } - // Classic Jonker-Volgenant style implementation adapted for - // square n×n. 1-based indexing in arrays of size n+1 for the - // canonical algorithm — cost is read as `cost[(i-1)*n + (j-1)]`. - let INF = Int.max / 4 - var u = [Int](repeating: 0, count: n + 1) - var v = [Int](repeating: 0, count: n + 1) - var p = [Int](repeating: 0, count: n + 1) - var way = [Int](repeating: 0, count: n + 1) - - for i in 1...n { - p[0] = i - var j0 = 0 - var minv = [Int](repeating: INF, count: n + 1) - var used = [Bool](repeating: false, count: n + 1) - repeat { - used[j0] = true - let i0 = p[j0] - var delta = INF - var j1 = 0 - for j in 1...n where !used[j] { - let cur = cost[(i0 - 1) * n + (j - 1)] - u[i0] - v[j] - if cur < minv[j] { - minv[j] = cur - way[j] = j0 - } - if minv[j] < delta { - delta = minv[j] - j1 = j - } - } - for j in 0...n { - if used[j] { - u[p[j]] += delta - v[j] -= delta - } else { - minv[j] -= delta - } - } - j0 = j1 - } while p[j0] != 0 - repeat { - let j1 = way[j0] - p[j0] = p[j1] - j0 = j1 - } while j0 != 0 - } - var assign = [Int](repeating: -1, count: n) - for j in 1...n { - if p[j] != 0 { assign[p[j] - 1] = j - 1 } - } - return assign - } -} diff --git a/Sources/FluidAudio/Diarizer/HungarianAssignment.swift b/Sources/FluidAudio/Diarizer/HungarianAssignment.swift new file mode 100644 index 000000000..2f89128c3 --- /dev/null +++ b/Sources/FluidAudio/Diarizer/HungarianAssignment.swift @@ -0,0 +1,98 @@ +import Foundation + +// MARK: - Hungarian (O(n^3) min-cost assignment on a square cost matrix) + +enum HungarianAssignment { + /// Kuhn-Munkres with potentials. `cost` is row-major `n × n`, + /// non-negative integers. Returns `assign[row] = col`. + static func solve(costSquare cost: [Int], n: Int) -> [Int] { + if n == 0 { return [] } + // Classic Jonker-Volgenant style implementation adapted for + // square n×n. 1-based indexing in arrays of size n+1 for the + // canonical algorithm — cost is read as `cost[(i-1)*n + (j-1)]`. + let INF = Int.max / 4 + var u = [Int](repeating: 0, count: n + 1) + var v = [Int](repeating: 0, count: n + 1) + var p = [Int](repeating: 0, count: n + 1) + var way = [Int](repeating: 0, count: n + 1) + + for i in 1...n { + p[0] = i + var j0 = 0 + var minv = [Int](repeating: INF, count: n + 1) + var used = [Bool](repeating: false, count: n + 1) + repeat { + used[j0] = true + let i0 = p[j0] + var delta = INF + var j1 = 0 + for j in 1...n where !used[j] { + let cur = cost[(i0 - 1) * n + (j - 1)] - u[i0] - v[j] + if cur < minv[j] { + minv[j] = cur + way[j] = j0 + } + if minv[j] < delta { + delta = minv[j] + j1 = j + } + } + for j in 0...n { + if used[j] { + u[p[j]] += delta + v[j] -= delta + } else { + minv[j] -= delta + } + } + j0 = j1 + } while p[j0] != 0 + repeat { + let j1 = way[j0] + p[j0] = p[j1] + j0 = j1 + } while j0 != 0 + } + var assign = [Int](repeating: -1, count: n) + for j in 1...n { + if p[j] != 0 { assign[p[j] - 1] = j - 1 } + } + return assign + } + + /// Max-total-score assignment on a rectangular score matrix + /// (`scores[row][col]`, higher is better). Returns `assign[row] = col`, + /// or `-1` for rows left unassigned when there are more rows than columns. + /// Non-finite scores are treated as worse than any finite score. + static func maxScoreAssignment(scores: [[Double]]) -> [Int] { + let rows = scores.count + guard rows > 0 else { return [] } + let cols = scores[0].count + guard cols > 0 else { return Array(repeating: -1, count: rows) } + + let finite = scores.flatMap { $0 }.filter { $0.isFinite } + let maxScore = finite.max() ?? 0 + let minScore = finite.min() ?? 0 + let sentinel = minScore - 1 + + // Pad to square; dummy cells share a constant cost so they cannot + // influence which real pairs the solver prefers. Maximise score ⇒ + // minimise (maxScore − score), scaled to integers for the solver. + let n = max(rows, cols) + let scale = 1e6 + var cost = [Int](repeating: 0, count: n * n) + for r in 0.. Double { - guard !similarity.isNaN else { return Double.infinity } - if similarity < -1.0 || similarity > 1.0 { - logger.debug("Clustering threshold \(similarity) outside cosine range; clamping to [-1, 1]") + // MARK: - Distance Threshold Validation + // The threshold is a Euclidean cut distance on unit-normalized embeddings, + // applied directly to the dendrogram like pyannote's + // `fcluster(dendrogram, threshold, criterion="distance")`. Valid range is + // [0, 2]; larger values merge more aggressively (fewer clusters). + private func clampDistanceThreshold(_ threshold: Double) -> Double { + guard !threshold.isNaN else { + logger.warning("Clustering threshold is NaN; disabling AHC merging") + return 0 } - let clamped = max(-1.0, min(1.0, similarity)) - return sqrt(max(0, 2.0 - 2.0 * clamped)) + if threshold < 0 || threshold > 2.0 { + logger.warning("Clustering threshold \(threshold) outside [0, 2]; clamping") + } + return max(0, min(2.0, threshold)) } // MARK: - Dendrogram Parsing & Threshold-Based Cluster Assignment diff --git a/Sources/FluidAudio/Diarizer/Offline/Clustering/ConstrainedClusterAssignment.swift b/Sources/FluidAudio/Diarizer/Offline/Clustering/ConstrainedClusterAssignment.swift new file mode 100644 index 000000000..021ccea71 --- /dev/null +++ b/Sources/FluidAudio/Diarizer/Offline/Clustering/ConstrainedClusterAssignment.swift @@ -0,0 +1,43 @@ +import Foundation + +/// pyannote-parity constrained assignment (`constrained_argmax`): within a +/// segmentation chunk, distinct local speakers must map to distinct clusters. +/// Plain per-embedding argmax lets two speakers who share a chunk both snap to +/// the same centroid, silently absorbing one speaker's turns into the other's. +@available(macOS 14.0, iOS 17.0, *) +enum ConstrainedClusterAssignment { + + /// Assigns each embedding to a cluster, maximizing total similarity per + /// chunk under the one-cluster-per-local-speaker constraint. + /// + /// - Parameters: + /// - scores: Per-embedding similarity against each cluster centroid + /// (higher is better; non-finite values rank below all finite ones). + /// - chunkIndices: Chunk index of each embedding (parallel to `scores`). + /// - Returns: Cluster index per embedding, or `-2` when a chunk holds more + /// local speakers than there are clusters and the slot is dropped + /// (matching pyannote, which leaves such slots unassigned). + static func assign(scores: [[Double]], chunkIndices: [Int]) -> [Int] { + precondition( + scores.count == chunkIndices.count, + "scores and chunkIndices must be parallel arrays" + ) + var assignments = [Int](repeating: -2, count: scores.count) + + var rowsByChunk: [Int: [Int]] = [:] + for (row, chunk) in chunkIndices.enumerated() { + rowsByChunk[chunk, default: []].append(row) + } + + for rows in rowsByChunk.values { + let chunkScores = rows.map { scores[$0] } + let assigned = HungarianAssignment.maxScoreAssignment(scores: chunkScores) + for (localIndex, row) in rows.enumerated() { + let cluster = assigned[localIndex] + assignments[row] = cluster >= 0 ? cluster : -2 + } + } + + return assignments + } +} diff --git a/Sources/FluidAudio/Diarizer/Offline/Clustering/VBxClustering.swift b/Sources/FluidAudio/Diarizer/Offline/Clustering/VBxClustering.swift index 1ec518f5f..456988048 100644 --- a/Sources/FluidAudio/Diarizer/Offline/Clustering/VBxClustering.swift +++ b/Sources/FluidAudio/Diarizer/Offline/Clustering/VBxClustering.swift @@ -688,13 +688,18 @@ struct VBxClustering { initialClusters: [Int], constraints: SpeakerCountConstraints? ) -> VBxOutput { - var output = refine(rhoFeatures: rhoFeatures, initialClusters: initialClusters) + let output = refine(rhoFeatures: rhoFeatures, initialClusters: initialClusters) guard let constraints = constraints else { return output } - let detectedCount = output.numClusters + // Compare against the clusters VBx actually kept (pi above epsilon), not the + // AHC warm-start count: VBx collapses unused warm-start clusters, and that + // collapsed count is what the pipeline reports as the detected speaker count. + // Checking the warm-start count re-clusters even when VBx already agrees + // with the requested count, replacing its partition with K-Means (#801). + let detectedCount = output.activeClusterCount guard constraints.needsAdjustment(detectedCount: detectedCount) else { return output } diff --git a/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerManager.swift b/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerManager.swift index 184b5119c..17f04669b 100644 --- a/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerManager.swift +++ b/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerManager.swift @@ -351,10 +351,28 @@ public final class OfflineDiarizerManager { if centroids.isEmpty { centroids = computeFallbackCentroids(from: embeddingFeatures) } - let assignments = assignEmbeddings( - embeddingFeatures: embeddingFeatures, - centroids: centroids - ) + // pyannote parity: constrain co-chunk speakers to distinct clusters, but + // not when the count was forced via K-Means — the constraint can then + // artificially inflate the number of speakers. + let useConstrainedAssignment = + config.clustering.constrainedAssignment + && !vbxOutput.wasAdjusted + && centroids.count > 1 + let assignments: [Int] + if useConstrainedAssignment { + assignments = ConstrainedClusterAssignment.assign( + scores: centroidScores( + embeddingFeatures: embeddingFeatures, + centroids: centroids + ), + chunkIndices: timedEmbeddings.map(\.chunkIndex) + ) + } else { + assignments = assignEmbeddings( + embeddingFeatures: embeddingFeatures, + centroids: centroids + ) + } let chunkAssignments = buildChunkAssignments( segmentation: segmentation, @@ -767,6 +785,18 @@ public final class OfflineDiarizerManager { return [accumulator] } + /// Cosine similarity of every embedding against every centroid. + private func centroidScores( + embeddingFeatures: [[Double]], + centroids: [[Double]] + ) -> [[Double]] { + let normalizedCentroids = centroids.map(normalize) + return embeddingFeatures.map { embedding in + let normalizedEmbedding = normalize(embedding) + return normalizedCentroids.map { dot(normalizedEmbedding, $0) } + } + } + private func assignEmbeddings( embeddingFeatures: [[Double]], centroids: [[Double]] @@ -776,17 +806,16 @@ public final class OfflineDiarizerManager { return Array(repeating: 0, count: embeddingFeatures.count) } - let normalizedCentroids = centroids.map(normalize) - return embeddingFeatures.map { embedding in - let normalizedEmbedding = normalize(embedding) + let scores = centroidScores( + embeddingFeatures: embeddingFeatures, + centroids: centroids + ) + return scores.map { row in var bestIndex = 0 var bestScore = -Double.infinity - for (index, centroid) in normalizedCentroids.enumerated() { - let score = dot(normalizedEmbedding, centroid) - if score > bestScore { - bestScore = score - bestIndex = index - } + for (index, score) in row.enumerated() where score > bestScore { + bestScore = score + bestIndex = index } return bestIndex } diff --git a/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerTypes.swift b/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerTypes.swift index f2c7c58c3..3d608a33b 100644 --- a/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerTypes.swift +++ b/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerTypes.swift @@ -118,7 +118,11 @@ public struct OfflineDiarizerConfig: Sendable { } public struct Clustering: Sendable { - /// Euclidean distance threshold for unit-normalized embeddings. + /// Euclidean distance threshold for unit-normalized embeddings, in [0, 2]. + /// + /// Applied directly as the AHC dendrogram cut distance (pyannote parity): + /// embeddings closer than this merge into one cluster, so **larger values + /// merge more aggressively and yield fewer speakers**. public var threshold: Double /// VBx warm-start parameters (Fa controls precision, Fb controls recall) @@ -126,21 +130,36 @@ public struct OfflineDiarizerConfig: Sendable { public var warmStartFb: Double /// Minimum number of speakers. Ignored if `numSpeakers` is set. + /// + /// Only binds when the auto-detected count falls below it; it cannot + /// change a partition whose count already satisfies the bound. public var minSpeakers: Int? /// Maximum number of speakers. Ignored if `numSpeakers` is set. public var maxSpeakers: Int? /// Exact number of speakers. Overrides `minSpeakers` and `maxSpeakers` when set. + /// + /// Treated as a target, not a guarantee: when it differs from the + /// auto-detected count, embeddings are re-clustered with K-Means to this + /// count, but downstream assignment may leave some clusters unused. public var numSpeakers: Int? + /// When true (pyannote parity), local speakers sharing a segmentation + /// chunk are assigned to distinct clusters instead of each snapping to + /// its nearest centroid independently. Automatically disabled when the + /// speaker count is forced via `numSpeakers`/`minSpeakers`/`maxSpeakers` + /// re-clustering. + public var constrainedAssignment: Bool + public static let community = Clustering( threshold: 0.6, warmStartFa: 0.07, warmStartFb: 0.8, minSpeakers: nil, maxSpeakers: nil, - numSpeakers: nil + numSpeakers: nil, + constrainedAssignment: true ) public init( @@ -149,7 +168,8 @@ public struct OfflineDiarizerConfig: Sendable { warmStartFb: Double, minSpeakers: Int? = nil, maxSpeakers: Int? = nil, - numSpeakers: Int? = nil + numSpeakers: Int? = nil, + constrainedAssignment: Bool = true ) { self.threshold = threshold self.warmStartFa = warmStartFa @@ -157,6 +177,7 @@ public struct OfflineDiarizerConfig: Sendable { self.minSpeakers = minSpeakers self.maxSpeakers = maxSpeakers self.numSpeakers = numSpeakers + self.constrainedAssignment = constrainedAssignment } } @@ -334,10 +355,11 @@ public struct OfflineDiarizerConfig: Sendable { /// Validate configuration values and throw if they fall outside expected ranges. public func validate() throws { - let maxClusteringThreshold = sqrt(2.0) + // Euclidean distances between unit-normalized embeddings live in [0, 2]. + let maxClusteringThreshold = 2.0 guard clustering.threshold > 0, clustering.threshold <= maxClusteringThreshold else { throw OfflineDiarizationError.invalidConfiguration( - "clustering.threshold must be within (0, sqrt(2)], got \(clustering.threshold)" + "clustering.threshold must be within (0, 2], got \(clustering.threshold)" ) } @@ -637,6 +659,19 @@ public struct VBxOutput: Sendable { self.wasAdjusted = wasAdjusted self.originalClusterCount = originalClusterCount } + + /// Mixture-weight epsilon below which a VBx cluster counts as collapsed. + public static let activeClusterEpsilon = 1e-7 + + /// Number of clusters VBx actually kept (mixture weight above epsilon). + /// + /// VBx is warm-started with the AHC cluster count (`numClusters`) and prunes + /// clusters by driving their mixture weight to zero, so this — not + /// `numClusters` — is the auto-detected speaker count (pyannote parity). + public var activeClusterCount: Int { + guard !pi.isEmpty else { return numClusters } + return pi.filter { $0 > Self.activeClusterEpsilon }.count + } } /// Intermediate representation of an embedding associated with its timeline. diff --git a/Sources/FluidAudioCLI/Commands/DiarizationBenchmark.swift b/Sources/FluidAudioCLI/Commands/DiarizationBenchmark.swift index 1257c9400..8e070b4cd 100644 --- a/Sources/FluidAudioCLI/Commands/DiarizationBenchmark.swift +++ b/Sources/FluidAudioCLI/Commands/DiarizationBenchmark.swift @@ -115,7 +115,7 @@ enum StreamDiarizationBenchmark { --min-embed-update Min segment duration to update embeddings (default: 2.0) Offline Mode Options: - --threshold <0-√2> Clustering threshold (default: 0.6) + --threshold <0-2> AHC cut distance, higher = fewer speakers (default: 0.6) --fa VBx warm-start precision (default: 0.07) --fb VBx warm-start recall (default: 0.8) --window-duration Segmentation window size (default: 10.0) diff --git a/Sources/FluidAudioCLI/Commands/ProcessCommand.swift b/Sources/FluidAudioCLI/Commands/ProcessCommand.swift index 434cafa25..34ee98bb1 100644 --- a/Sources/FluidAudioCLI/Commands/ProcessCommand.swift +++ b/Sources/FluidAudioCLI/Commands/ProcessCommand.swift @@ -488,7 +488,7 @@ enum ProcessCommand { OFFLINE MODE OPTIONS (VBx clustering, all optional): --rttm Compute DER/JER against RTTM annotations - --threshold <0-√2> Euclidean clustering threshold (default: 0.6) + --threshold <0-2> AHC cut distance, higher = fewer speakers (default: 0.6) --fa VBx warm-start precision (default: 0.07) --fb VBx warm-start recall (default: 0.8) --window-duration Segmentation window size (default: 10.0) diff --git a/Tests/FluidAudioTests/Diarizer/HungarianAssignmentTests.swift b/Tests/FluidAudioTests/Diarizer/HungarianAssignmentTests.swift new file mode 100644 index 000000000..87a2d186c --- /dev/null +++ b/Tests/FluidAudioTests/Diarizer/HungarianAssignmentTests.swift @@ -0,0 +1,75 @@ +import XCTest + +@testable import FluidAudio + +final class HungarianAssignmentTests: XCTestCase { + + // MARK: - Square min-cost solve + + func testSolveEmptyMatrix() { + XCTAssertTrue(HungarianAssignment.solve(costSquare: [], n: 0).isEmpty) + } + + func testSolvePicksMinimumCostPerfectMatching() { + // Row 0 is cheapest on col 1, row 1 on col 0, row 2 on col 2. + let cost = [ + 4, 1, 3, + 2, 0, 5, + 3, 2, 2, + ] + let assign = HungarianAssignment.solve(costSquare: cost, n: 3) + XCTAssertEqual(assign, [1, 0, 2]) + } + + func testSolveResolvesGreedyConflict() { + // Both rows prefer col 0; the optimal total forces row 1 onto col 0. + let cost = [ + 1, 2, + 0, 10, + ] + let assign = HungarianAssignment.solve(costSquare: cost, n: 2) + XCTAssertEqual(assign, [1, 0]) + } + + // MARK: - Rectangular max-score assignment + + func testMaxScoreAssignmentSquare() { + let scores: [[Double]] = [ + [0.9, 0.1], + [0.8, 0.2], + ] + // Greedy would put both rows on col 0; the constraint forces row 1 + // (the weaker match) onto col 1. + XCTAssertEqual(HungarianAssignment.maxScoreAssignment(scores: scores), [0, 1]) + } + + func testMaxScoreAssignmentMoreColumnsThanRows() { + let scores: [[Double]] = [ + [0.1, 0.9, 0.3] + ] + XCTAssertEqual(HungarianAssignment.maxScoreAssignment(scores: scores), [1]) + } + + func testMaxScoreAssignmentMoreRowsThanColumnsDropsWeakestRow() { + let scores: [[Double]] = [ + [0.9], + [0.5], + [0.7], + ] + // Only one column exists; the strongest row keeps it, the rest drop. + XCTAssertEqual(HungarianAssignment.maxScoreAssignment(scores: scores), [0, -1, -1]) + } + + func testMaxScoreAssignmentTreatsNonFiniteAsWorst() { + let scores: [[Double]] = [ + [Double.nan, 0.2], + [0.6, 0.5], + ] + XCTAssertEqual(HungarianAssignment.maxScoreAssignment(scores: scores), [1, 0]) + } + + func testMaxScoreAssignmentEmpty() { + XCTAssertTrue(HungarianAssignment.maxScoreAssignment(scores: []).isEmpty) + XCTAssertEqual(HungarianAssignment.maxScoreAssignment(scores: [[], []]), [-1, -1]) + } +} diff --git a/Tests/FluidAudioTests/Diarizer/Offline/AHCClusteringTests.swift b/Tests/FluidAudioTests/Diarizer/Offline/AHCClusteringTests.swift index dd85ff7cd..cda0aaca6 100644 --- a/Tests/FluidAudioTests/Diarizer/Offline/AHCClusteringTests.swift +++ b/Tests/FluidAudioTests/Diarizer/Offline/AHCClusteringTests.swift @@ -36,8 +36,10 @@ final class AHCClusteringTests: XCTestCase { // MARK: - Orthogonal Embeddings - func testOrthogonalEmbeddingsSeparateAtHighThreshold() { - // Two groups: pointing in very different directions + func testOrthogonalGroupsSeparateWhenThresholdBelowGroupDistance() { + // Two groups: pointing in very different directions. Unit-normalized, + // the inter-group Euclidean distance is ~sqrt(2) ≈ 1.414, so a cut at + // 0.8 merges within groups but not across them. let group1: [[Double]] = [ [1.0, 0.0, 0.0], [0.9, 0.1, 0.0], @@ -65,6 +67,24 @@ final class AHCClusteringTests: XCTestCase { ) } + func testLargerThresholdMergesMore() { + // pyannote semantics: the threshold is the dendrogram cut distance, + // so raising it can only reduce the number of clusters. + let embeddings: [[Double]] = [ + [1.0, 0.0, 0.0], + [0.9, 0.1, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.9, 0.1], + ] + + let tight = Set(clustering.cluster(embeddingFeatures: embeddings, threshold: 0.5)).count + let loose = Set(clustering.cluster(embeddingFeatures: embeddings, threshold: 1.5)).count + + XCTAssertEqual(tight, 2, "Cut below the inter-group distance should keep two clusters") + XCTAssertEqual(loose, 1, "Cut above the inter-group distance should merge everything") + XCTAssertLessThanOrEqual(loose, tight) + } + // MARK: - Cluster ID Properties func testClusterIdsAreContiguousFromZero() { @@ -84,19 +104,32 @@ final class AHCClusteringTests: XCTestCase { } } - // MARK: - Low Threshold + // MARK: - Threshold Extremes + + func testMaxThresholdMergesAll() { + // 2.0 is the largest possible distance between unit vectors, so + // everything should merge into one cluster. + let embeddings: [[Double]] = [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ] + + let result = clustering.cluster(embeddingFeatures: embeddings, threshold: 2.0) + let uniqueClusters = Set(result) + XCTAssertEqual(uniqueClusters.count, 1, "Maximum threshold should merge all into one cluster") + } - func testLowThresholdMergesAll() { - // At a very low threshold (cosine similarity), everything should merge into one cluster + func testZeroThresholdKeepsDistinctEmbeddingsApart() { let embeddings: [[Double]] = [ [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], ] - let result = clustering.cluster(embeddingFeatures: embeddings, threshold: -1.0) + let result = clustering.cluster(embeddingFeatures: embeddings, threshold: 0.0) let uniqueClusters = Set(result) - XCTAssertEqual(uniqueClusters.count, 1, "Very low threshold should merge all into one cluster") + XCTAssertEqual(uniqueClusters.count, 3, "Zero threshold should not merge distinct embeddings") } // MARK: - Zero Vector diff --git a/Tests/FluidAudioTests/Diarizer/Offline/ConstrainedClusterAssignmentTests.swift b/Tests/FluidAudioTests/Diarizer/Offline/ConstrainedClusterAssignmentTests.swift new file mode 100644 index 000000000..ee8c0f197 --- /dev/null +++ b/Tests/FluidAudioTests/Diarizer/Offline/ConstrainedClusterAssignmentTests.swift @@ -0,0 +1,77 @@ +import XCTest + +@testable import FluidAudio + +@available(macOS 14.0, iOS 17.0, *) +final class ConstrainedClusterAssignmentTests: XCTestCase { + + func testCoChunkSpeakersGetDistinctClusters() { + // Both embeddings score highest against cluster 0. Plain argmax would + // merge them; the per-chunk constraint forces the weaker match onto + // cluster 1 (the #801 failure mode). + let scores: [[Double]] = [ + [0.9, 0.3], + [0.8, 0.6], + ] + let assignments = ConstrainedClusterAssignment.assign( + scores: scores, + chunkIndices: [0, 0] + ) + XCTAssertEqual(assignments, [0, 1]) + } + + func testSeparateChunksAreUnconstrained() { + let scores: [[Double]] = [ + [0.9, 0.3], + [0.8, 0.6], + ] + let assignments = ConstrainedClusterAssignment.assign( + scores: scores, + chunkIndices: [0, 1] + ) + XCTAssertEqual(assignments, [0, 0], "Embeddings in different chunks may share a cluster") + } + + func testMoreSpeakersThanClustersDropsWeakestSlot() { + let scores: [[Double]] = [ + [0.9], + [0.2], + ] + let assignments = ConstrainedClusterAssignment.assign( + scores: scores, + chunkIndices: [0, 0] + ) + XCTAssertEqual(assignments, [0, -2], "Surplus co-chunk speaker should be left unassigned") + } + + func testSingleEmbeddingPerChunkIsPlainArgmax() { + let scores: [[Double]] = [ + [0.1, 0.7, 0.4], + [0.5, 0.2, 0.9], + ] + let assignments = ConstrainedClusterAssignment.assign( + scores: scores, + chunkIndices: [3, 7] + ) + XCTAssertEqual(assignments, [1, 2]) + } + + func testEmptyInput() { + XCTAssertTrue(ConstrainedClusterAssignment.assign(scores: [], chunkIndices: []).isEmpty) + } + + func testMaximizesTotalScoreAcrossChunk() { + // Row 0 slightly prefers cluster 1, but giving cluster 1 to row 1 + // yields a higher total; the matching should be globally optimal, + // not first-come-first-served. + let scores: [[Double]] = [ + [0.50, 0.55], + [0.10, 0.90], + ] + let assignments = ConstrainedClusterAssignment.assign( + scores: scores, + chunkIndices: [0, 0] + ) + XCTAssertEqual(assignments, [0, 1]) + } +} diff --git a/Tests/FluidAudioTests/Diarizer/Offline/OfflineConfigTests.swift b/Tests/FluidAudioTests/Diarizer/Offline/OfflineConfigTests.swift index f79d141a2..4c68fc97e 100644 --- a/Tests/FluidAudioTests/Diarizer/Offline/OfflineConfigTests.swift +++ b/Tests/FluidAudioTests/Diarizer/Offline/OfflineConfigTests.swift @@ -9,6 +9,7 @@ final class OfflineConfigTests: XCTestCase { XCTAssertNil(clustering.minSpeakers) XCTAssertNil(clustering.maxSpeakers) XCTAssertNil(clustering.numSpeakers) + XCTAssertTrue(clustering.constrainedAssignment) } func testClusteringAcceptsSpeakerConstraints() { diff --git a/Tests/FluidAudioTests/Diarizer/Offline/OfflineModuleTests.swift b/Tests/FluidAudioTests/Diarizer/Offline/OfflineModuleTests.swift index 03fe11eae..8e69a342f 100644 --- a/Tests/FluidAudioTests/Diarizer/Offline/OfflineModuleTests.swift +++ b/Tests/FluidAudioTests/Diarizer/Offline/OfflineModuleTests.swift @@ -21,7 +21,7 @@ final class OfflineDiarizerConfigTests: XCTestCase { } func testValidateThrowsForInvalidClusteringThreshold() { - let config = OfflineDiarizerConfig(clusteringThreshold: 1.5) + let config = OfflineDiarizerConfig(clusteringThreshold: 2.5) XCTAssertThrowsError(try config.validate()) { error in guard case OfflineDiarizationError.invalidConfiguration(let message) = error else { diff --git a/Tests/FluidAudioTests/Diarizer/Offline/VBxConstraintTests.swift b/Tests/FluidAudioTests/Diarizer/Offline/VBxConstraintTests.swift index 1f71ae65d..89472d635 100644 --- a/Tests/FluidAudioTests/Diarizer/Offline/VBxConstraintTests.swift +++ b/Tests/FluidAudioTests/Diarizer/Offline/VBxConstraintTests.swift @@ -47,4 +47,44 @@ final class VBxConstraintTests: XCTestCase { XCTAssertEqual(output.numClusters, 2) XCTAssertEqual(output.originalClusterCount, 8) } + + // MARK: - Active cluster count (pyannote auto_num_clusters parity) + + func testActiveClusterCountIgnoresCollapsedClusters() { + // VBx warm-started with 5 AHC clusters but collapsed 3 of them + // (mixture weight ~0). The detected speaker count is 2, not 5. + let output = VBxOutput( + gamma: [], + pi: [0.63, 0.0, 1e-12, 0.37, 0.0], + hardClusters: [], + centroids: [], + numClusters: 5, + elbos: [] + ) + XCTAssertEqual(output.activeClusterCount, 2) + } + + func testActiveClusterCountWithAllClustersActive() { + let output = VBxOutput( + gamma: [], + pi: [0.5, 0.3, 0.2], + hardClusters: [], + centroids: [], + numClusters: 3, + elbos: [] + ) + XCTAssertEqual(output.activeClusterCount, 3) + } + + func testActiveClusterCountFallsBackToNumClustersWithoutPi() { + let output = VBxOutput( + gamma: [], + pi: [], + hardClusters: [], + centroids: [], + numClusters: 4, + elbos: [] + ) + XCTAssertEqual(output.activeClusterCount, 4) + } }