Update app.py

app.py

@@ -47,42 +47,63 @@ pipe = pipeline(
 )
 
 
-
-
 def associate_speakers_with_timestamps(transcription_result, diarization):
     word_segments = transcription_result['chunks']
     diarization_segments = list(diarization.itertracks(yield_label=True))
+    
+    # Pré-calcul des chevauchements
+    overlaps = set()
+    for i in range(1, len(diarization_segments)):
+        if diarization_segments[i-1][0].end > diarization_segments[i][0].start:
+            overlaps.add(diarization_segments[i][0].start)
+
     speaker_transcription = []
-    current_segment_index = 0
-    previous_speaker = None
+    current_speaker = None
+    current_text = []
+    
+    diar_iter = iter(diarization_segments)
+    current_segment, _, current_speaker = next(diar_iter)
 
     for word in word_segments:
-        word_start, word_end = word['timestamp']
-
-        # Trouver le segment de diarisation correspondant
-        while current_segment_index < len(diarization_segments) - 1 and diarization_segments[current_segment_index][0].end <= word_start:
-            current_segment_index += 1
-
-        current_segment, _, current_speaker = diarization_segments[current_segment_index]
-
-        # Vérifier s'il y a un chevauchement avec le segment précédent
-        if current_segment_index > 0:
-            previous_segment, _, previous_speaker = diarization_segments[current_segment_index - 1]
-            if previous_segment.end > word_start and previous_speaker != current_speaker:
-                word_text = f"[{word['text']}]"
-            else:
-                word_text = word['text']
+        word_start, _ = word['timestamp']
+        
+        while word_start >= current_segment.end:
+            if current_text:
+                speaker_transcription.append((current_speaker, ' '.join(current_text)))
+                current_text = []
+            current_segment, _, current_speaker = next(diar_iter, (None, None, None))
+            if current_segment is None:
+                break
+
+        if current_segment is None:
+            break
+
+        if word_start in overlaps:
+            current_text.append(f"[{word['text']}]")
         else:
-            word_text = word['text']
+            current_text.append(word['text'])
 
-        # Ajouter le mot au segment du locuteur actuel
-        if not speaker_transcription or speaker_transcription[-1][0] != current_speaker:
-            speaker_transcription.append((current_speaker, word_text))
-        else:
-            speaker_transcription[-1] = (current_speaker, speaker_transcription[-1][1] + " " + word_text)
+    if current_text:
+        speaker_transcription.append((current_speaker, ' '.join(current_text)))
 
     return speaker_transcription
 
+def merge_short_segments(speaker_transcription, min_words=3):
+    def merge_group(group):
+        speaker, texts = group
+        merged = []
+        current = []
+        for text in texts:
+            current.extend(text.split())
+            if len(current) >= min_words:
+                merged.append(' '.join(current))
+                current = []
+        if current:
+            merged.append(' '.join(current))
+        return [(speaker, text) for text in merged]
+
+    return [item for group in groupby(speaker_transcription, key=lambda x: x[0]) for item in merge_group(group)]
+
 def merge_short_segments(speaker_transcription, min_words=3):
     merged = []
     for speaker, text in speaker_transcription: