scripts/benchmark_stt.py

"""
STT Benchmarking Script

This script benchmarks different Speech-to-Text (STT) handlers to compare their performance.
Measures: inference time, warmup time, memory usage, and transcription quality.

Usage:
    python benchmark_stt.py --audio_file path/to/audio.wav --iterations 10
    python benchmark_stt.py --audio_file path/to/audio.wav --handlers whisper mlx-audio-whisper
"""

import argparse
import json
import logging
import time
from pathlib import Path
from queue import Queue
from threading import Event
from typing import Any, Dict, List, Optional

import numpy as np
import soundfile as sf

from speech_to_speech.pipeline.messages import VADAudio

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


class BenchmarkResult:
    """Stores benchmark results for a single STT handler."""

    def __init__(self, handler_name: str):
        self.handler_name = handler_name
        self.warmup_time = 0.0
        self.inference_times: list[float] = []
        self.time_to_first_token: list[float] = []
        self.transcriptions: list[str] = []
        self.errors: list[str] = []

    def add_inference(self, time_taken: float, transcription: Any, ttft: Optional[float] = None):
        self.inference_times.append(time_taken)
        self.transcriptions.append(transcription)
        if ttft is not None:
            self.time_to_first_token.append(ttft)

    def add_error(self, error: str):
        self.errors.append(error)

    def get_stats(self) -> Dict[str, Any]:
        """Calculate statistics from benchmark results."""
        if not self.inference_times:
            return {
                "handler": self.handler_name,
                "status": "failed",
                "errors": self.errors,
            }

        stats = {
            "handler": self.handler_name,
            "warmup_time": self.warmup_time,
            "avg_inference_time": np.mean(self.inference_times),
            "min_inference_time": np.min(self.inference_times),
            "max_inference_time": np.max(self.inference_times),
            "std_inference_time": np.std(self.inference_times),
            "total_iterations": len(self.inference_times),
            "errors": self.errors,
            "sample_transcription": self.transcriptions[0] if self.transcriptions else None,
        }

        # Add time to first token stats if available
        if self.time_to_first_token:
            stats["avg_time_to_first_token"] = np.mean(self.time_to_first_token)
            stats["min_time_to_first_token"] = np.min(self.time_to_first_token)
            stats["max_time_to_first_token"] = np.max(self.time_to_first_token)
            stats["std_time_to_first_token"] = np.std(self.time_to_first_token)

        return stats


def load_audio(audio_path: str) -> np.ndarray:
    """Load audio file and return as numpy array."""
    logger.info(f"Loading audio from: {audio_path}")
    audio, sample_rate = sf.read(audio_path)

    # Convert to mono if stereo
    if len(audio.shape) > 1:
        audio = audio.mean(axis=1)

    # Resample to 16kHz if needed (most whisper models expect 16kHz)
    if sample_rate != 16000:
        logger.warning(f"Audio sample rate is {sample_rate}Hz, resampling to 16000Hz")
        try:
            import librosa
            audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=16000)
        except ImportError:
            logger.error("librosa not installed. Please install it with: pip install librosa")
            logger.error("Attempting scipy resampling as fallback...")
            from scipy import signal
            # Calculate resampling ratio
            num_samples = int(len(audio) * 16000 / sample_rate)
            audio = signal.resample(audio, num_samples)

    return audio.astype(np.float32)


