k_compare_transcripts_performance.py

# STEPS
# - Download all files.
# - get each of those memories
# - read local json with each memory audio file
# - call whisper groq (whisper-largev3)
# - Create a table df, well printed, with each transcript result side by side
# - prompt for computing WER using groq whisper as baseline (if better, but most likely)
# - Run for deepgram and generate comparison result
import asyncio

import json
import os
import re
from collections import defaultdict
from itertools import islice
from typing import Any, Dict, List, cast

import firebase_admin
import requests
from dotenv import load_dotenv
from pydub import AudioSegment
from tabulate import tabulate

load_dotenv('../../.dev.env')
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = '../../' + (os.getenv('GOOGLE_APPLICATION_CREDENTIALS') or '')
cast(Any, firebase_admin).initialize_app()

from models.transcript_segment import TranscriptSegment
from utils.stt.streaming import process_audio_dg
from groq import Groq

from utils.other.storage import upload_postprocessing_audio


def add_model_result_segments(model: str, new_segments: List[Dict[str, Any]], result: Dict[str, Any]) -> None:
    segments = [TranscriptSegment(**s) for s in result[model]]
    new_seg_objs = [TranscriptSegment(**s) for s in new_segments]
    segments, _, _ = TranscriptSegment.combine_segments(segments, new_seg_objs)
    result[model] = [s.dict() for s in segments]


def execute_groq(file_path: str) -> str:
    file_size = os.path.getsize(file_path)
    print('execute_groq file_size', file_size / 1024 / 1024, 'MB')
    split_files: List[str] = []
    if file_size / 1024 / 1024 > 25:
        # split file
        aseg = cast(Any, AudioSegment).from_wav(file_path)
        # split every 10 minutes
        split_duration = 10 * 60 * 1000
        for i in range(0, len(aseg), split_duration):
            split_file_path = f'{file_path}_{i}.wav'
            split_files.append(split_file_path)
            aseg[i : i + split_duration].export(split_file_path, format="wav")
    else:
        split_files.append(file_path)

    client = Groq(api_key=os.getenv('GROQ_API_KEY'))
    result = ''
    for split_file_path in split_files:
        with open(split_file_path, "rb") as file:
            transcription = client.audio.transcriptions.create(
                file=(split_file_path, file.read()),
                model="whisper-large-v3",
                response_format="text",
                language="en",
                temperature=0.0,
            )
            result += ' ' + str(transcription)
    return result.strip().lower().replace('  ', ' ')


async def _execute_single(file_path: str) -> None:
    aseg = cast(Any, AudioSegment).from_wav(file_path)
    duration: float = aseg.duration_seconds
    memory_id = file_path.split('/')[-1].split('.')[0]

    if os.path.exists(f'results/{memory_id}.json'):
        print('Already processed', memory_id)
        return
    if aseg.duration_seconds < 5:
        print('Skipping', memory_id, 'duration', aseg.duration_seconds)
        return

    print('Started processing', memory_id, 'duration', aseg.duration_seconds)
    result: Dict[str, Any] = {'deepgram': []}

    def stream_transcript_deepgram(new_segments: List[Dict[str, Any]]) -> None:
        print('stream_transcript_deepgram', new_segments)
        add_model_result_segments('deepgram', new_segments, result)

    # streaming models
    socket = await process_audio_dg(stream_transcript_deepgram, language='en', sample_rate=16000, channels=1)
    if socket is None:
        print('Failed to obtain socket for', memory_id)
        return
    print('duration', duration)
    with open(file_path, "rb") as file:
        while True:
            chunk = file.read(320)
            if not chunk:
                break
            socket.send(bytes(chunk))
            await asyncio.sleep(0.005)

    print('Finished sending audio')
    groq_result: str = execute_groq(file_path)  # source of truth
    result['whisper-large-v3'] = groq_result

    print('Waiting for sockets to finish', min(60, duration), 'seconds')
    await asyncio.sleep(min(30, duration))

    os.makedirs('results', exist_ok=True)
    with open(f'results/{memory_id}.json', 'w') as f:
        json.dump(result, f, indent=2)

    socket.finish()


