evaluate_benchmarks.sh

#!/bin/bash
# =============================================================================
# ASR Model Benchmark Evaluation Script
# =============================================================================
# This script evaluates multiple ASR models across various benchmark datasets.
# 
# Usage:
#   ./evaluate_benchmarks.sh                    # Run with defaults
#   MAX_SAMPLES=100 ./evaluate_benchmarks.sh    # Custom sample count
#   MODELS="whisper" ./evaluate_benchmarks.sh   # Run specific models only
#
# Environment Variables:
#   MAX_SAMPLES     - Number of samples per benchmark (default: 10)
#   MODELS          - Comma-separated list of models to run (default: all)
#   OUTPUT_DIR      - Base output directory (default: results/benchmark)
# =============================================================================

set -e

# =============================================================================
# Configuration
# =============================================================================

# Defaults (can be overridden by environment variables)
MAX_SAMPLES="${MAX_SAMPLES:-10}" # None, 10, 100, 1000, etc.
OUTPUT_DIR="${OUTPUT_DIR:-results/benchmarks}"
CLI="python -m asr_evalkit"

# Model configurations: "evaluator|model_id|batch_size|extra_args"
#
# --- HuggingFace Models ---
#    --use-fp16           : Use mixed precision for HF models.
#    --use-flash-attention: Speed up HF inference (if supported/installed).
# Example: ["meralion"]="meralion|MERaLiON/MERaLiON-2-10B-ASR|8|--use-fp16 --revision #your_revision"
#
# --- vLLM Models ---
#    No extra args needed, memory utilization managed internally.
# Example: ["qwen3_asr"]="qwen3_asr|Qwen/Qwen3-ASR-0.6B|16|"
#
# --- NeMo Models ---
#    Precision managed automatically; generally no extra args required.
# Example: ["parakeet"]="parakeet|nvidia/parakeet-tdt-0.6b-v3|16|"
#
declare -A MODELS=(
    ["whisper"]="whisper|openai/whisper-large-v3-turbo|16|"
    # ["meralion"]="meralion|MERaLiON/MERaLiON-2-10B-ASR|1|--use-fp16"
    # ["qwen3_asr_0.6b"]="qwen3_asr|Qwen/Qwen3-ASR-0.6B|16|"
    # ["qwen3_asr_1.7b"]="qwen3_asr|Qwen/Qwen3-ASR-1.7B|16|"
    # ["polyglot_lion_1.7b"]="qwen3_asr|knoveleng/polyglot-lion-1.7b|16|"
    # ["polyglot_lion_0.6b"]="qwen3_asr|knoveleng/polyglot-lion-0.6b|16|"
    # # ["parakeet"]="parakeet|nvidia/parakeet-tdt-0.6b-v3|16|"
    # ["qwen_omni_7b"]="qwen_omni|Qwen/Qwen2.5-Omni-7B|1|--use-fp16"
    # ["qwen_omni_3b"]="qwen_omni|Qwen/Qwen2.5-Omni-3B|1|--use-fp16"
    # ["seallm"]="seallm_audio|SeaLLMs/SeaLLMs-Audio-7B|1|--use-fp16"
    # # ["phi4"]="phi4|microsoft/Phi-4-multimodal-instruct|1|"
)

# Benchmark datasets: "name|dataset_id|config|split|audio_col|text_col"
#
# --- Dataset Formats ---
# HuggingFace (Default): Uses dataset ID from the Hugging Face Hub.
#   Example: "English (LibriSpeech)|openslr/librispeech_asr|clean|test|audio|text"
#
# NeMo Manifests: To use local NeMo manifests, you need to add '--dataset-source nemo' 
#   to the CLI command in run_evaluation(). The dataset_id should be the path to the manifest.
#   Example: "NeMo Data|/path/to/manifest.jsonl|||audio_filepath|text"
#
DATASET_SOURCE="huggingface"
BENCHMARKS=(
    "cv_mandarin|knoveleng/cv-mandarin|default|test|audio|text"
    "aishell1_mandarin|knoveleng/aishell1-mandarin|default|test|audio|text"
    "aishell3_mandarin|knoveleng/aishell3-mandarin|default|test|audio|text"
    "fleurs_mandarin|knoveleng/fleurs-mandarin|default|test|audio|text"
    "fleurs_tamil|knoveleng/fleurs-tamil|default|test|audio|text"
    "slr127_tamil|knoveleng/slr127-tamil|default|test|audio|text"
    "slr65_tamil|knoveleng/slr65-tamil|default|test|audio|text"
    "cv_tamil|knoveleng/cv-tamil|default|test|audio|text"
    "mesolitica_malay|knoveleng/mesolitica-malay|default|test|audio|text"
    "fleurs_malay|knoveleng/fleurs-malay|default|test|audio|text"
    "nsc_singlish|knoveleng/nsc-singlish|default|test|audio|text"
    "librispeech_english|knoveleng/librispeech-english|default|test|audio|text" 
    # You can also use other hf datasets: E.g "librispeech_english|openslr/librispeech_asr|clean|test|audio|text"
)