def benchmark_handler(
    handler_name: str,
    audio: np.ndarray,
    iterations: int,
    handler_kwargs: Optional[Dict[str, Any]] = None
) -> BenchmarkResult:
    """Benchmark a single STT handler."""
    logger.info(f"Benchmarking {handler_name}...")
    result = BenchmarkResult(handler_name)

    try:
        # Create queues and events for handler
        stop_event = Event()
        queue_in: Queue[Any] = Queue()
        queue_out: Queue[Any] = Queue()

        handler: Any = None
        if handler_name == "whisper":
            from speech_to_speech.STT.whisper_stt_handler import WhisperSTTHandler
            setup_kwargs = handler_kwargs or {
                "model_name": "distil-whisper/distil-large-v3",
                "device": "cuda",
                "torch_dtype": "float16",
            }
            handler = WhisperSTTHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
                setup_kwargs=setup_kwargs
            )

        elif handler_name == "whisper-mlx":
            from speech_to_speech.STT.lightning_whisper_mlx_handler import LightningWhisperSTTHandler
            setup_kwargs = handler_kwargs or {
                "model_name": "large-v3",
                "device": "mps",
            }
            handler = LightningWhisperSTTHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
                setup_kwargs=setup_kwargs
            )

        elif handler_name == "mlx-audio-whisper":
            from speech_to_speech.STT.mlx_audio_whisper_handler import MLXAudioWhisperSTTHandler
            setup_kwargs = handler_kwargs or {
                "model_name": "mlx-community/whisper-large-v3-turbo",
            }
            handler = MLXAudioWhisperSTTHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
                setup_kwargs=setup_kwargs
            )

        elif handler_name == "faster-whisper":
            from speech_to_speech.STT.faster_whisper_handler import FasterWhisperSTTHandler
            setup_kwargs = handler_kwargs or {
                "model_name": "large-v3",
                "device": "auto",
                "compute_type": "float16",
            }
            handler = FasterWhisperSTTHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
                setup_kwargs=setup_kwargs
            )

        elif handler_name == "moonshine":
            from archive.STT.moonshine_handler import MoonshineSTTHandler
            handler = MoonshineSTTHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
            )

        elif handler_name == "parakeet-tdt":
            from speech_to_speech.STT.parakeet_tdt_handler import ParakeetTDTSTTHandler
            setup_kwargs = handler_kwargs or {
                "device": "mps",
                "enable_live_transcription": False,
            }
            handler = ParakeetTDTSTTHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
                setup_kwargs=setup_kwargs
            )

        elif handler_name == "parakeet-tdt-progressive":
            from speech_to_speech.STT.parakeet_tdt_handler import ParakeetTDTSTTHandler
            setup_kwargs = handler_kwargs or {
                "device": "mps",
                "enable_live_transcription": True,
                "live_transcription_update_interval": 0.25,
            }
            handler = ParakeetTDTSTTHandler(
                stop_event,
                queue_in=queue_in,
                queue_out=queue_out,
                setup_kwargs=setup_kwargs
            )
        else:
            raise ValueError(f"Unknown handler: {handler_name}")

        # Warmup is done in handler setup
        logger.info(f"Handler {handler_name} initialized and warmed up")

        # Additional warmup on the actual audio (excluded from timings)
        for _ in handler.process(VADAudio(audio=audio)):
            pass

        # Run benchmark iterations
        for i in range(iterations):
            logger.info(f"Iteration {i+1}/{iterations} for {handler_name}")

            start_time = time.perf_counter()

            # Process audio
            transcription = None
            time_to_first_token = None
            first_output = True

            for output in handler.process(VADAudio(audio=audio)):
                # Measure time to first token
                if first_output:
                    time_to_first_token = time.perf_counter() - start_time
                    first_output = False

                if isinstance(output, tuple):
                    transcription = output[0]  # (text, language)
                else:
                    transcription = output

            end_time = time.perf_counter()

            time_taken = end_time - start_time
            result.add_inference(time_taken, transcription, time_to_first_token)

            ttft_str = f", TTFT: {time_to_first_token:.4f}s" if time_to_first_token else ""
            text_preview = str(transcription)[:50] if transcription is not None else "(none)"
            logger.info(f"  Time: {time_taken:.4f}s{ttft_str}, Text: {text_preview}...")

        # Cleanup
        handler.cleanup()
        stop_event.set()

    except Exception as e:
        logger.error(f"Error benchmarking {handler_name}: {e}", exc_info=True)
        result.add_error(str(e))

    return result