def batched(iterable: Any, n: int) -> Any:
    """
    Generator that yields lists of size 'n' from 'iterable'.
    """
    it = iter(iterable)
    while True:
        batch: List[Any] = list(islice(it, n))
        if not batch:
            break
        yield batch


async def process_memories_audio_files() -> None:
    uids = os.listdir('_temp2')
    for uid in uids:
        memories = os.listdir(f'_temp2/{uid}')
        memories = [f'_temp2/{uid}/{memory}' for memory in memories]
        # batch_size = 5
        for memory in memories:
            await _execute_single(memory)
        # for batch_num, batch in enumerate(batched(memories, batch_size), start=1):
        #     tasks = [asyncio.create_task(_execute_single(file_path)) for file_path in batch]
        #     await asyncio.gather(*tasks)
        #     print(f'Batch {batch_num} processed')


from jiwer import wer


def compute_wer() -> None:
    """
    Computes the Word Error Rate (WER) for each transcription model against a reference model
    across all JSON files in the specified directory. Outputs detailed results and overall rankings.
    """
    dir_path = 'results/'  # Directory containing JSON files
    reference_model = 'whisper-large-v3'  # Reference model key
    table_data: List[List[Any]] = []  # List to hold detailed table rows
    wer_accumulator: defaultdict[str, List[float]] = defaultdict(list)  # To accumulate WERs per model
    points_counter: defaultdict[str, int] = defaultdict(int)  # To count points per model based on WER rankings

    # Define detailed table headers
    detailed_headers = [
        "File",
        "Model",
        "WER",
        "Source Words",
        "Model Words",
        "Source Characters",
        "Model Characters",
        "Transcript",
    ]

    # Check if the directory exists
    if not os.path.isdir(dir_path):
        print(f"Directory '{dir_path}' does not exist.")
        return

    # Iterate through all files in the specified directory
    for file in os.listdir(dir_path):
        if not file.endswith('.json'):
            continue  # Skip non-JSON files

        file_path = os.path.join(dir_path, file)
        with open(file_path, 'r', encoding='utf-8') as f:
            try:
                result: Any = json.load(f)
            except json.JSONDecodeError:
                print(f"Error decoding JSON in file: {file}")
                continue  # Skip files with invalid JSON

        # Check if the reference model exists in the JSON
        if reference_model not in result:
            print(f"Reference model '{reference_model}' not found in file: {file}")
            continue  # Skip files without the reference model

        # Assemble the reference transcript
        reference_text: Any = regex_fix(result.get(reference_model, ''))
        if isinstance(reference_text, list):
            # If reference_text is a list of segments
            reference_text = (
                ' '.join([str(segment.get('text', '')) for segment in cast(List[Any], reference_text)]).strip().lower()
            )
        else:
            # If reference_text is a single string
            reference_text = str(reference_text).strip().lower()
        reference_text = ' '.join(reference_text.split())  # Normalize whitespace

        # Calculate source words and characters
        source_words = len(reference_text.split())
        source_characters = len(reference_text)

        print(f"Processing file: {file}")

        # Temporary storage for current file's model WERs to determine ranking points
        current_file_wer: Dict[str, float] = {}

        # Iterate through each model in the JSON
        for model, segments in result.items():
            if model == reference_model:
                model_text = reference_text  # Reference model's transcript
            else:
                if isinstance(segments, list):
                    # Assemble the model's transcript from segments
                    model_text = (
                        ' '.join([str(segment.get('text', '')) for segment in cast(List[Any], segments)])
                        .strip()
                        .lower()
                    )
                else:
                    # If segments is a single string
                    model_text = str(segments).strip().lower()
                model_text = ' '.join(model_text.split())  # Normalize whitespace

            # Calculate model words and characters
            model_words = len(model_text.split())
            model_characters = len(model_text)

            # Compute WER
            current_wer: float = wer(reference_text, model_text)

            # Accumulate WER for overall statistics (exclude reference model)
            if model != reference_model:
                wer_accumulator[model].append(current_wer)

            # Store WER for current file's ranking
            if model != reference_model:
                current_file_wer[model] = current_wer

            # Append the data to the detailed table
            table_data.append(
                [
                    file,
                    model,
                    f"{current_wer:.2%}",
                    source_words,
                    model_words,
                    source_characters,
                    model_characters,
                    model_text,
                ]
            )

        # Determine which model(s) had the lowest WER in the current file
        if current_file_wer:
            min_wer = min(current_file_wer.values())
            best_models = [model for model, w in current_file_wer.items() if w == min_wer]
            for model in best_models:
                points_counter[model] += 1  # Assign 1 point to each best model

        print('-----------------------------------------')

    # Generate the detailed WER table using tabulate
    if table_data:
        print("\nDetailed WER Results:")
        detailed_table = tabulate(table_data, headers=detailed_headers, tablefmt="grid", stralign="left")
        with open('results/detailed_wer.txt', 'w') as f:
            f.write(detailed_table)
    else:
        print("No data to display.")

    # Compute overall WER per model (average)
    overall_wer: Dict[str, float] = {}
    for model, wer_list in wer_accumulator.items():
        if wer_list:
            overall_wer[model] = sum(wer_list) / len(wer_list)

    # Create a list for overall WER table
    overall_wer_table: List[List[Any]] = []
    for model, avg_wer in overall_wer.items():
        overall_wer_table.append([model, f"{avg_wer:.2%}"])

    # Sort the overall WER table by average WER ascending (lower is better)
    overall_wer_table_sorted = sorted(overall_wer_table, key=lambda x: x[1])

    # Define overall WER table headers
    overall_wer_headers = ["Model", "Average WER"]

    # Generate the overall WER table
    if overall_wer_table_sorted:
        print("\nOverall WER per Model:")
        overall_wer_formatted = tabulate(
            overall_wer_table_sorted, headers=overall_wer_headers, tablefmt="grid", stralign="left"
        )
        print(overall_wer_formatted)
        with open('results/wer.txt', 'w') as f:
            f.write(overall_wer_formatted)
    else:
        print("No overall WER data to display.")

    # Create a ranking table based on points
    ranking_table: List[List[Any]] = []
    for model, points in points_counter.items():
        ranking_table.append([model, points])

    # Sort the ranking table by points descending (more points are better)
    ranking_table_sorted = sorted(ranking_table, key=lambda x: x[1], reverse=True)

    # Assign rankings
    ranking_table_with_rank: List[List[Any]] = []
    current_rank = 1
    previous_points: Any = None
    for model, points in ranking_table_sorted:
        if points != previous_points:
            rank = current_rank
        else:
            rank = current_rank - 1  # Same rank as previous
        ranking_table_with_rank.append([rank, model, points])
        previous_points = points
        current_rank += 1

    # Define ranking table headers
    ranking_headers = ["Rank", "Model", "Points"]

    # Generate the ranking table
    if ranking_table_with_rank:
        print("\nModel Rankings Based on WER Performance:")
        ranking_table_formatted = tabulate(
            ranking_table_with_rank, headers=ranking_headers, tablefmt="grid", stralign="left"
        )
        print(ranking_table_formatted)
        with open('results/ranking.txt', 'w') as f:
            f.write(ranking_table_formatted)
    else:
        print("No ranking data to display.")


