scripts/run_batch_pipeline.py

#!/usr/bin/env python3
"""
run_batch_pipeline.py — Production batch runner for all 192 WAV files.

Processes all .wav files from /root/audiodown through the full Voice AI pipeline:
  1. Audio Normalization (FFmpeg EBU R128)
  2. VAD Segmentation (Silero VAD on CPU)
  3. Speaker Diarization (pyannote → ECAPA-TDNN → simple fallback)
  4. Emotion Tagging (Wav2Vec2 on GPU)
  5. ASR Transcription (Whisper medium on GPU with cuDNN disabled)
  6. Quality Validation (SNR/silence/clipping filters)
  7. Structured Export (JSONL + CSV + RTTM + manifest.json)
  8. Quality Reporting (plots + JSON summary)

Features:
- GPU-accelerated Whisper (cuDNN disabled for Conv1d compatibility)
- Checkpoint-based idempotency (safe to re-run after crash)
- Progress tracking with ETA estimates
- Detailed per-file timing logs

Usage:
    cd /workspace/Voice AI-int
    python3 scripts/run_batch_pipeline.py [--limit N] [--hf-token TOKEN]
    python3 scripts/run_batch_pipeline.py --limit 5          # Test first 5 files
    python3 scripts/run_batch_pipeline.py --skip-small       # Skip files < 10MB
"""

from __future__ import annotations

import argparse
import os
import sys
import time
import json
from pathlib import Path

# ─── Ensure we run from repo root ───────────────────────────────────────────
REPO_ROOT = Path(__file__).resolve().parent.parent
os.chdir(REPO_ROOT)
sys.path.insert(0, str(REPO_ROOT))

# ─── Config defaults ─────────────────────────────────────────────────────────
AUDIO_DIR      = Path("/root/audiodown")
CONFIG_PATH    = REPO_ROOT / "voice_pipeline" / "configs" / "pipeline_config.yaml"
LANGUAGE       = "hi"
CONTENT_TYPE   = "debate_podcast"

# ─── Import pipeline (triggers cuDNN fix in __init__.py) ─────────────────────
import voice_pipeline  # noqa: E402 — must be after path setup
from voice_pipeline.pipeline import SpeechDataPipeline
from voice_pipeline.utils.logger import setup_logging, get_logger

# ─────────────────────────────────────────────────────────────────────────────

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Batch-process all WAV files through the Voice AI Speech Data Pipeline.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("--audio-dir",   default=str(AUDIO_DIR),
                   help=f"Directory containing .wav files (default: {AUDIO_DIR})")
    p.add_argument("--config",      default=str(CONFIG_PATH),
                   help="Path to pipeline_config.yaml")
    p.add_argument("--limit",       type=int, default=None,
                   help="Process only the first N files (for testing)")
    p.add_argument("--hf-token",    default=None,
                   help="HuggingFace token for pyannote diarization (or set HF_TOKEN env var)")
    p.add_argument("--language",    default=LANGUAGE,
                   help="ISO-639-1 language code (default: hi)")
    p.add_argument("--log-level",   default="INFO",
                   choices=["DEBUG", "INFO", "WARNING", "ERROR"])
    p.add_argument("--skip-asr",    action="store_true",
                   help="Skip Whisper transcription (faster processing)")
    p.add_argument("--skip-small",  action="store_true",
                   help="Skip files smaller than 5MB (likely incomplete/corrupt)")
    p.add_argument("--min-size-mb", type=float, default=0.1,
                   help="Minimum file size in MB to process (default: 0.1 = 100KB)")
    p.add_argument("--max-size-mb", type=float, default=None,
                   help="Maximum file size in MB to process (default: None)")
    p.add_argument("--resume",      action="store_true", default=True,
                   help="Skip already-completed files (default: True)")
    p.add_argument("--no-resume",   dest="resume", action="store_false",
                   help="Re-process all files regardless of checkpoint")
    return p.parse_args()


def get_completed_file_ids(checkpoint_dir: Path) -> set[str]:
    """Read checkpoints and return set of file_ids that completed all stages."""
    # A file is "done" if its validation stage is marked done (final stage before export)
    val_ckpt = checkpoint_dir / "validation_done.jsonl"
    done_ids: set[str] = set()
    if val_ckpt.exists():
        for line in val_ckpt.open():
            try:
                rec = json.loads(line)
                if rec.get("status") == "done":
                    done_ids.add(rec["file_id"])
            except Exception:
                pass
    return done_ids


