Generate Custom Speaker Labels with Pyannote

1
import os
2
import assemblyai as aai
3
from pyannote.audio import Pipeline
4
import torch
5
import pandas as pd
6
import numpy as np
7
8
# Assign your API keys
9
HUGGING_FACE_TOKEN = os.getenv("HF_TOKEN")
10
ASSEMBLYAI_API_KEY = os.getenv("ASSEMBLYAI_API_KEY")
11
12
# Authenticate with AssemblyAI
13
aai.settings.api_key = ASSEMBLYAI_API_KEY
14
15
def transcribe_audio(audio_file, language="en"):
16
    """
17
    Transcribe an audio file using AssemblyAI.
18
19
    Args:
20
        audio_file (str): Path to the audio file.
21
        language (str, optional): Language code for transcription. Defaults to "en".
22
23
    Returns:
24
        aai.Transcript: The transcription result.
25
    """
26
27
    transcriber = aai.Transcriber(config=aai.TranscriptionConfig(speech_model='nano', language_code=language))
28
    transcript = transcriber.transcribe(audio_file)
29
    print(f"Transcript ID: {transcript.id}")
30
    return transcript
31
32
def get_speaker_labels(audio_file, transcript: aai.Transcript):
33
    """
34
    Perform speaker diarization on an audio file and combine results with the transcript.
35
36
    Args:
37
        audio_file (str): Path to the audio file.
38
        transcript (aai.Transcript): The transcription result from AssemblyAI.
39
40
    Returns:
41
        str: A formatted string containing the transcript with speaker labels and timestamps.
42
    """
43
    # Initialize the speaker diarization pipeline with GPU support
44
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
45
    pipeline = Pipeline.from_pretrained(
46
        "pyannote/speaker-diarization",
47
        use_auth_token=HUGGING_FACE_TOKEN,
48
    )
49
50
    if pipeline is None:
51
        raise ValueError("Failed to initialize the pipeline. Please check your authentication token and internet connection.")
52
    else:
53
        pipeline = pipeline.to(device)
54
55
    # Apply the pipeline to the audio file
56
    diarization = pipeline(audio_file)
57
58
    # Create a dictionary to store speaker segments
59
    speaker_segments = {}
60
61
    # Process diarization results
62
    for turn, _, speaker in diarization.itertracks(yield_label=True):
63
        start, end = turn.start, turn.end
64
        if speaker not in speaker_segments:
65
            speaker_segments[speaker] = []
66
        speaker_segments[speaker].append((start, end))
67
68
    # Convert speaker_segments to a DataFrame
69
    diarize_df = pd.DataFrame([(speaker, start, end)
70
                               for speaker, segments in speaker_segments.items()
71
                               for start, end in segments],
72
                              columns=['speaker', 'start', 'end'])
73
74
    # Assign speakers to transcript words
75
    for word in transcript.words:
76
        word_start = float(word.start) / 1000
77
        word_end = float(word.end) / 1000
78
79
        overlaps = diarize_df[
80
            (diarize_df['start'] <= word_end) & (diarize_df['end'] >= word_start)
81
        ].copy()
82
83
        if not overlaps.empty:
84
            overlaps['overlap'] = np.minimum(overlaps['end'], word_end) - np.maximum(overlaps['start'], word_start)
85
            word.speaker = overlaps.loc[overlaps['overlap'].idxmax(), 'speaker']
86
        else:
87
            word.speaker = "Unknown"
88
89
    full_transcript = ''
90
91
    # Update segment speakers based on the majority speaker of its words
92
    for segment in transcript.get_sentences():
93
        segment_start = float(segment.start) / 1000
94
        segment_end = float(segment.end) / 1000
95
96
        overlaps = diarize_df[
97
            (diarize_df['start'] <= segment_end) & (diarize_df['end'] >= segment_start)
98
        ].copy()
99
100
        if not overlaps.empty:
101
            overlaps['overlap'] = np.minimum(overlaps['end'], segment_end) - np.maximum(overlaps['start'], segment_start)
102
            segment.speaker = overlaps.loc[overlaps['overlap'].idxmax(), 'speaker']
103
            speaker_label = segment.speaker.replace('SPEAKER_', 'SPEAKER ')
104
            full_transcript += f'[{format_timestamp(segment_start)}] {speaker_label}: {segment.text}\n'
105
        else:
106
            segment.speaker = "Unknown"
107
            full_transcript += f'[{format_timestamp(segment_start)}] Unknown: {segment.text}\n'
108
109
    return full_transcript
110
111
def format_timestamp(seconds):
112
    """
113
    Convert seconds to a formatted timestamp string (HH:MM:SS).
114
115
    Args:
116
        seconds (float): Time in seconds.
117
118
    Returns:
119
        str: Formatted timestamp string.
120
    """
121
    hours, remainder = divmod(int(seconds), 3600)
122
    minutes, seconds = divmod(remainder, 60)
123
    return f"{hours:02d}:{minutes:02d}:{seconds:02d}"
124
125
audio_file = "audio.wav" # your local file path
126
transcript: aai.Transcript = transcribe_audio(audio_file, language="hr") # select a language code
127
transcript_with_speakers = get_speaker_labels(audio_file, transcript)
128
print(transcript_with_speakers)

