speech-data-factory / scripts /run_batch_pipeline.py

Initial commit: Production-grade Voice AI speech data factory pipeline

73fed57 verified 5 months ago

12.5 kB

#!/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()

	#!/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()