UPDATE fixing multiple speakers detections

app.py

@@ -51,6 +51,7 @@ initial_prompt = os.environ.get(
 
 beam_size_default = int(os.environ.get("WHISPER_BEAM_SIZE", 5))
 best_of_default = int(os.environ.get("WHISPER_BEST_OF", 5))
+expected_speakers_default = int(os.environ.get("EXPECTED_SPEAKERS", 2))
 
 # ---------------------------------------------------------------------------
 # Lazy singletons for the heavy models
@@ -147,6 +148,7 @@ def transcribe(
     audio_path: str,
     language: str,
     enable_diarization: bool,
+    expected_speakers: int,
     beam_size: int,
     best_of: int,
 ) -> Dict:
@@ -305,6 +307,57 @@ def transcribe(
         # Sort speaker turns by start time
         speaker_turns.sort(key=lambda x: x["start"])
         
+        # Consolidate speakers if we detected more than expected
+        unique_speakers = set(turn["speaker"] for turn in speaker_turns)
+        print(f"[DEBUG] Detected {len(unique_speakers)} unique speakers: {unique_speakers}")
+        
+        if expected_speakers > 0 and len(unique_speakers) > expected_speakers:
+            print(f"[DEBUG] Consolidating from {len(unique_speakers)} to {expected_speakers} speakers")
+            # Create a mapping to merge speakers
+            # Strategy: Merge speakers by order of first appearance, keeping the most active ones
+            speaker_stats = {}
+            for turn in speaker_turns:
+                spk = turn["speaker"]
+                if spk not in speaker_stats:
+                    speaker_stats[spk] = {"first_appear": turn["start"], "duration": 0, "count": 0}
+                speaker_stats[spk]["duration"] += turn["end"] - turn["start"]
+                speaker_stats[spk]["count"] += 1
+            
+            # Sort speakers by total speaking duration (most active first)
+            sorted_speakers = sorted(speaker_stats.items(), key=lambda x: x[1]["duration"], reverse=True)
+            print(f"[DEBUG] Speaker activity: {[(s, round(stats['duration'], 1)) for s, stats in sorted_speakers]}")
+            
+            # Keep the top N most active speakers, map others to them
+            kept_speakers = [s[0] for s in sorted_speakers[:expected_speakers]]
+            speaker_mapping = {}
+            
+            for spk, stats in sorted_speakers:
+                if spk in kept_speakers:
+                    speaker_mapping[spk] = spk
+                else:
+                    # Map this speaker to the closest kept speaker by first appearance time
+                    closest_kept = min(kept_speakers, 
+                                     key=lambda k: abs(speaker_stats[k]["first_appear"] - stats["first_appear"]))
+                    speaker_mapping[spk] = closest_kept
+                    print(f"[DEBUG] Mapping {spk} -> {closest_kept}")
+            
+            # Apply the mapping
+            for turn in speaker_turns:
+                turn["speaker"] = speaker_mapping[turn["speaker"]]
+            
+            # Merge consecutive turns from the same speaker
+            merged_turns = []
+            for turn in speaker_turns:
+                if merged_turns and merged_turns[-1]["speaker"] == turn["speaker"] and \
+                   turn["start"] - merged_turns[-1]["end"] < 1.0:  # Less than 1 second gap
+                    # Extend the previous turn
+                    merged_turns[-1]["end"] = turn["end"]
+                else:
+                    merged_turns.append(turn.copy())
+            
+            speaker_turns = merged_turns
+            print(f"[DEBUG] After consolidation: {len(speaker_turns)} speaker turns")
+        
         # Assign speakers to each transcript segment
         for segment in response["segments"]:
             seg_start = segment["start"]
@@ -406,6 +459,14 @@ def build_interface() -> gr.Blocks:
                 value=False,
                 info="Uses NVIDIA Sortformer model (max 4 speakers, downloads ~700MB on first use).",
             )
+            expected_speakers_slider = gr.Slider(
+                label="Expected Number of Speakers",
+                minimum=0,
+                maximum=4,
+                step=1,
+                value=expected_speakers_default,
+                info="Set to 0 for automatic detection, or specify 2-4 to consolidate speakers.",
+            )
             beam_slider = gr.Slider(
                 label="Beam Size",
                 minimum=1,
@@ -432,6 +493,7 @@ def build_interface() -> gr.Blocks:
                 audio_input,
                 language_input,
                 diarization_toggle,
+                expected_speakers_slider,
                 beam_slider,
                 best_of_slider,
             ],