1
import os
2
import assemblyai as aai
3
from pyannote.audio import Pipeline
4
import torch
5
import pandas as pd
6
import numpy as np
7
8
# Assign your API keys
9
HUGGING_FACE_TOKEN = os.getenv("HF_TOKEN")
10
ASSEMBLYAI_API_KEY = os.getenv("ASSEMBLYAI_API_KEY")
11
12
# Authenticate with AssemblyAI
13
aai.settings.api_key = ASSEMBLYAI_API_KEY

1
def transcribe_audio(audio_file, language="en"):
2
    """
3
    Transcribe an audio file using AssemblyAI.
4
5
    Args:
6
        audio_file (str): Path to the audio file.
7
        language (str, optional): Language code for transcription. Defaults to "en".
8
9
    Returns:
10
        aai.Transcript: The transcription result.
11
    """
12
13
    transcriber = aai.Transcriber(config=aai.TranscriptionConfig(speech_model='nano', language_code=language))
14
    transcript = transcriber.transcribe(audio_file)
15
    print(f"Transcript ID: {transcript.id}")
16
    return transcript

1
def get_speaker_labels(audio_file, transcript: aai.Transcript):
2
    """
3
    Perform speaker diarization on an audio file and combine results with the transcript.
4
5
    Args:
6
        audio_file (str): Path to the audio file.
7
        transcript (aai.Transcript): The transcription result from AssemblyAI.
8
9
    Returns:
10
        str: A formatted string containing the transcript with speaker labels and timestamps.
11
    """
12
    # Initialize the speaker diarization pipeline with GPU support
13
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14
    pipeline = Pipeline.from_pretrained(
15
        "pyannote/speaker-diarization",
16
        use_auth_token=HUGGING_FACE_TOKEN,
17
    )
18
19
    if pipeline is None:
20
        raise ValueError("Failed to initialize the pipeline. Please check your authentication token and internet connection.")
21
    else:
22
        pipeline = pipeline.to(device)
23
24
    # Apply the pipeline to the audio file
25
    diarization = pipeline(audio_file)
26
27
    # Create a dictionary to store speaker segments
28
    speaker_segments = {}
29
30
    # Process diarization results
31
    for turn, _, speaker in diarization.itertracks(yield_label=True):
32
        start, end = turn.start, turn.end
33
        if speaker not in speaker_segments:
34
            speaker_segments[speaker] = []
35
        speaker_segments[speaker].append((start, end))
36
37
    # Convert speaker_segments to a DataFrame
38
    diarize_df = pd.DataFrame([(speaker, start, end)
39
                               for speaker, segments in speaker_segments.items()
40
                               for start, end in segments],
41
                              columns=['speaker', 'start', 'end'])
42
43
    # Assign speakers to transcript words
44
    for word in transcript.words:
45
        word_start = float(word.start) / 1000
46
        word_end = float(word.end) / 1000
47
48
        overlaps = diarize_df[
49
            (diarize_df['start'] <= word_end) & (diarize_df['end'] >= word_start)
50
        ].copy()
51
52
        if not overlaps.empty:
53
            overlaps['overlap'] = np.minimum(overlaps['end'], word_end) - np.maximum(overlaps['start'], word_start)
54
            word.speaker = overlaps.loc[overlaps['overlap'].idxmax(), 'speaker']
55
        else:
56
            word.speaker = "Unknown"
57
58
    full_transcript = ''
59
60
    # Update segment speakers based on the majority speaker of its words
61
    for segment in transcript.get_sentences():
62
        segment_start = float(segment.start) / 1000
63
        segment_end = float(segment.end) / 1000
64
65
        overlaps = diarize_df[
66
            (diarize_df['start'] <= segment_end) & (diarize_df['end'] >= segment_start)
67
        ].copy()
68
69
        if not overlaps.empty:
70
            overlaps['overlap'] = np.minimum(overlaps['end'], segment_end) - np.maximum(overlaps['start'], segment_start)
71
            segment.speaker = overlaps.loc[overlaps['overlap'].idxmax(), 'speaker']
72
            speaker_label = segment.speaker.replace('SPEAKER_', 'SPEAKER ')
73
            full_transcript += f'[{format_timestamp(segment_start)}] {speaker_label}: {segment.text}\n'
74
        else:
75
            segment.speaker = "Unknown"
76
            full_transcript += f'[{format_timestamp(segment_start)}] Unknown: {segment.text}\n'
77
78
    return full_transcript

1
def format_timestamp(seconds):
2
    """
3
    Convert seconds to a formatted timestamp string (HH:MM:SS).
4
5
    Args:
6
        seconds (float): Time in seconds.
7
8
    Returns:
9
        str: Formatted timestamp string.
10
    """
11
    hours, remainder = divmod(int(seconds), 3600)
12
    minutes, seconds = divmod(remainder, 60)
13
    return f"{hours:02d}:{minutes:02d}:{seconds:02d}"

1
audio_file = "audio.wav" # your local file path
2
transcript: aai.Transcript = transcribe_audio(audio_file, language="hr") # select a language code
3
transcript_with_speakers = get_speaker_labels(audio_file, transcript)
4
print(transcript_with_speakers)