Generate Custom Speaker Labels with Pyannote

import os
import assemblyai as aai
from pyannote.audio import Pipeline
import torch
import pandas as pd
import numpy as np

# Assign your API keys
HUGGING_FACE_TOKEN = os.getenv("HF_TOKEN")
ASSEMBLYAI_API_KEY = os.getenv("ASSEMBLYAI_API_KEY")

# Authenticate with AssemblyAI
aai.settings.api_key = ASSEMBLYAI_API_KEY

def transcribe_audio(audio_file, language="en"):
    """
    Transcribe an audio file using AssemblyAI.

    Args:
        audio_file (str): Path to the audio file.
        language (str, optional): Language code for transcription. Defaults to "en".

    Returns:
        aai.Transcript: The transcription result.
    """

    transcriber = aai.Transcriber(config=aai.TranscriptionConfig(speech_models=["universal-3-pro", "universal-2"], language_code=language))
    transcript = transcriber.transcribe(audio_file)
    print(f"Transcript ID: {transcript.id}")
    return transcript

def get_speaker_labels(audio_file, transcript: aai.Transcript):
    """
    Perform speaker diarization on an audio file and combine results with the transcript.

    Args:
        audio_file (str): Path to the audio file.
        transcript (aai.Transcript): The transcription result from AssemblyAI.

    Returns:
        str: A formatted string containing the transcript with speaker labels and timestamps.
    """
    # Initialize the speaker diarization pipeline with GPU support
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    pipeline = Pipeline.from_pretrained(
        "pyannote/speaker-diarization",
        use_auth_token=HUGGING_FACE_TOKEN,
    )

    if pipeline is None:
        raise ValueError("Failed to initialize the pipeline. Please check your authentication token and internet connection.")
    else:
        pipeline = pipeline.to(device)

    # Apply the pipeline to the audio file
    diarization = pipeline(audio_file)

    # Create a dictionary to store speaker segments
    speaker_segments = {}

    # Process diarization results
    for turn, _, speaker in diarization.itertracks(yield_label=True):
        start, end = turn.start, turn.end
        if speaker not in speaker_segments:
            speaker_segments[speaker] = []
        speaker_segments[speaker].append((start, end))

    # Convert speaker_segments to a DataFrame
    diarize_df = pd.DataFrame([(speaker, start, end)
                               for speaker, segments in speaker_segments.items()
                               for start, end in segments],
                              columns=['speaker', 'start', 'end'])

    # Assign speakers to transcript words
    for word in transcript.words:
        word_start = float(word.start) / 1000
        word_end = float(word.end) / 1000

        overlaps = diarize_df[
            (diarize_df['start'] <= word_end) & (diarize_df['end'] >= word_start)
        ].copy()

        if not overlaps.empty:
            overlaps['overlap'] = np.minimum(overlaps['end'], word_end) - np.maximum(overlaps['start'], word_start)
            word.speaker = overlaps.loc[overlaps['overlap'].idxmax(), 'speaker']
        else:
            word.speaker = "Unknown"

    full_transcript = ''

    # Update segment speakers based on the majority speaker of its words
    for segment in transcript.get_sentences():
        segment_start = float(segment.start) / 1000
        segment_end = float(segment.end) / 1000

        overlaps = diarize_df[
            (diarize_df['start'] <= segment_end) & (diarize_df['end'] >= segment_start)
        ].copy()

        if not overlaps.empty:
            overlaps['overlap'] = np.minimum(overlaps['end'], segment_end) - np.maximum(overlaps['start'], segment_start)
            segment.speaker = overlaps.loc[overlaps['overlap'].idxmax(), 'speaker']
            speaker_label = segment.speaker.replace('SPEAKER_', 'SPEAKER ')
            full_transcript += f'[{format_timestamp(segment_start)}] {speaker_label}: {segment.text}\n'
        else:
            segment.speaker = "Unknown"
            full_transcript += f'[{format_timestamp(segment_start)}] Unknown: {segment.text}\n'

    return full_transcript

def format_timestamp(seconds):
    """
    Convert seconds to a formatted timestamp string (HH:MM:SS).

    Args:
        seconds (float): Time in seconds.

    Returns:
        str: Formatted timestamp string.
    """
    hours, remainder = divmod(int(seconds), 3600)
    minutes, seconds = divmod(remainder, 60)
    return f"{hours:02d}:{minutes:02d}:{seconds:02d}"

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

def get_speaker_labels(audio_file, transcript: aai.Transcript):
    """
    Perform speaker diarization on an audio file and combine results with the transcript.

    Args:
        audio_file (str): Path to the audio file.
        transcript (aai.Transcript): The transcription result from AssemblyAI.

    Returns:
        str: A formatted string containing the transcript with speaker labels and timestamps.
    """
    # Initialize the speaker diarization pipeline with GPU support
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    pipeline = Pipeline.from_pretrained(
        "pyannote/speaker-diarization",
        use_auth_token=HUGGING_FACE_TOKEN,
    )

    if pipeline is None:
        raise ValueError("Failed to initialize the pipeline. Please check your authentication token and internet connection.")
    else:
        pipeline = pipeline.to(device)

    # Apply the pipeline to the audio file
    diarization = pipeline(audio_file)

    # Create a dictionary to store speaker segments
    speaker_segments = {}

    # Process diarization results
    for turn, _, speaker in diarization.itertracks(yield_label=True):
        start, end = turn.start, turn.end
        if speaker not in speaker_segments:
            speaker_segments[speaker] = []
        speaker_segments[speaker].append((start, end))

    # Convert speaker_segments to a DataFrame
    diarize_df = pd.DataFrame([(speaker, start, end)
                               for speaker, segments in speaker_segments.items()
                               for start, end in segments],
                              columns=['speaker', 'start', 'end'])

    # Assign speakers to transcript words
    for word in transcript.words:
        word_start = float(word.start) / 1000
        word_end = float(word.end) / 1000

        overlaps = diarize_df[
            (diarize_df['start'] <= word_end) & (diarize_df['end'] >= word_start)
        ].copy()

        if not overlaps.empty:
            overlaps['overlap'] = np.minimum(overlaps['end'], word_end) - np.maximum(overlaps['start'], word_start)
            word.speaker = overlaps.loc[overlaps['overlap'].idxmax(), 'speaker']
        else:
            word.speaker = "Unknown"

    full_transcript = ''

    # Update segment speakers based on the majority speaker of its words
    for segment in transcript.get_sentences():
        segment_start = float(segment.start) / 1000
        segment_end = float(segment.end) / 1000

        overlaps = diarize_df[
            (diarize_df['start'] <= segment_end) & (diarize_df['end'] >= segment_start)
        ].copy()

        if not overlaps.empty:
            overlaps['overlap'] = np.minimum(overlaps['end'], segment_end) - np.maximum(overlaps['start'], segment_start)
            segment.speaker = overlaps.loc[overlaps['overlap'].idxmax(), 'speaker']
            speaker_label = segment.speaker.replace('SPEAKER_', 'SPEAKER ')
            full_transcript += f'[{format_timestamp(segment_start)}] {speaker_label}: {segment.text}\n'
        else:
            segment.speaker = "Unknown"
            full_transcript += f'[{format_timestamp(segment_start)}] Unknown: {segment.text}\n'

    return full_transcript