def regex_fix(text: str) -> str:
    """Fix some of the stored JSON in results/$id.json from the Groq API."""
    pattern = r'(?<=transcription\(text=["\'])(.*?)(?=["\'],\s*task=)'
    match = re.search(pattern, text)
    if match:
        extracted_text = match.group(0)
        return extracted_text
    else:
        print("No match found.")
        return text


def pyannote_diarize(file_path: str) -> None:
    memory_id = file_path.split('/')[-1].split('.')[0]
    with open('diarization.json', 'r') as f:
        results = json.loads(f.read())

    if memory_id in results:
        print('Already diarized', memory_id)
        return

    url = "https://api.pyannote.ai/v1/diarize"
    headers = {"Authorization": f"Bearer {os.getenv('PYANNOTE_API_KEY')}"}
    webhook = 'https://camel-lucky-reliably.ngrok-free.app/webhook'
    signed_url = upload_postprocessing_audio(file_path)
    data = {'webhook': webhook, 'url': signed_url}
    response = requests.post(url, headers=headers, json=data)
    print(memory_id, response.json()['jobId'])
    # update diarization.json, and set jobId=memoryId
    with open('diarization.json', 'r') as f:
        diarization = json.loads(f.read())

    diarization[response.json()['jobId']] = memory_id
    with open('diarization.json', 'w') as f:
        json.dump(diarization, f, indent=2)