# DATASET_SOURCE="nemo"
# BENCHMARKS=(
#     # Mandarin
#     "cv_mandarin|../dataset/knoveleng__cv-mandarin/test/knoveleng_cv_mandarin_test_asr.json|||audio|text"
#     "aishell1_mandarin|../dataset/knoveleng__aishell1-mandarin/test/knoveleng_aishell1_mandarin_test_asr.json|||audio|text"
#     "aishell3_mandarin|../dataset/knoveleng__aishell3-mandarin/test/knoveleng_aishell3_mandarin_test_asr.json|||audio|text"
#     "fleurs_mandarin|../dataset/knoveleng__fleurs-mandarin/test/knoveleng_fleurs_mandarin_test_asr.json|||audio|text"
#     # Tamil
#     "fleurs_tamil|../dataset/knoveleng__fleurs-tamil/test/knoveleng_fleurs_tamil_test_asr.json|||audio|text"
#     "slr127_tamil|../dataset/knoveleng__slr127-tamil/test/knoveleng_slr127_tamil_test_asr.json|||audio|text"
#     "slr65_tamil|../dataset/knoveleng__slr65-tamil/test/knoveleng_slr65_tamil_test_asr.json|||audio|text"
#     "cv_tamil|../dataset/knoveleng__cv-tamil/test/knoveleng_cv_tamil_test_asr.json|||audio|text"
#     # Malay
#     "mesolitica_malay|../dataset/knoveleng__mesolitica-malay/test/knoveleng_mesolitica_malay_test_asr.json|||audio|text"
#     "fleurs_malay|../dataset/knoveleng__fleurs-malay/test/knoveleng_fleurs_malay_test_asr.json|||audio|text"
#     # Singlish
#     "nsc_singlish|../dataset/knoveleng__nsc-singlish/test/knoveleng_nsc_singlish_test_asr.json|||audio|text"
#     # English
#     "librispeech_english|../dataset/knoveleng__librispeech-english/test/knoveleng_librispeech_english_test_asr.json|||audio|text"
# )

# =============================================================================
# Helper Functions
# =============================================================================

log_header() {
    echo ""
    echo "============================================================"
    echo "$1"
    echo "============================================================"
}

log_section() {
    echo ""
    echo "------------------------------------------------------------"
    echo "$1"
    echo "------------------------------------------------------------"
}

log_info() {
    echo "[INFO] $1"
}

log_error() {
    echo "[ERROR] $1" >&2
}

log_success() {
    echo "[✓] $1"
}

log_warning() {
    echo "[WARNING] $1"
}

# Generate safe filename from benchmark name
get_output_filename() {
    echo "$1" | tr '[:upper:]' '[:lower:]' | tr ' ' '_' | tr -cd '[:alnum:]_'
}

# Parse model configuration string
parse_model_config() {
    local config="$1"
    IFS='|' read -r EVAL_NAME MODEL_ID BATCH_SIZE EXTRA_ARGS <<< "$config"
}

# Parse benchmark configuration string
parse_benchmark() {
    local benchmark="$1"
    IFS='|' read -r BENCH_NAME DATASET_ID DATASET_CONFIG SPLIT AUDIO_COL TEXT_COL <<< "$benchmark"
}

# =============================================================================
# Evaluation Function
# =============================================================================