def print_results(results: List[BenchmarkResult]):
    """Print benchmark results in a formatted table."""
    print("\n" + "="*80)
    print("BENCHMARK RESULTS")
    print("="*80)

    for result in results:
        stats = result.get_stats()
        print(f"\nHandler: {stats['handler']}")
        print("-" * 80)

        if stats.get("status") == "failed":
            print("  Status: FAILED")
            print(f"  Errors: {stats['errors']}")
            continue

        print(f"  Warmup Time:          {stats['warmup_time']:.4f}s")
        print(f"  Avg Inference Time:   {stats['avg_inference_time']:.4f}s")
        print(f"  Min Inference Time:   {stats['min_inference_time']:.4f}s")
        print(f"  Max Inference Time:   {stats['max_inference_time']:.4f}s")
        print(f"  Std Deviation:        {stats['std_inference_time']:.4f}s")

        # Print time to first token stats if available
        if 'avg_time_to_first_token' in stats:
            print("\n  Time to First Token:")
            print(f"    Avg TTFT:           {stats['avg_time_to_first_token']:.4f}s")
            print(f"    Min TTFT:           {stats['min_time_to_first_token']:.4f}s")
            print(f"    Max TTFT:           {stats['max_time_to_first_token']:.4f}s")
            print(f"    Std TTFT:           {stats['std_time_to_first_token']:.4f}s")

        print(f"\n  Total Iterations:     {stats['total_iterations']}")
        print(f"  Sample Transcription: {stats['sample_transcription']}")

        if stats['errors']:
            print(f"  Errors: {stats['errors']}")

    # Comparison table
    print("\n" + "="*80)
    print("COMPARISON (Average Inference Time)")
    print("="*80)

    successful_results = [r for r in results if r.inference_times]
    if successful_results:
        sorted_results = sorted(successful_results, key=lambda x: np.mean(x.inference_times))

        fastest = sorted_results[0]
        fastest_time = np.mean(fastest.inference_times)

        for result in sorted_results:
            avg_time = np.mean(result.inference_times)
            speedup = avg_time / fastest_time
            print(f"  {result.handler_name:25s}: {avg_time:.4f}s  ({speedup:.2f}x slower than fastest)")


def save_results(results: List[BenchmarkResult], output_file: str):
    """Save benchmark results to JSON file."""
    data = {
        "results": [r.get_stats() for r in results],
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
    }

    with open(output_file, 'w') as f:
        json.dump(data, f, indent=2)

    logger.info(f"Results saved to: {output_file}")


def main():
    parser = argparse.ArgumentParser(description="Benchmark STT handlers")
    parser.add_argument(
        "--audio_file",
        type=str,
        required=True,
        help="Path to audio file for benchmarking"
    )
    parser.add_argument(
        "--handlers",
        nargs="+",
        default=["whisper", "whisper-mlx", "mlx-audio-whisper", "faster-whisper", "parakeet-tdt", "parakeet-tdt-progressive"],
        help="List of handlers to benchmark (default: all)"
    )
    parser.add_argument(
        "--iterations",
        type=int,
        default=5,
        help="Number of iterations per handler (default: 5)"
    )
    parser.add_argument(
        "--output",
        type=str,
        default="stt_benchmark_results.json",
        help="Output JSON file for results (default: stt_benchmark_results.json)"
    )

    args = parser.parse_args()

    # Validate audio file exists
    if not Path(args.audio_file).exists():
        logger.error(f"Audio file not found: {args.audio_file}")
        return

    # Load audio
    audio = load_audio(args.audio_file)
    logger.info(f"Audio loaded: {len(audio)} samples, {len(audio)/16000:.2f}s duration")

    # Run benchmarks
    results = []
    for handler_name in args.handlers:
        result = benchmark_handler(handler_name, audio, args.iterations)
        results.append(result)

    # Print and save results
    print_results(results)
    save_results(results, args.output)

    logger.info("Benchmarking complete!")


if __name__ == "__main__":
    main()