From 0f89a626f872a7e3d60b307c3e01e77258f24a65 Mon Sep 17 00:00:00 2001 From: Christian Brickhouse Date: Tue, 30 Jan 2024 13:28:39 -0800 Subject: [PATCH] [hotfix] Fix speaker assignment bug Speaker assignment treated the result as a list when it was actually a dictionary. --- pipeline.py | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/pipeline.py b/pipeline.py index fcd3c77..0af857d 100644 --- a/pipeline.py +++ b/pipeline.py @@ -94,17 +94,21 @@ def assign_speakers( result_segments = assign_word_speakers( diarization_result, aligned_segments ) - results_segments_w_speakers: List[Dict[str, Any]] = [] - for result_segment in result_segments: - results_segments_w_speakers.append( - { - "start": result_segment["start"], - "end": result_segment["end"], - "text": result_segment["text"], - "speaker": result_segment["speaker"], - } - ) - return results_segments_w_speakers + # Upstream uses this, but it's bugged and I think upstream's upstream has since adopted the + # output that it tries to create making it redundant + # + #results_segments_w_speakers: List[Dict[str, Any]] = [] + #for result_segment in result_segments['segments']: + # results_segments_w_speakers.append( + # { + # "start": result_segment["start"], + # "end": result_segment["end"], + # "text": result_segment["text"], + # "speaker": result_segment["speaker"], + # "words": result_segment["words"] + # } + # ) + return result_segments def transcribe_and_diarize( audio_file: str, @@ -133,7 +137,7 @@ def transcribe_and_diarize( results_segments_w_speakers = assign_speakers(diarization_result, aligned_segments) # Print the results in a user-friendly way - for i, segment in enumerate(results_segments_w_speakers): + for i, segment in enumerate(results_segments_w_speakers['segments']): print(f"Segment {i + 1}:") print(f"Start time: {segment['start']:.2f}") print(f"End time: {segment['end']:.2f}")