run_evaluation() {
    local evaluator="$1"
    local model_id="$2"
    local batch_size="$3"
    local extra_args="$4"
    local bench_name="$5"
    local dataset_id="$6"
    local dataset_config="$7"
    local split="$8"
    local audio_col="$9"
    local text_col="${10}"
    local output_file="${11}"
    
    log_section "$bench_name"
    log_info "Evaluator: $evaluator"
    log_info "Model: $model_id"
    log_info "Dataset: $dataset_id (config: ${dataset_config:-default})"
    log_info "Batch size: $batch_size | Max samples: $MAX_SAMPLES"
    
    # Build command
    local cmd="$CLI"
    cmd+=" --evaluator \"$evaluator\""
    cmd+=" --model \"$model_id\""
    cmd+=" --dataset \"$dataset_id\""
    cmd+=" --dataset-split \"$split\""
    cmd+=" --audio-column \"$audio_col\""
    cmd+=" --text-column \"$text_col\""
    if [[ "${MAX_SAMPLES,,}" != "none" ]]; then
        cmd+=" --max-samples $MAX_SAMPLES"
    fi
    cmd+=" --dataset-source \"$DATASET_SOURCE\""
    cmd+=" --batch-size $batch_size"
    cmd+=" --streaming"
    cmd+=" --output-file \"$output_file\""
    cmd+=" --lowercase"
    cmd+=" --normalize-unicode"
    cmd+=" --remove-extra-whitespace"
    cmd+=" --remove-punctuation"
    cmd+=" --save-predictions"
    
    # Add dataset config if specified
    [[ -n "$dataset_config" ]] && cmd+=" --dataset-config \"$dataset_config\""
    
    # Add extra args
    [[ -n "$extra_args" ]] && cmd+=" $extra_args"
    
    # Execute
    if eval $cmd; then
        log_success "Completed: $output_file"
        return 0
    else
        log_warning "Evaluation failed for $bench_name"
        return 1
    fi
}

# =============================================================================
# Main Execution
# =============================================================================

main() {
    log_header "ASR Model Benchmark Evaluation"
    log_info "Max samples per benchmark: $MAX_SAMPLES"
    log_info "Output directory: $OUTPUT_DIR"
    # Extract available models in the order they are defined in this script
    local available_models
    available_models=$(awk -F'"' '/^[[:space:]]*\["/ {print $2}' "$0" | paste -sd " " -)
    log_info "Available models: $available_models"
    
    # Determine which models to run
    local models_to_run=()
    if [[ -n "${RUN_MODELS:-}" ]]; then
        IFS=',' read -ra models_to_run <<< "$RUN_MODELS"
    else
        # Extract model names in the order they are defined
        readarray -t models_to_run < <(awk -F'"' '/^[[:space:]]*\["/ {print $2}' "$0")
    fi
    
    log_info "Running models: ${models_to_run[*]}"
    
    local total_success=0
    local total_failed=0
    
    # Run evaluations for each model
    for model_key in "${models_to_run[@]}"; do
        # Skip if model not found
        if [[ -z "${MODELS[$model_key]:-}" ]]; then
            log_warning "Model '$model_key' not found, skipping"
            continue
        fi
        
        # Parse model configuration
        parse_model_config "${MODELS[$model_key]}"
        
        # Create output directory for this model (use model name from MODEL_ID)
        local model_name="${MODEL_ID##*/}"  # Extract part after last /
        local model_output_dir="$OUTPUT_DIR/$model_name"
        mkdir -p "$model_output_dir"
        
        log_header "Evaluating: $model_key ($MODEL_ID)"
        
        # Run each benchmark
        for benchmark in "${BENCHMARKS[@]}"; do
            parse_benchmark "$benchmark"
            
            local output_filename=$(get_output_filename "$BENCH_NAME")
            local output_file="$model_output_dir/${output_filename}.json"
            
            if run_evaluation \
                "$EVAL_NAME" \
                "$MODEL_ID" \
                "$BATCH_SIZE" \
                "$EXTRA_ARGS" \
                "$BENCH_NAME" \
                "$DATASET_ID" \
                "$DATASET_CONFIG" \
                "$SPLIT" \
                "$AUDIO_COL" \
                "$TEXT_COL" \
                "$output_file"; then
                ((total_success++)) || true
            else
                ((total_failed++)) || true
            fi
        done
    done
    
    # Summary
    log_header "Evaluation Complete"
    log_info "Successful: $total_success"
    log_info "Failed: $total_failed"
    echo ""
    log_info "Results saved to: $OUTPUT_DIR/"
    
    # List result files
    echo ""
    find "$OUTPUT_DIR" -name "*.json" -type f 2>/dev/null | sort | while read -r file; do
        echo "  $file"
    done
    
    echo ""
    log_success "Done!"
}

# Run main function
main "$@"