fix(offline-diarizer): pyannote-parity clustering — threshold semantics, constraint count, constrained assignment

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
 

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..<H {
                 let r = assign[h]
                 if r < R && overlap[h * R + r] > 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
-    }
-}

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..<rows {
+            precondition(scores[r].count == cols, "Jagged score matrix is not supported")
+            for c in 0..<cols {
+                let score = scores[r][c].isFinite ? scores[r][c] : sentinel
+                cost[r * n + c] = Int(((maxScore - score) * scale).rounded())
+            }
+        }
+
+        let assign = solve(costSquare: cost, n: n)
+        return (0..<rows).map { row in
+            let col = assign[row]
+            return col < cols ? col : -1
+        }
+    }
+}

Sources/FluidAudio/Diarizer/Offline/Clustering/AHCClustering.swift

@@ -54,7 +54,7 @@ struct AHCClustering {
             return Array(0..<count)
         }
 
-        let distanceThreshold = convertThresholdToDistance(threshold)
+        let distanceThreshold = clampDistanceThreshold(threshold)
         let assignments = assignmentsFromDendrogram(
             dendrogram,
             count: count,
@@ -104,14 +104,20 @@ struct AHCClustering {
         return normalized
     }
 
-    // MARK: - Similarity-to-Distance Conversion
-    private func convertThresholdToDistance(_ similarity: Double) -> 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

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
+    }
+}

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
         }

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
         }