diff --git a/Documentation/Benchmarks.md b/Documentation/Benchmarks.md index 74be2866..07463064 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 89c47108..f5879c48 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 00000000..2f89128c --- /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 00000000..021ccea7 --- /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 1ec518f5..45698804 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 184b5119..17f04669 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 f2c7c58c..3d608a33 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 1257c940..8e070b4c 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 434cafa2..34ee98bb 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 00000000..87a2d186 --- /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 dd85ff7c..cda0aaca 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 00000000..ee8c0f19 --- /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 f79d141a..4c68fc97 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 03fe11ea..8e69a342 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 1f71ae65..89472d63 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) + } }