def generate_diarizations() -> None:
    uids = os.listdir('_temp2')
    for uid in uids:
        memories = os.listdir(f'_temp2/{uid}')
        memories = [f'_temp2/{uid}/{memory}' for memory in memories]
        for memory in memories:
            memory_id = memory.split('/')[-1].split('.')[0]
            if os.path.exists(f'results/{memory_id}.json'):
                pyannote_diarize(memory)
            else:
                print('Skipping', memory_id)


from pyannote.metrics.diarization import DiarizationErrorRate  # type: ignore[reportMissingImports]
from pyannote.core import Annotation, Segment  # type: ignore[reportMissingImports]

der_metric = cast(Any, DiarizationErrorRate())


def compute_der() -> None:
    """
    Computes the Diarization Error Rate (DER) for each model across all JSON files in the 'results/' directory.
    Outputs a summary table and rankings to 'der_report.txt'.
    """
    dir_path = 'results/'  # Directory containing result JSON files and 'diarization.json'
    output_file = os.path.join(dir_path, 'der_report.txt')  # Output report file
    excluded_model = 'whisper-large-v3'  # Model to exclude from analysis

    # Initialize DER metric
    local_der_metric = cast(Any, DiarizationErrorRate())

    # Check if the directory exists
    if not os.path.isdir(dir_path):
        print(f"Directory '{dir_path}' does not exist.")
        return

    # Path to 'diarization.json'
    diarization_path = 'diarization.json'

    # Load reference diarization data
    with open(diarization_path, 'r', encoding='utf-8') as f:
        try:
            diarization: Any = json.load(f)
        except json.JSONDecodeError:
            print(f"Error decoding JSON in 'diarization.json'.")
            return

    # Prepare to collect DER results
    der_results: List[List[Any]] = []  # List to store [Memory ID, Model, DER]
    model_der_accumulator: defaultdict[str, List[float]] = defaultdict(list)  # To calculate average DER per model

    # Iterate through all JSON files in 'results/' directory
    for file in os.listdir(dir_path):
        if not file.endswith('.json') or file == 'diarization.json':
            continue  # Skip non-JSON files and 'diarization.json' itself

        memory_id = file.split('.')[0]  # Extract memory ID from filename

        # Check if memory_id exists in 'diarization.json'
        if memory_id not in diarization:
            print(f"Memory ID '{memory_id}' not found in 'diarization.json'. Skipping file: {file}")
            continue

        # Load reference segments for the current memory_id
        ref_segments = diarization[memory_id]
        ref_annotation = cast(Any, Annotation())
        for seg in ref_segments:
            speaker, start, end = seg['speaker'], seg['start'], seg['end']
            ref_annotation[Segment(start, end)] = speaker

        # Load hypothesis segments from the result JSON file
        file_path = os.path.join(dir_path, file)
        with open(file_path, 'r', encoding='utf-8') as f:
            try:
                data: Any = json.load(f)
            except json.JSONDecodeError:
                print(f"Error decoding JSON in file: {file}. Skipping.")
                continue

        # Iterate through each model's segments in the result
        for model, segments in data.items():
            if model == excluded_model:
                continue  # Skip the excluded model

            hyp_annotation = cast(Any, Annotation())
            for seg in segments:
                speaker, start, end = seg['speaker'], seg['start'], seg['end']
                # Optional: Normalize speaker labels if necessary
                if speaker == 'SPEAKER_0':
                    speaker = 'SPEAKER_00'
                elif speaker == 'SPEAKER_1':
                    speaker = 'SPEAKER_01'
                elif speaker == 'SPEAKER_2':
                    speaker = 'SPEAKER_02'
                elif speaker == 'SPEAKER_3':
                    speaker = 'SPEAKER_03'
                hyp_annotation[Segment(start, end)] = speaker

            # Compute DER between reference and hypothesis
            der = cast(float, local_der_metric(ref_annotation, hyp_annotation))

            # Store the result
            der_results.append([memory_id, model, f"{der:.2%}"])
            model_der_accumulator[model].append(der)

    # Generate the detailed DER table
    der_table = tabulate(der_results, headers=["Memory ID", "Model", "DER"], tablefmt="grid", stralign="left")

    # Calculate average DER per model
    average_der: List[List[Any]] = []
    for model, ders in model_der_accumulator.items():
        avg = sum(ders) / len(ders)
        average_der.append([model, f"{avg:.2%}"])

    # Sort models by average DER ascending (lower is better)
    average_der_sorted = sorted(average_der, key=lambda x: float(str(x[1]).strip('%')))

    # Determine the winner (model with the lowest average DER)
    winner = average_der_sorted[0][0] if average_der_sorted else "N/A"

    # Prepare rankings (1st, 2nd, etc.)
    rankings: List[List[Any]] = []
    rank = 1
    previous_der: Any = None
    for model, avg in average_der_sorted:
        current_der = float(str(avg).strip('%'))
        if previous_der is None or current_der < previous_der:
            current_rank = rank
        else:
            current_rank = rank - 1  # Same rank as previous if DER is equal
        rankings.append([current_rank, model, avg])
        previous_der = current_der
        rank += 1

    # Generate the rankings table
    ranking_table = tabulate(rankings, headers=["Rank", "Model", "Average DER"], tablefmt="grid", stralign="left")

    # Write all results to the output file
    with open(output_file, 'w', encoding='utf-8') as out_f:
        out_f.write("Diarization Error Rate (DER) Analysis Report\n")
        out_f.write("=" * 50 + "\n\n")
        out_f.write("Detailed DER Results:\n")
        out_f.write(der_table + "\n\n")
        out_f.write("Average DER per Model:\n")
        out_f.write(
            tabulate(average_der_sorted, headers=["Model", "Average DER"], tablefmt="grid", stralign="left") + "\n\n"
        )
        out_f.write("Model Rankings Based on Average DER:\n")
        out_f.write(ranking_table + "\n\n")
        out_f.write(f"Winner: {winner}\n")

    # Print a confirmation message
    print(f"Diarization Error Rate (DER) analysis completed. Report saved to '{output_file}'.")

    # Optionally, print the tables to the console as well
    if der_results:
        print("\nDetailed DER Results:")
        print(der_table)
    if average_der_sorted:
        print("\nAverage DER per Model:")
        print(tabulate(average_der_sorted, headers=["Model", "Average DER"], tablefmt="grid", stralign="left"))
    if rankings:
        print("\nModel Rankings Based on Average DER:")
        print(ranking_table)
        print(f"\nWinner: {winner}")


if __name__ == '__main__':
    # asyncio.run(process_memories_audio_files())
    # generate_diarizations()

    # compute_wer()
    compute_der()