def main() -> None:
    args = parse_args()

    # ── Inject HF token ──────────────────────────────────────────────────────
    if args.hf_token:
        os.environ["HF_TOKEN"] = args.hf_token
        print(f"[INFO] HF_TOKEN set from --hf-token argument.")
    if os.environ.get("HF_TOKEN"):
        print(f"[INFO] HF_TOKEN found — pyannote diarization will be attempted.")
    else:
        print("[WARN] HF_TOKEN not set — ECAPA-TDNN clustering will be used for diarization.")

    # ── Setup logging ────────────────────────────────────────────────────────
    setup_logging(log_level=args.log_level)
    log = get_logger("batch_runner")

    # ── Print GPU info ───────────────────────────────────────────────────────
    try:
        import torch
        log.info(f"PyTorch: {torch.__version__} | "
                 f"CUDA: {torch.cuda.is_available()} | "
                 f"cuDNN: {torch.backends.cudnn.enabled}")
        if torch.cuda.is_available():
            gpu_props = torch.cuda.get_device_properties(0)
            log.info(f"GPU: {gpu_props.name} | "
                     f"VRAM: {gpu_props.total_memory / 1e9:.1f} GB")
    except Exception:
        pass

    # ── Collect WAV files ────────────────────────────────────────────────────
    audio_dir = Path(args.audio_dir)
    min_bytes  = int(args.min_size_mb * 1_000_000)

    wav_files = sorted(audio_dir.glob("*.wav"))
    wav_files = [f for f in wav_files if f.stat().st_size >= min_bytes]

    if args.max_size_mb:
        max_bytes = int(args.max_size_mb * 1_000_000)
        wav_files = [f for f in wav_files if f.stat().st_size <= max_bytes]

    if args.skip_small:
        # Skip files < 5MB (probably truncated or empty)
        wav_files = [f for f in wav_files if f.stat().st_size >= 5_000_000]

    if args.limit:
        wav_files = wav_files[:args.limit]

    total_size_gb = sum(f.stat().st_size for f in wav_files) / 1e9

    log.info(
        f"Found {len(wav_files)} WAV files "
        f"({total_size_gb:.1f} GB total) in {audio_dir}"
    )

    # ── Check existing checkpoints ───────────────────────────────────────────
    checkpoint_dir = REPO_ROOT / "data" / ".checkpoints"
    already_done: set[str] = set()

    if args.resume:
        already_done = get_completed_file_ids(checkpoint_dir)
        if already_done:
            log.info(f"Resuming: {len(already_done)} files already completed — will skip.")

    # ── Initialize Pipeline ──────────────────────────────────────────────────
    log.info(f"Initializing Voice AI Speech Pipeline from: {args.config}")
    pipeline = SpeechDataPipeline(config_path=args.config)
    log.info("Pipeline ready.")

    # ── Batch Processing ─────────────────────────────────────────────────────
    log.info("=" * 70)
    log.info(f"STARTING BATCH PROCESSING: {len(wav_files)} files")
    log.info(f"Language: {args.language} | Content: {CONTENT_TYPE}")
    log.info("=" * 70)

    start_total = time.time()
    succeeded   = 0
    failed      = 0
    skipped     = 0
    file_timings: list[dict] = []

    for i, wav_path in enumerate(wav_files):
        from voice_pipeline.utils.file_utils import get_file_hash
        file_hash = get_file_hash(wav_path, algo="md5")
        file_id = file_hash[:16]

        if args.resume and file_id in already_done:
            log.info(f"[{i+1}/{len(wav_files)}] {wav_path.name} → Skipping (already completed)")
            skipped += 1
            continue

        file_start = time.time()
        file_label = f"[{i+1}/{len(wav_files)}] {wav_path.name}"
        file_size_mb = wav_path.stat().st_size / 1e6

        # ── Estimate audio duration (heuristic: 16kHz mono WAV = ~1.92MB/min) ──
        est_duration_min = file_size_mb / 1.92

        log.info(
            f"\n{file_label} "
            f"({file_size_mb:.1f}MB ~{est_duration_min:.0f}min)"
        )

        try:
            export_dir = pipeline.run_on_file(
                file_path    = wav_path,
                language     = args.language,
                content_type = CONTENT_TYPE,
            )
            elapsed      = time.time() - file_start
            files_done   = i + 1 - skipped
            elapsed_total = time.time() - start_total
            avg_per_file  = elapsed_total / max(files_done, 1)
            remaining     = avg_per_file * (len(wav_files) - i - 1)
            eta_mins      = remaining / 60

            log.info(
                f"  ✓ Done in {elapsed:.1f}s "
                f"(RTF={elapsed/max(est_duration_min*60,1):.2f}x) | "
                f"ETA: {eta_mins:.0f}min"
            )
            succeeded += 1
            file_timings.append({
                "file": wav_path.name,
                "size_mb": round(file_size_mb, 1),
                "elapsed_s": round(elapsed, 1),
                "status": "ok"
            })

        except Exception as e:
            elapsed = time.time() - file_start
            log.error(f"  ✗ Failed after {elapsed:.1f}s: {e}", exc_info=True)
            failed += 1
            file_timings.append({
                "file": wav_path.name,
                "size_mb": round(file_size_mb, 1),
                "elapsed_s": round(elapsed, 1),
                "status": "failed",
                "error": str(e)[:200]
            })

    # ── Final Summary ─────────────────────────────────────────────────────────
    total_elapsed = time.time() - start_total
    log.info("\n" + "=" * 70)
    log.info("BATCH PROCESSING COMPLETE")
    log.info("=" * 70)
    log.info(f"  Total files processed: {len(wav_files)}")
    log.info(f"  Succeeded:             {succeeded}")
    log.info(f"  Failed:                {failed}")
    log.info(f"  Skipped (checkpointed):{skipped}")
    log.info(f"  Total wall time:       {total_elapsed/60:.1f} minutes")
    log.info(f"  Avg per file:          {total_elapsed/max(len(wav_files),1):.1f} seconds")
    log.info(f"  Exports:               {REPO_ROOT}/data/exports/")
    log.info(f"  Reports:               {REPO_ROOT}/data/reports/")

    # Write batch summary
    summary_path = REPO_ROOT / "data" / "reports" / "batch_run_summary.json"
    summary_path.parent.mkdir(parents=True, exist_ok=True)
    with open(summary_path, "w") as f:
        json.dump({
            "total_files": len(wav_files),
            "succeeded": succeeded,
            "failed": failed,
            "skipped": skipped,
            "total_elapsed_minutes": round(total_elapsed / 60, 2),
            "avg_per_file_seconds": round(total_elapsed / max(len(wav_files), 1), 1),
            "file_timings": file_timings,
        }, f, indent=2)
    log.info(f"  Batch summary:         {summary_path}")
    log.info("=" * 70)


if __name__ == "__main__":
    main()