diff --git a/pyproject.toml b/pyproject.toml index 4384069868..afd1d4385e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -442,6 +442,48 @@ fixable = ["ALL"] "INP001", # no __init__.py is required "PLE2515", # ignore \u200b complaint ] +"tutorials/audio/david_ai_redelivered_mfa/tests/**/*.py" = [ + "ANN001", # pytest fixtures and local fakes do not need public API annotations + "ANN201", # tests do not need return annotations + "ARG001", # fake callback signatures intentionally mirror production callables + "PLR2004", # expected numeric values are clearer inline in tests + "S101", # assertions are required in tests +] +"tutorials/audio/david_ai_redelivered_mfa/*/david_ai_common.py" = [ + "ANN001", # numpy audio arrays intentionally avoid a hard runtime annotation dependency + "BLE001", # file/subprocess boundaries convert arbitrary failures into pipeline errors + "PLR0913", # audio/MFA helpers expose explicit configuration rather than opaque dictionaries + "PLR2004", # audio formats and RTTM parsing use domain constants + "S603", # subprocess argument lists use resolved executables and no shell + "TRY300", # compact pipeline helpers keep success returns beside subprocess checks + "TRY301", # parse failures are translated at the narrow file boundary +] +"tutorials/audio/david_ai_redelivered_mfa/*/david_ai_manifest.py" = [ + "BLE001", # one malformed segment must be logged without aborting manifest construction + "PLR2004", # transcript timing and number normalization use domain constants +] +"tutorials/audio/david_ai_redelivered_mfa/*/david_ai_mfa_align.py" = [ + "BLE001", # alignment failures are isolated per segment/recording + "C901", # the alignment function intentionally owns one complete MFA subprocess lifecycle + "PLR0912", # explicit fallback branches document alignment failure modes + "PLR0913", # MFA options remain explicit at the process boundary + "PLR0915", # one lifecycle function guarantees cleanup and fallback consistency + "S603", # MFA is invoked with an argument list and resolved executable +] +"tutorials/audio/david_ai_redelivered_mfa/*/david_ai_ram_session.py" = [ + "BLE001", # process workers must return structured failures instead of crashing the pool + "PLR0913", # session resources and output paths remain explicit + "PLR2004", # PCM and masking parameters are domain constants + "PLW0603", # the process-local MFA model cache is intentionally initialized once per worker +] +"tutorials/audio/david_ai_redelivered_mfa/*/stage_ram_session_pipeline.py" = [ + "C901", # CLI validation, sharding, resume filtering, and worker orchestration are kept together + "PLR0915", # the CLI entry point owns the complete process-pool lifecycle +] +"tutorials/audio/david_ai_redelivered_mfa/opus/convert_opus_to_wav.py" = [ + "PLR0913", # conversion controls are explicit CLI options + "S603", # ffmpeg is invoked with an argument list and resolved executable +] "fern/**/*.py" = [ "INP001", # Fern CLI helper scripts; not an installable package ] diff --git a/tutorials/audio/david_ai_redelivered_mfa/.gitignore b/tutorials/audio/david_ai_redelivered_mfa/.gitignore new file mode 100644 index 0000000000..269aa30600 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/.gitignore @@ -0,0 +1,33 @@ +# Raw or generated audio/data +*.wav +*.opus +*.rttm +*.TextGrid +*.json +*.jsonl +*.csv +*.tsv +*.tar +*.tar.gz +*.tar.zst +*.zst + +# Pipeline outputs and scratch +.done/ +audio_16k/ +audio_16k_masked/ +audio_mixed/ +logs/ +textgrids/ +workdir*/ + +# Credentials and local environments +.env +.env.* +*.key +*.pem + +# Python/tool caches +__pycache__/ +.pytest_cache/ +.ruff_cache/ diff --git a/tutorials/audio/david_ai_redelivered_mfa/README.md b/tutorials/audio/david_ai_redelivered_mfa/README.md new file mode 100644 index 0000000000..9c507652c4 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/README.md @@ -0,0 +1,368 @@ +# David AI redelivered MFA tutorial + +This tutorial provides two independent, on-the-fly end-to-end pipelines for +multi-speaker David AI sessions: + +- [`opus/`](opus/README.md): writes per-speaker and mixed-session Opus audio. +- [`wav/`](wav/README.md): writes masked per-speaker and mixed-session mono + 16 kHz PCM WAV audio. +- [`parakeet_wer/`](parakeet_wer/README.md): runs segment-level Parakeet ASR, + computes WER distributions, filters high-WER segments, and writes aligned + per-speaker manifests. + +Both variants: + +1. Read raw source WAVs and `machine_generated_transcript.json`. +2. Normalize transcript text in memory. +3. Run Montreal Forced Aligner (MFA) with the `english_us_arpa` acoustic model, + base dictionary, and runtime G2P for OOV words. +4. Generate session RTTM. +5. Preserve speech inside the original manifest boundaries ±0.5 seconds. +6. Replace pauses outside those boundaries with white noise at amplitude + `0.0002`, using a 5 ms boundary crossfade. +7. Mix all speakers on the common session timeline. +8. Write ordinary TextGrids (MFA + fallback) and FastMSS TextGrids (MFA only). +9. Validate all required outputs before writing + `.done/sessions/.done`. + +A session with zero MFA-aligned segments is treated as failed even when manifest +fallback intervals exist, preventing infrastructure/model failures from being +marked complete. + +The pipelines do not reuse persisted normalized manifests, shared lexicons, or +partial RTTM/TextGrid/audio outputs. A validated session done flag is the only +resume state: sessions with a done flag are skipped, while unfinished sessions +are rebuilt from raw inputs. + +## Data privacy + +No David AI dataset content is included in this tutorial. All paths, IDs, and +transcript values in the documentation and tests are synthetic placeholders. + +Do not commit or publish raw manifests, audio, RTTM, TextGrid, logs, archives, +or generated outputs. The tutorial `.gitignore` excludes these artifact types, +but generated changes must still be reviewed before committing. + +## Requirements + +### Supported environment + +- Linux x86_64 +- Python 3.11–3.13 (Python 3.12 is recommended and tested) +- ffmpeg and ffprobe +- Montreal Forced Aligner 3.3.9 +- NeMo Curator from this repository +- Approximately 1–2 GB of node-local scratch per concurrent session + +GPU access is not required for this pipeline. + +### Required MFA models + +The following English US ARPA models must be installed: + +- Dictionary: `english_us_arpa` +- Acoustic model: `english_us_arpa` +- G2P model: `english_us_arpa` + +The pipeline expects this layout under `$MFA_ROOT_DIR`: + +```text +MFA_ROOT_DIR/ +└── pretrained_models/ + ├── acoustic/english_us_arpa.zip + ├── dictionary/english_us_arpa.dict + └── g2p/english_us_arpa.zip +``` + +## Environment setup + +From the repository root: + +```bash +cd ~/Curator_my_fork + +conda env create \ + -f tutorials/audio/david_ai_redelivered_mfa/environment.yml + +conda activate david-ai-mfa + +python -m pip install -e . +python -m pip install \ + -r tutorials/audio/david_ai_redelivered_mfa/requirements.txt +``` + +Download the MFA models: + +```bash +export MFA_ROOT_DIR="$HOME/MFA_models" +mkdir -p "$MFA_ROOT_DIR" + +mfa model download dictionary english_us_arpa +mfa model download acoustic english_us_arpa +mfa model download g2p english_us_arpa +``` + +Confirm the tools: + +```bash +python --version +mfa version +ffmpeg -version +ffprobe -version +``` + +For development and tests: + +```bash +python -m pip install \ + -r tutorials/audio/david_ai_redelivered_mfa/requirements-dev.txt +``` + +## Input layout + +`DATA_ROOT` must contain one directory per session: + +```text +/ +└── / + ├── machine_generated_transcript.json + ├── _postprocess.wav + ├── .wav + └── ... +``` + +The transcript JSON must contain a `transcript` list: + +```json +{ + "transcript": [ + { + "text": "Example utterance.", + "start": 1.25, + "end": 2.85, + "speaker": "" + } + ] +} +``` + +Segment times use the shared session timeline. Speaker IDs must match the source +WAV filename prefixes. + +For each speaker, the pipeline selects the first existing WAV in this order: + +1. `_postprocess.wav` +2. `_postprocessed.wav` +3. `.wav` +4. `_preprocessed.wav` + +The session fails explicitly if none of these files exists for a speaker named +in the transcript. + +## Choose an output format + +### Opus + +```bash +cd ~/Curator_my_fork/tutorials/audio/david_ai_redelivered_mfa/opus + +DATA_ROOT=/path/to/raw/sessions \ +WORK_DIR=/path/to/opus-output \ +MFA_ROOT_DIR="$HOME/MFA_models" \ +MFA_ENV="$CONDA_PREFIX" \ +WORKERS=16 \ +MFA_NUM_JOBS=2 \ +SEG_EXTRACT_WORKERS=8 \ +bash run_david_ai_mfa_ram_session.sh +``` + +See [`opus/README.md`](opus/README.md) for the complete output layout. + +### WAV + +```bash +cd ~/Curator_my_fork/tutorials/audio/david_ai_redelivered_mfa/wav + +DATA_ROOT=/path/to/raw/sessions \ +WORK_DIR=/path/to/wav-output \ +MFA_ROOT_DIR="$HOME/MFA_models" \ +MFA_ENV="$CONDA_PREFIX" \ +WORKERS=16 \ +MFA_NUM_JOBS=2 \ +SEG_EXTRACT_WORKERS=8 \ +bash run_david_ai_mfa_ram_session.sh +``` + +The WAV variant writes: + +```text +/ +├── audio_16k_masked/ +│ ├── __postprocessed.wav +│ └── __postprocessed.rttm +├── audio_mixed/ +│ ├── .wav +│ └── .rttm +├── textgrids/ +├── logs/ +└── .done/sessions/ +``` + +All WAV outputs are mono, 16 kHz, signed 16-bit PCM. + +## Multi-node cluster run + +Cluster submission scripts are kept separately in [`cluster/`](cluster/README.md). +They support both output variants through `VARIANT=opus` or `VARIANT=wav`. + +Example: + +```bash +VARIANT=wav \ +DATA_ROOT=/shared/data/david_ai_sessions \ +WORK_DIR=/shared/output/david_ai_wav \ +MFA_ENV=/shared/envs/david-ai-mfa \ +MFA_ROOT_DIR=/shared/models/MFA_models \ +NUM_NODES=8 \ +CPUS_PER_NODE=64 \ +WORKERS_PER_NODE=16 \ +MFA_NUM_JOBS=2 \ +SLURM_ACCOUNT=my-account \ +SLURM_PARTITION=cpu \ +bash cluster/run_multinode.sh +``` + +The launcher uses one SLURM array task per node and exports only explicitly +required variables. It does not copy data to or from the cluster. + +## Runtime configuration + +| Variable | Default | Purpose | +|---|---:|---| +| `DATA_ROOT` | variant-specific example path | Raw session root | +| `WORK_DIR` | local variant workdir | Persistent output root | +| `SESSIONS_FILE` | unset | Optional absolute session-ID list | +| `MFA_ROOT_DIR` | `~/MFA_models` | MFA model root | +| `MFA_ENV` | `~/miniconda3/envs/curator_pain_1` | Environment containing MFA | +| `WORKERS` | `4` | Concurrent sessions | +| `MFA_NUM_JOBS` | `2` | MFA jobs per speaker recording | +| `SEG_EXTRACT_WORKERS` | `8` | Parallel ffmpeg segment extraction | +| `MIX_PREP_WORKERS` | number of session speakers | Parallel speaker masking | +| `RAM_DIR` | unique `/tmp` directory | Node-local ephemeral scratch | +| `FFMPEG_BIN` | `ffmpeg` from `PATH` | Explicit ffmpeg executable | +| `FFMPEG_TIMEOUT_S` | `600` | Per-ffmpeg timeout | + +Approximate peak MFA parallelism is `WORKERS × MFA_NUM_JOBS`. Keep that value +near or below the available CPU count when other stages need CPU concurrently. + +Scratch must be node-local (`/tmp` or equivalent), not a shared network +filesystem. + +## Run a session subset + +Create a text file containing one session directory name per line. Empty lines +and lines beginning with `#` are ignored: + +```text +# validation subset +session-a +session-b +``` + +Pass it to either local variant: + +```bash +SESSIONS_FILE=/absolute/path/to/sessions.txt \ +DATA_ROOT=/path/to/raw/sessions \ +WORK_DIR=/path/to/output \ +bash run_david_ai_mfa_ram_session.sh +``` + +For a cluster run, pass the same shared absolute `SESSIONS_FILE` path to +`cluster/run_multinode.sh`. The subset is applied before deterministic sharding +and done-flag filtering. + +## Success and restart behavior + +A done flag is written only after all expected outputs for the current session +exist and are non-empty: + +```text +/.done/sessions/.done +``` + +Done flags are the parallel resume authority. Starting the same local or +multi-node command again skips sessions that already have a validated flag and +processes only sessions without one. + +To intentionally rebuild a completed session, remove only its flag: + +```bash +rm /.done/sessions/.done +``` + +## Tests + +Tests are separated by output variant: + +```bash +cd ~/Curator_my_fork/tutorials/audio/david_ai_redelivered_mfa + +pytest tests/opus +pytest tests/wav +pytest tests/cluster +pytest tests/parakeet_wer +``` + +Run lint checks: + +```bash +ruff check opus wav tests +``` + +## Troubleshooting + +### `mfa` not found + +Activate the environment and confirm: + +```bash +conda activate david-ai-mfa +which mfa +mfa version +``` + +### MFA model not found + +Verify `$MFA_ROOT_DIR/pretrained_models/` matches the required model layout +above. + +### ffmpeg not found + +Install ffmpeg in the conda environment or set an explicit executable: + +```bash +export FFMPEG_BIN=/path/to/ffmpeg +``` + +### A session has no done flag + +Search its ID in `/logs/run_e2e_*.log`. A missing flag means the +session failed or at least one required output was missing/empty. + +### Unexpectedly high memory or CPU use + +Reduce `WORKERS`, `MFA_NUM_JOBS`, or `SEG_EXTRACT_WORKERS`. Start with: + +```bash +WORKERS=4 MFA_NUM_JOBS=2 SEG_EXTRACT_WORKERS=4 +``` + +## Security and cluster use + +The local pipeline requires no outbound network access after dependencies and +MFA models are installed. Do not pass credentials, `.env` files, or unnecessary +shell environment variables to the pipeline. + +Cluster runners should receive only explicitly required variables. This +tutorial does not require or perform data-copy operations to a cluster. diff --git a/tutorials/audio/david_ai_redelivered_mfa/cluster/README.md b/tutorials/audio/david_ai_redelivered_mfa/cluster/README.md new file mode 100644 index 0000000000..bd85570a5b --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/cluster/README.md @@ -0,0 +1,157 @@ +# Multi-node SLURM execution + +The cluster launcher uses a SLURM job array with one task per node. Session +discovery is deterministic, and task `i` processes sessions whose sorted index +satisfies: + +```text +session_index % NUM_NODES == i +``` + +Each node runs multiple session processes (`WORKERS_PER_NODE`), while each MFA +recording may use multiple MFA jobs (`MFA_NUM_JOBS`). + +## Requirements + +- `DATA_ROOT`, `WORK_DIR`, the repository, MFA environment, and MFA models must + already be visible at identical absolute paths on every node. +- `$MFA_ENV/bin/python`, `$MFA_ENV/bin/mfa`, and ffmpeg must be executable on + compute nodes. +- `WORK_DIR` must be writable by all array tasks. +- Scratch uses `$SLURM_TMPDIR` when available, otherwise `/tmp`. +- No input/output copying is performed by these scripts. +- Optional containers must mount every required shared path explicitly. + +## Submit + +From the tutorial root: + +```bash +cd ~/Curator_my_fork/tutorials/audio/david_ai_redelivered_mfa + +VARIANT=wav \ +DATA_ROOT=/shared/data/david_ai_sessions \ +WORK_DIR=/shared/output/david_ai_wav \ +MFA_ENV=/shared/envs/david-ai-mfa \ +MFA_ROOT_DIR=/shared/models/MFA_models \ +NUM_NODES=8 \ +MAX_CONCURRENT_NODES=8 \ +CPUS_PER_NODE=64 \ +WORKERS_PER_NODE=16 \ +MFA_NUM_JOBS=2 \ +SEG_EXTRACT_WORKERS=8 \ +SLURM_ACCOUNT=my-account \ +SLURM_PARTITION=cpu \ +TIME_LIMIT=08:00:00 \ +bash cluster/run_multinode.sh +``` + +Use `VARIANT=opus` for the Opus pipeline. + +To process a subset, add a shared absolute list path: + +```bash +SESSIONS_FILE=/shared/config/session_subset.txt \ +... \ +bash cluster/run_multinode.sh +``` + +## Optional container + +```bash +CONTAINER_IMAGE=/shared/containers/pipeline.sqsh \ +CONTAINER_MOUNTS=/shared:/shared \ +... \ +bash cluster/run_multinode.sh +``` + +The image must have access to the submitted repository path, data, outputs, MFA +environment, models, and node-local scratch. + +## Parallelism + +Approximate MFA process slots per node: + +```text +WORKERS_PER_NODE × MFA_NUM_JOBS +``` + +Start with that product at or below `CPUS_PER_NODE`. Leave CPU headroom for +ffmpeg extraction, pause masking, mixing, and filesystem work. + +Example for a 64-CPU node: + +```text +WORKERS_PER_NODE=16 +MFA_NUM_JOBS=2 +``` + +This creates up to 32 MFA slots and leaves headroom for other stages. + +## MFA directory isolation + +`MFA_ROOT_DIR` is treated as a read-only source for pretrained models. Runtime +state is never written there. + +Isolation hierarchy: + +```text +/ +└── david_ai___/ + ├── model_source/ # one shard-local copy from shared storage + └── mfa_workers/ + └── worker_/ + ├── models/ # private worker model copies + ├── mfa_root/ # private MFA config/database root + └── align_temp/ + └── / +``` + +- Every array shard has a unique scratch root based on job and task IDs. +- Every shard stages the shared dictionary, acoustic, and G2P source once into + its own node-local `model_source`. +- Every session worker process has a private MFA root and private model copies. +- Every worker pre-extracts the G2P archive and validates that it contains one + non-empty `model.fst` before launching MFA, avoiding concurrent MFA extraction. +- A worker processes its speaker recordings sequentially. +- Every session uses a separate alignment temp directory. +- MFA subprocesses receive the private worker root through `MFA_ROOT_DIR`. + +Consequently, nodes and concurrent session workers do not share writable MFA +model, database, or temporary directories. + +## Logs and status + +SLURM logs: + +```text +/logs/slurm/__.out +/logs/slurm/__.err +``` + +Per-shard pipeline logs include the SLURM job and array task IDs, preventing +multiple nodes from writing the same log filename. + +Monitor: + +```bash +squeue -j +sacct -j --format=JobID,State,Elapsed,ExitCode +``` + +Successful sessions write: + +```text +/.done/sessions/.done +``` + +If an array task fails or times out, submit the same command again. Every node +recomputes its deterministic shard, skips sessions with done flags, and runs only +unfinished sessions. Changing `NUM_NODES` is also safe because done flags are +checked after the new deterministic sharding assignment. + +## Security + +`run_multinode.sh` exports only explicitly required, non-secret variables. It +does not pass the submit shell's full environment. Do not place credentials in +these variables, command arguments, logs, or container mounts. diff --git a/tutorials/audio/david_ai_redelivered_mfa/cluster/run_multinode.sh b/tutorials/audio/david_ai_redelivered_mfa/cluster/run_multinode.sh new file mode 100755 index 0000000000..a6004c8e34 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/cluster/run_multinode.sh @@ -0,0 +1,122 @@ +#!/bin/bash +# Submit the David AI E2E pipeline as one SLURM array task per node. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TUTORIAL_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" + +VARIANT="${VARIANT:-wav}" +DATA_ROOT="${DATA_ROOT:?Set DATA_ROOT to the shared raw-session directory}" +WORK_DIR="${WORK_DIR:?Set WORK_DIR to the shared output directory}" +MFA_ENV="${MFA_ENV:?Set MFA_ENV to an environment visible on compute nodes}" +MFA_ROOT_DIR="${MFA_ROOT_DIR:-$HOME/MFA_models}" +SESSIONS_FILE="${SESSIONS_FILE:-}" + +NUM_NODES="${NUM_NODES:-1}" +MAX_CONCURRENT_NODES="${MAX_CONCURRENT_NODES:-$NUM_NODES}" +CPUS_PER_NODE="${CPUS_PER_NODE:-64}" +WORKERS_PER_NODE="${WORKERS_PER_NODE:-16}" +MFA_NUM_JOBS="${MFA_NUM_JOBS:-2}" +SEG_EXTRACT_WORKERS="${SEG_EXTRACT_WORKERS:-8}" +MIX_PREP_WORKERS="${MIX_PREP_WORKERS:-4}" +FFMPEG_TIMEOUT_S="${FFMPEG_TIMEOUT_S:-2400}" + +SLURM_ACCOUNT="${SLURM_ACCOUNT:-}" +SLURM_PARTITION="${SLURM_PARTITION:-}" +TIME_LIMIT="${TIME_LIMIT:-08:00:00}" +MEMORY_PER_NODE="${MEMORY_PER_NODE:-}" +JOB_NAME="${JOB_NAME:-david-ai-${VARIANT}}" +LOG_DIR="${LOG_DIR:-$WORK_DIR/logs/slurm}" + +CONTAINER_IMAGE="${CONTAINER_IMAGE:-}" +CONTAINER_MOUNTS="${CONTAINER_MOUNTS:-}" +FFMPEG_BIN="${FFMPEG_BIN:-}" +CONTAINER_MOUNTS_B64="" +if [[ -n "$CONTAINER_MOUNTS" ]]; then + CONTAINER_MOUNTS_B64="$(printf "%s" "$CONTAINER_MOUNTS" | base64 | tr -d "\n")" +fi + +case "$VARIANT" in + opus | wav) ;; + *) + echo "ERROR: VARIANT must be 'opus' or 'wav', got: $VARIANT" >&2 + exit 2 + ;; +esac +for value in "$NUM_NODES" "$MAX_CONCURRENT_NODES" "$CPUS_PER_NODE" \ + "$WORKERS_PER_NODE" "$MFA_NUM_JOBS"; do + if [[ ! "$value" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: node/thread counts must be positive integers" >&2 + exit 2 + fi +done +for path in "$DATA_ROOT" "$WORK_DIR" "$MFA_ENV" "$MFA_ROOT_DIR" "$TUTORIAL_ROOT"; do + if [[ "$path" != /* ]]; then + echo "ERROR: cluster paths must be absolute: $path" >&2 + exit 2 + fi + if [[ "$path" == *,* ]]; then + echo "ERROR: cluster paths cannot contain commas: $path" >&2 + exit 2 + fi +done +if [[ -n "$SESSIONS_FILE" ]]; then + if [[ "$SESSIONS_FILE" != /* || "$SESSIONS_FILE" == *,* ]]; then + echo "ERROR: SESSIONS_FILE must be an absolute path without commas" >&2 + exit 2 + fi + if [[ ! -f "$SESSIONS_FILE" ]]; then + echo "ERROR: SESSIONS_FILE does not exist: $SESSIONS_FILE" >&2 + exit 1 + fi +fi +if ! command -v sbatch >/dev/null 2>&1; then + echo "ERROR: sbatch is not available" >&2 + exit 1 +fi +if [[ ! -d "$DATA_ROOT" ]]; then + echo "ERROR: DATA_ROOT does not exist: $DATA_ROOT" >&2 + exit 1 +fi +if [[ ! -x "$MFA_ENV/bin/python" || ! -x "$MFA_ENV/bin/mfa" ]]; then + echo "ERROR: MFA_ENV must contain executable python and mfa: $MFA_ENV" >&2 + exit 1 +fi + +mkdir -p "$LOG_DIR" "$WORK_DIR" + +# Export only task-required values. Do not propagate the submit shell environment. +EXPORTS="VARIANT=$VARIANT" +EXPORTS+=",DATA_ROOT=$DATA_ROOT,WORK_DIR=$WORK_DIR" +EXPORTS+=",MFA_ENV=$MFA_ENV,MFA_ROOT_DIR=$MFA_ROOT_DIR" +EXPORTS+=",TUTORIAL_ROOT=$TUTORIAL_ROOT" +EXPORTS+=",SHARD_COUNT=$NUM_NODES" +EXPORTS+=",WORKERS_PER_NODE=$WORKERS_PER_NODE,MFA_NUM_JOBS=$MFA_NUM_JOBS" +EXPORTS+=",SEG_EXTRACT_WORKERS=$SEG_EXTRACT_WORKERS,MIX_PREP_WORKERS=$MIX_PREP_WORKERS" +EXPORTS+=",FFMPEG_TIMEOUT_S=$FFMPEG_TIMEOUT_S" +EXPORTS+=",CONTAINER_IMAGE=$CONTAINER_IMAGE,CONTAINER_MOUNTS_B64=$CONTAINER_MOUNTS_B64" +EXPORTS+=",FFMPEG_BIN=$FFMPEG_BIN" +[[ -n "$SESSIONS_FILE" ]] && EXPORTS+=",SESSIONS_FILE=$SESSIONS_FILE" + +SBATCH_ARGS=( + --job-name "$JOB_NAME" + --nodes 1 + --ntasks 1 + --cpus-per-task "$CPUS_PER_NODE" + --time "$TIME_LIMIT" + --array "0-$((NUM_NODES - 1))%$MAX_CONCURRENT_NODES" + --output "$LOG_DIR/${JOB_NAME}_%A_%a.out" + --error "$LOG_DIR/${JOB_NAME}_%A_%a.err" + --export "$EXPORTS" +) +[[ -n "$SLURM_ACCOUNT" ]] && SBATCH_ARGS+=(--account "$SLURM_ACCOUNT") +[[ -n "$SLURM_PARTITION" ]] && SBATCH_ARGS+=(--partition "$SLURM_PARTITION") +[[ -n "$MEMORY_PER_NODE" ]] && SBATCH_ARGS+=(--mem "$MEMORY_PER_NODE") + +echo "Submitting $NUM_NODES-node $VARIANT pipeline" +echo "Per node: CPUs=$CPUS_PER_NODE session_workers=$WORKERS_PER_NODE MFA_jobs=$MFA_NUM_JOBS" +echo "Input: $DATA_ROOT" +echo "Output: $WORK_DIR" + +sbatch "${SBATCH_ARGS[@]}" "$SCRIPT_DIR/run_node.sh" diff --git a/tutorials/audio/david_ai_redelivered_mfa/cluster/run_node.sh b/tutorials/audio/david_ai_redelivered_mfa/cluster/run_node.sh new file mode 100755 index 0000000000..ba5f956fb8 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/cluster/run_node.sh @@ -0,0 +1,145 @@ +#!/bin/bash +# Execute one deterministic session shard on one SLURM node. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +: "${VARIANT:?Missing VARIANT}" +: "${DATA_ROOT:?Missing DATA_ROOT}" +: "${WORK_DIR:?Missing WORK_DIR}" +: "${MFA_ENV:?Missing MFA_ENV}" +: "${MFA_ROOT_DIR:?Missing MFA_ROOT_DIR}" +: "${TUTORIAL_ROOT:?Missing TUTORIAL_ROOT}" +: "${SHARD_COUNT:?Missing SHARD_COUNT}" +: "${SLURM_ARRAY_TASK_ID:?This script must run as a SLURM array task}" + +SHARD_INDEX="$SLURM_ARRAY_TASK_ID" +WORKERS_PER_NODE="${WORKERS_PER_NODE:-16}" +MFA_NUM_JOBS="${MFA_NUM_JOBS:-2}" +SEG_EXTRACT_WORKERS="${SEG_EXTRACT_WORKERS:-8}" +MIX_PREP_WORKERS="${MIX_PREP_WORKERS:-4}" +FFMPEG_TIMEOUT_S="${FFMPEG_TIMEOUT_S:-2400}" +SESSIONS_FILE="${SESSIONS_FILE:-}" +CONTAINER_IMAGE="${CONTAINER_IMAGE:-}" +CONTAINER_MOUNTS_B64="${CONTAINER_MOUNTS_B64:-}" +CONTAINER_MOUNTS="" +if [[ -n "$CONTAINER_MOUNTS_B64" ]]; then + CONTAINER_MOUNTS="$(printf "%s" "$CONTAINER_MOUNTS_B64" | base64 --decode)" +fi +FFMPEG_BIN="${FFMPEG_BIN:-}" + +if [[ -n "$CONTAINER_IMAGE" && "${IN_CONTAINER:-0}" != "1" ]]; then + SRUN_ARGS=( + --nodes 1 + --ntasks 1 + --container-image "$CONTAINER_IMAGE" + ) + [[ -n "$CONTAINER_MOUNTS" ]] && + SRUN_ARGS+=(--container-mounts "$CONTAINER_MOUNTS") + exec srun "${SRUN_ARGS[@]}" \ + env IN_CONTAINER=1 bash "$TUTORIAL_ROOT/cluster/run_node.sh" +fi + +PIPELINE_DIR="$TUTORIAL_ROOT/$VARIANT" +RUNNER="$PIPELINE_DIR/run_david_ai_mfa_ram_session.sh" +if [[ ! -x "$RUNNER" ]]; then + echo "ERROR: variant runner is missing or not executable: $RUNNER" >&2 + exit 1 +fi +if [[ ! -x "$MFA_ENV/bin/python" || ! -x "$MFA_ENV/bin/mfa" ]]; then + echo "ERROR: compute node cannot access MFA_ENV=$MFA_ENV" >&2 + exit 1 +fi + +export PATH="$MFA_ENV/bin:/usr/local/bin:/usr/bin:/bin" +export PYTHON="$MFA_ENV/bin/python" +export PYTHONPATH="$(cd "$TUTORIAL_ROOT/../../.." && pwd)" +export MFA_ROOT_DIR +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export NUMEXPR_MAX_THREADS=1 + +if [[ -n "$FFMPEG_BIN" ]]; then + if [[ ! -x "$FFMPEG_BIN" ]]; then + echo "ERROR: FFMPEG_BIN is not executable: $FFMPEG_BIN" >&2 + exit 1 + fi + export FFMPEG_BIN + export PATH="$(dirname "$FFMPEG_BIN"):$PATH" +elif [[ ! -x "$MFA_ENV/bin/ffmpeg" ]]; then + echo "ERROR: ffmpeg is missing from MFA_ENV and FFMPEG_BIN is unset" >&2 + exit 1 +fi + +SCRATCH_ROOT="${SLURM_TMPDIR:-/tmp}" +RAM_DIR="$SCRATCH_ROOT/david_ai_${VARIANT}_${SLURM_JOB_ID}_${SHARD_INDEX}" +rm -rf "$RAM_DIR" +mkdir -p "$RAM_DIR" +cleanup() { + rm -rf "$RAM_DIR" +} +trap cleanup EXIT + +SHARED_MFA_ROOT_DIR="$MFA_ROOT_DIR" +NODE_MFA_ROOT_DIR="$RAM_DIR/model_source" +mkdir -p \ + "$NODE_MFA_ROOT_DIR/pretrained_models/dictionary" \ + "$NODE_MFA_ROOT_DIR/pretrained_models/acoustic" \ + "$NODE_MFA_ROOT_DIR/pretrained_models/g2p" + +stage_model() { + local destination_dir="$1" + shift + local candidate + for candidate in "$@"; do + if [[ -e "$candidate" ]]; then + cp -a "$candidate" "$destination_dir/" + return 0 + fi + done + echo "ERROR: no model candidate found: $*" >&2 + return 1 +} + +stage_model \ + "$NODE_MFA_ROOT_DIR/pretrained_models/dictionary" \ + "$SHARED_MFA_ROOT_DIR/pretrained_models/dictionary/english_us_arpa.dict" \ + "$SHARED_MFA_ROOT_DIR/pretrained_models/dictionary/english_us_arpa.txt" +stage_model \ + "$NODE_MFA_ROOT_DIR/pretrained_models/acoustic" \ + "$SHARED_MFA_ROOT_DIR/pretrained_models/acoustic/english_us_arpa.zip" \ + "$SHARED_MFA_ROOT_DIR/pretrained_models/acoustic/english_us_arpa" +stage_model \ + "$NODE_MFA_ROOT_DIR/pretrained_models/g2p" \ + "$SHARED_MFA_ROOT_DIR/pretrained_models/g2p/english_us_arpa.zip" \ + "$SHARED_MFA_ROOT_DIR/pretrained_models/g2p/english_us_arpa" \ + "$SHARED_MFA_ROOT_DIR/extracted_models/g2p/english_us_arpa_g2p" + +export MFA_ROOT_DIR="$NODE_MFA_ROOT_DIR" + +echo "[$(date -Is)] Node=$(hostname) job=$SLURM_JOB_ID shard=$SHARD_INDEX/$SHARD_COUNT" +echo "Variant=$VARIANT workers=$WORKERS_PER_NODE MFA_jobs=$MFA_NUM_JOBS" +echo "Scratch=$RAM_DIR" +echo "Shared model source=$SHARED_MFA_ROOT_DIR" +echo "Shard-local model source=$MFA_ROOT_DIR" + +env \ + DATA_ROOT="$DATA_ROOT" \ + WORK_DIR="$WORK_DIR" \ + MFA_ENV="$MFA_ENV" \ + MFA_ROOT_DIR="$MFA_ROOT_DIR" \ + PYTHON="$PYTHON" \ + WORKERS="$WORKERS_PER_NODE" \ + MFA_NUM_JOBS="$MFA_NUM_JOBS" \ + SEG_EXTRACT_WORKERS="$SEG_EXTRACT_WORKERS" \ + MIX_PREP_WORKERS="$MIX_PREP_WORKERS" \ + FFMPEG_TIMEOUT_S="$FFMPEG_TIMEOUT_S" \ + RAM_DIR="$RAM_DIR" \ + SHARD_COUNT="$SHARD_COUNT" \ + SHARD_INDEX="$SHARD_INDEX" \ + SESSIONS_FILE="$SESSIONS_FILE" \ + bash "$RUNNER" + +echo "[$(date -Is)] Shard $SHARD_INDEX completed" diff --git a/tutorials/audio/david_ai_redelivered_mfa/environment.yml b/tutorials/audio/david_ai_redelivered_mfa/environment.yml new file mode 100644 index 0000000000..a6023a8895 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/environment.yml @@ -0,0 +1,12 @@ +name: david-ai-mfa +channels: + - conda-forge +dependencies: + - python=3.12 + - ffmpeg + - montreal-forced-aligner=3.3.9 + - pip + - pip: + - numpy + - num2words>=0.5.12 + - textgrid diff --git a/tutorials/audio/david_ai_redelivered_mfa/opus/README.md b/tutorials/audio/david_ai_redelivered_mfa/opus/README.md new file mode 100644 index 0000000000..3364663e0f --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/opus/README.md @@ -0,0 +1,83 @@ +# David AI on-the-fly MFA E2E + +This directory contains the Opus pipeline. Every unfinished session is rebuilt +from its raw inputs: + +1. Read `machine_generated_transcript.json` and resolve one supported WAV per speaker + (`_postprocess`, `_postprocessed`, ordinary `.wav`, then `_preprocessed`). +2. Normalize transcript rows in memory; no normalized manifests are saved or read. +3. Align with the base `english_us_arpa` dictionary and runtime MFA G2P for OOV words. +4. Save session RTTM. +5. Replace pauses using original manifest boundaries protected by ±0.5 seconds. + Noise amplitude is `0.0002`; boundary smoothing is 5 ms. +6. Save every masked per-speaker Opus and the mixed session Opus. +7. Save ordinary and FastMSS TextGrids at session and recording level. +8. Validate every required output and write `.done/sessions/.done`. + +The pipeline never reads shared lexicons, persisted manifests, cached RTTM, +previous alignment results, or partial mixed audio. Validated done flags are +used for resume: completed sessions are skipped and sessions without flags are +processed from raw inputs. + +## Entrypoints + +Local: + +```bash +DATA_ROOT=/path/to/raw/sessions \ +WORK_DIR=/path/to/output \ +WORKERS=16 \ +MFA_NUM_JOBS=2 \ +bash run_david_ai_mfa_ram_session.sh +``` + +SLURM multi-node runs use the shared cluster launcher: + +```bash +cd .. + +VARIANT=opus \ +DATA_ROOT=/shared/path/to/raw/sessions \ +WORK_DIR=/shared/path/to/output \ +NUM_NODES=8 \ +WORKERS_PER_NODE=16 \ +MFA_NUM_JOBS=2 \ +bash cluster/run_multinode.sh +``` + +See `../cluster/README.md` for environment and scheduler options. + +## Outputs + +```text +/ +├── audio_mixed/ +│ ├── .opus +│ ├── .rttm +│ └── speakers/ +│ └── __postprocessed.opus +├── textgrids/ +│ ├── .TextGrid +│ ├── _fastmss.TextGrid +│ ├── .TextGrid +│ └── _fastmss.TextGrid +├── .done/ +│ └── sessions/ +│ └── .done +└── logs/ +``` + +## Runtime controls + +- `DATA_ROOT`, `WORK_DIR`: raw input and output roots. +- `WORKERS`: concurrent sessions. +- `MFA_NUM_JOBS`: MFA jobs per recording. +- `SEG_EXTRACT_WORKERS`: concurrent segment extraction per recording. +- `MIX_PREP_WORKERS`: concurrent speaker preparation per session. +- `RAM_DIR`: ephemeral node-local scratch. +- `FFMPEG_BIN`, `FFMPEG_TIMEOUT_S`: ffmpeg runtime controls. +- `MFA_ROOT_DIR`, `MFA_ENV`: MFA model and environment locations. +- `RAM_ARRAY_COUNT`: cluster scheduling shards. + +There are deliberately no `FORCE`, `SKIP_LEXICON`, session-subset, persisted-manifest, +resume, stage, or cache flags. diff --git a/tutorials/audio/david_ai_redelivered_mfa/opus/convert_opus_to_wav.py b/tutorials/audio/david_ai_redelivered_mfa/opus/convert_opus_to_wav.py new file mode 100755 index 0000000000..73e2011a8c --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/opus/convert_opus_to_wav.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Convert an Opus directory tree to mono PCM WAV while preserving relative paths.""" + +from __future__ import annotations + +import argparse +import os +import shutil +import subprocess +import threading +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from pathlib import Path + + +@dataclass +class ConversionResult: + source: Path + destination: Path + status: str + error: str = "" + + +def destination_for(source: Path, input_dir: Path, output_dir: Path) -> Path: + return (output_dir / source.relative_to(input_dir)).with_suffix(".wav") + + +def convert_one( + source: Path, + destination: Path, + *, + ffmpeg: str, + sample_rate: int, + channels: int, + overwrite: bool, +) -> ConversionResult: + if destination.is_file() and not overwrite: + return ConversionResult(source, destination, "skipped") + + destination.parent.mkdir(parents=True, exist_ok=True) + temp = destination.with_name( + f".{destination.stem}.{os.getpid()}.{threading.get_ident()}.tmp.wav" + ) + command = [ + ffmpeg, + "-nostdin", + "-y", + "-i", + str(source), + "-ar", + str(sample_rate), + "-ac", + str(channels), + "-c:a", + "pcm_s16le", + str(temp), + ] + try: + result = subprocess.run(command, capture_output=True, check=False, text=True) + if result.returncode != 0: + return ConversionResult(source, destination, "failed", result.stderr[-1000:]) + os.replace(temp, destination) + return ConversionResult(source, destination, "converted") + except OSError as exc: + return ConversionResult(source, destination, "failed", str(exc)) + finally: + temp.unlink(missing_ok=True) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("input_dir", type=Path, help="Directory containing .opus files") + parser.add_argument("output_dir", type=Path, help="Destination directory for .wav files") + parser.add_argument("--workers", type=int, default=min(16, os.cpu_count() or 1)) + parser.add_argument("--sample-rate", type=int, default=16000) + parser.add_argument("--channels", type=int, default=1) + parser.add_argument("--overwrite", action="store_true") + parser.add_argument( + "--ffmpeg", + default=os.environ.get("FFMPEG_BIN", "ffmpeg"), + help="ffmpeg executable (default: $FFMPEG_BIN or ffmpeg)", + ) + args = parser.parse_args() + + input_dir = args.input_dir.resolve() + output_dir = args.output_dir.resolve() + if not input_dir.is_dir(): + parser.error(f"input directory does not exist: {input_dir}") + ffmpeg = shutil.which(args.ffmpeg) + if ffmpeg is None: + parser.error(f"ffmpeg executable not found: {args.ffmpeg}") + + sources = sorted( + path for path in input_dir.rglob("*") if path.is_file() and path.suffix.lower() == ".opus" + ) + if not sources: + print(f"No Opus files found under {input_dir}") + return 0 + + workers = max(1, args.workers) + counts = {"converted": 0, "skipped": 0, "failed": 0} + failures: list[ConversionResult] = [] + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = [ + pool.submit( + convert_one, + source, + destination_for(source, input_dir, output_dir), + ffmpeg=ffmpeg, + sample_rate=args.sample_rate, + channels=args.channels, + overwrite=args.overwrite, + ) + for source in sources + ] + for completed, future in enumerate(as_completed(futures), start=1): + result = future.result() + counts[result.status] += 1 + if result.status == "failed": + failures.append(result) + if completed % 100 == 0 or completed == len(futures): + print( + f"Progress {completed}/{len(futures)}: " + f"converted={counts['converted']} skipped={counts['skipped']} " + f"failed={counts['failed']}" + ) + + for failure in failures: + print(f"FAILED {failure.source} -> {failure.destination}: {failure.error}") + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_common.py b/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_common.py new file mode 100644 index 0000000000..c3d2fb1c6e --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_common.py @@ -0,0 +1,1695 @@ +"""Shared helpers for the David AI MFA pipeline.""" + +from __future__ import annotations + +import json +import logging +import os +import re +import shutil +import subprocess +import threading +import traceback +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from pathlib import Path +from typing import TYPE_CHECKING, TypeVar + +if TYPE_CHECKING: + from collections.abc import Callable + +MFA_ROOT_DIR_DEFAULT = "~/MFA_models" + +POSTPROCESSED_RE = re.compile(r"^(.+)_postprocessed\.wav$") +SILENCE_TOKENS = {"", "sil", "sp", "spn", ""} + +logger = logging.getLogger(__name__) + +# Hard wall-clock cap for any single ffmpeg subprocess. Without this, a wedged +# ffmpeg (e.g. an internal futex deadlock seen when many run in a worker pool) +# blocks its caller forever and hangs the whole shard. On timeout, subprocess +# kills the child and raises TimeoutExpired, which callers treat as a failure. +# Configurable via FFMPEG_TIMEOUT_S: the final multi-speaker `amix` of long +# sessions can exceed the default when all node CPUs are saturated (mix reruns +# use a higher value and lower concurrency). +def _ffmpeg_timeout_s() -> int: + raw = os.environ.get("FFMPEG_TIMEOUT_S", "").strip() + if raw: + try: + value = int(float(raw)) + if value > 0: + return value + except ValueError: + pass + return 600 + + +FFMPEG_TIMEOUT_S = _ffmpeg_timeout_s() + + +def ffmpeg_executable() -> str: + """Return ffmpeg binary path; honors FFMPEG_BIN for cluster static builds.""" + env_bin = os.environ.get("FFMPEG_BIN", "").strip() + if env_bin: + return env_bin + return shutil.which("ffmpeg") or "ffmpeg" + + +T = TypeVar("T") +R = TypeVar("R") + + +class PipelineError(Exception): + """Raised when a pipeline stage cannot continue.""" + + +def mfa_models_root() -> Path: + return Path(os.environ.get("MFA_ROOT_DIR", MFA_ROOT_DIR_DEFAULT)).expanduser().resolve() + + +def resolve_mfa_dict(mfa_dict: str) -> Path: + for candidate in ( + mfa_models_root() / "pretrained_models" / "dictionary" / f"{mfa_dict}.dict", + mfa_models_root() / "pretrained_models" / "dictionary" / f"{mfa_dict}.txt", + Path(mfa_dict).expanduser(), + ): + if candidate.is_file(): + return candidate.resolve() + msg = f"MFA dictionary not found for {mfa_dict!r}" + raise FileNotFoundError(msg) + + +def resolve_mfa_acoustic_model(mfa_acoustic: str) -> str: + """Resolve an acoustic model name to its pretrained zip path. + + Passing the full zip path (instead of the bare model name) lets each MFA + invocation use its own ``--temporary_directory`` without breaking model + lookup, which is required for safe parallel execution. + """ + direct = Path(mfa_acoustic).expanduser() + if direct.is_file() or direct.is_dir(): + return str(direct.resolve()) + for candidate in ( + mfa_models_root() / "pretrained_models" / "acoustic" / f"{mfa_acoustic}.zip", + mfa_models_root() / "pretrained_models" / "acoustic" / mfa_acoustic, + ): + if candidate.is_file() or candidate.is_dir(): + return str(candidate.resolve()) + return mfa_acoustic + + +def partition_list(items: list[T], num_parts: int) -> list[list[T]]: + if num_parts <= 1 or not items: + return [items] + parts: list[list[T]] = [[] for _ in range(num_parts)] + for i, item in enumerate(items): + parts[i % num_parts].append(item) + return [part for part in parts if part] + + +def append_mfa_g2p_args(cmd: list[str], *, g2p_path: str | Path | None) -> None: + if g2p_path: + cmd.extend(["--g2p_model_path", str(g2p_path)]) + + +def mfa_subprocess_env( + *, + temp_root: Path, + mfa_root: Path, +) -> dict[str, str]: + """Build env for ``mfa`` subprocesses. + + Inside the pyxis container, ``PYTHONPATH`` includes ``/opt/venv/...`` which + shadows the packed conda ``pynini``/OpenFST with a different wheel build. + Mixed imports (conda ``montreal_forced_aligner`` + container ``pynini``) + cause ``FstIOError: Read failed`` on G2P ``model.fst``. + """ + env = os.environ.copy() + env["TMPDIR"] = str(temp_root.parent) + env["MFA_ROOT_DIR"] = str(mfa_root) + + conda_lib: Path | None = None + mfa_env_dir = os.environ.get("MFA_ENV", "").strip() + if mfa_env_dir: + candidate = Path(mfa_env_dir) / "lib" + if candidate.is_dir(): + conda_lib = candidate + if conda_lib is None: + mfa_bin = shutil.which("mfa") + if mfa_bin: + candidate = Path(mfa_bin).resolve().parent.parent / "lib" + if candidate.is_dir(): + conda_lib = candidate + if conda_lib is not None: + prev = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = f"{conda_lib}:{prev}" if prev else str(conda_lib) + + pp = env.get("PYTHONPATH", "") + if pp: + kept = [ + p + for p in pp.split(os.pathsep) + if p and "/opt/venv" not in p and "/opt/Export-Deploy" not in p + ] + if kept: + env["PYTHONPATH"] = os.pathsep.join(kept) + else: + env.pop("PYTHONPATH", None) + return env + + +def _extract_g2p_archive(g2p_src: Path, extract_root: Path) -> Path: + import zipfile + + extract_root.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(g2p_src) as archive: + for member in archive.infolist(): + relative = Path(member.filename) + if relative.is_absolute() or ".." in relative.parts: + msg = f"unsafe path in G2P archive: {member.filename}" + raise PipelineError(msg) + if member.is_dir(): + continue + destination = extract_root / relative + destination.parent.mkdir(parents=True, exist_ok=True) + with archive.open(member) as source, destination.open("wb") as target: + shutil.copyfileobj(source, target) + fst_paths = list(extract_root.rglob("model.fst")) + if len(fst_paths) != 1 or fst_paths[0].stat().st_size == 0: + msg = f"invalid G2P archive contents: {g2p_src}" + raise PipelineError(msg) + return fst_paths[0].parent + + +def _worker_g2p_arg(models_dir: Path, mfa_g2p: str) -> str | None: + g2p_src = resolve_mfa_g2p_model(mfa_g2p) + if g2p_src.is_dir(): + local_dir = models_dir / "g2p" / g2p_src.name + if not (local_dir / "model.fst").is_file(): + local_dir.parent.mkdir(parents=True, exist_ok=True) + if local_dir.exists(): + shutil.rmtree(local_dir) + shutil.copytree(g2p_src, local_dir) + return str(local_dir) + if g2p_src.is_file(): + if g2p_src.suffix == ".zip": + return str(_extract_g2p_archive(g2p_src, models_dir / "g2p")) + local = models_dir / g2p_src.name + if not local.is_file(): + shutil.copy2(g2p_src, local) + return str(local) + return None + + +def setup_mfa_worker_root( + worker_dir: Path, + *, + mfa_dict: Path, + mfa_acoustic: str, + mfa_g2p: str | None = None, + source_mfa_root: Path | None = None, +) -> tuple[Path, Path, str, str | None]: + """Prepare an isolated MFA root with local copies of lexicon and acoustic model. + + Returns (mfa_root, local_dict_path, acoustic_model_arg, g2p_model_arg) for ``mfa align``. + """ + worker_dir = worker_dir.resolve() + if worker_dir.exists(): + shutil.rmtree(worker_dir, ignore_errors=True) + mfa_root = worker_dir / "mfa_root" + models_dir = worker_dir / "models" + + models_dir.mkdir(parents=True, exist_ok=True) + mfa_root.mkdir(parents=True, exist_ok=True) + + local_dict = models_dir / mfa_dict.name + shutil.copy2(mfa_dict, local_dict) + + acoustic_src = Path(resolve_mfa_acoustic_model(mfa_acoustic)) + if acoustic_src.is_file() and acoustic_src.suffix == ".zip": + local_zip = models_dir / acoustic_src.name + shutil.copy2(acoustic_src, local_zip) + acoustic_arg = str(local_zip) + _extract_acoustic_zip(local_zip, mfa_root, source_mfa_root=source_mfa_root) + elif acoustic_src.is_dir(): + local_acoustic = models_dir / "acoustic" / acoustic_src.name + local_acoustic.parent.mkdir(parents=True, exist_ok=True) + shutil.copytree(acoustic_src, local_acoustic) + acoustic_arg = str(local_acoustic) + else: + acoustic_arg = resolve_mfa_acoustic_model(mfa_acoustic) + + g2p_arg = _worker_g2p_arg(models_dir, mfa_g2p) if mfa_g2p else None + + global_config = mfa_root / "global_config.yaml" + global_config.write_text( + "\n".join( + [ + "auto_server: true", + "blas_num_threads: 1", + "clean: false", + "cleanup_textgrids: true", + "database_limited_mode: false", + "debug: false", + "num_jobs: 3", + "overwrite: false", + "quiet: false", + "seed: 0", + "single_speaker: false", + f"temporary_directory: {mfa_root}", + "use_mp: true", + "use_postgres: false", + "use_threading: true", + "verbose: false", + ] + ) + + "\n", + encoding="utf-8", + ) + + cmd_hist = mfa_root / "command_history.yaml" + if cmd_hist.exists() or cmd_hist.is_symlink(): + cmd_hist.unlink(missing_ok=True) + cmd_hist.symlink_to("/dev/null") + + return mfa_root, local_dict, acoustic_arg, g2p_arg + + +def _extract_acoustic_zip( + zip_path: Path, + mfa_root: Path, + *, + source_mfa_root: Path | None = None, +) -> None: + import zipfile + + extracted_root = mfa_root / "extracted_models" / "acoustic" + if source_mfa_root is not None: + src_acoustic = source_mfa_root / "extracted_models" / "acoustic" + if src_acoustic.is_dir(): + for src_dir in src_acoustic.iterdir(): + if not src_dir.is_dir(): + continue + dst_dir = extracted_root / src_dir.name + if (dst_dir / "final.mdl").is_file(): + continue + shutil.copytree(src_dir, dst_dir) + return + + extracted_root.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(zip_path) as zf: + zf.extractall(extracted_root) + + for _path in extracted_root.rglob("final.mdl"): + return + msg = f"acoustic zip did not contain final.mdl: {zip_path}" + raise PipelineError(msg) + + +def resolve_mfa_g2p_model(mfa_g2p: str) -> Path: + direct = Path(mfa_g2p).expanduser() + if direct.is_file() or direct.is_dir(): + return direct.resolve() + root = mfa_models_root() + for candidate in ( + root / "extracted_models" / "g2p" / f"{mfa_g2p}_g2p", + root / "extracted_models" / "g2p" / mfa_g2p, + root / "pretrained_models" / "g2p" / mfa_g2p, + root / "pretrained_models" / "g2p" / f"{mfa_g2p}.zip", + ): + if candidate.is_file() or candidate.is_dir(): + return candidate.resolve() + msg = ( + f"MFA G2P model not found for {mfa_g2p!r} under " + f"{root / 'pretrained_models' / 'g2p'} or {root / 'extracted_models' / 'g2p'}" + ) + raise FileNotFoundError( + msg + ) + + +def log_exception(context: str, exc: BaseException) -> None: + logger.error("%s: %s", context, exc) + logger.debug(traceback.format_exc()) + + +def load_jsonl(path: Path) -> list[dict]: + rows: list[dict] = [] + try: + with path.open(encoding="utf-8") as f: + for line_no, raw_line in enumerate(f, start=1): + line = raw_line.strip() + if not line: + continue + try: + rows.append(json.loads(line)) + except json.JSONDecodeError as exc: + msg = f"{path}:{line_no}: invalid JSON: {exc}" + raise ValueError(msg) from exc + except OSError as exc: + msg = f"cannot read {path}: {exc}" + raise PipelineError(msg) from exc + return rows + + +def write_jsonl(path: Path, rows: list[dict]) -> None: + try: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + for row in rows: + f.write(json.dumps(row, ensure_ascii=False) + "\n") + except OSError as exc: + msg = f"cannot write {path}: {exc}" + raise PipelineError(msg) from exc + + +def append_jsonl(path: Path, row: dict, *, lock: threading.Lock | None = None) -> None: + def _write() -> None: + try: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as f: + f.write(json.dumps(row, ensure_ascii=False) + "\n") + except OSError as exc: + msg = f"cannot append to {path}: {exc}" + raise PipelineError(msg) from exc + + if lock is not None: + with lock: + _write() + else: + _write() + + +def thread_temp_root(base: Path) -> Path: + """Per-thread scratch directory under *base* for parallel workers.""" + root = base / f"thread_{threading.get_ident()}" + root.mkdir(parents=True, exist_ok=True) + return root + + +def run_thread_pool( + items: list[T], + fn: Callable[[T], R], + *, + workers: int = 1, +) -> list[R]: + if workers <= 1 or len(items) <= 1: + return [fn(item) for item in items] + with ThreadPoolExecutor(max_workers=workers) as pool: + return list(pool.map(fn, items)) + + +def run_process_pool( + items: list[T], + fn: Callable[[T], R], + *, + workers: int = 1, +) -> list[R]: + if workers <= 1 or len(items) <= 1: + return [fn(item) for item in items] + with ProcessPoolExecutor(max_workers=workers) as pool: + return list(pool.map(fn, items)) + + +def ffprobe_duration(path: Path) -> float: + cmd = [ + "ffprobe", + "-v", + "error", + "-show_entries", + "format=duration", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(path), + ] + try: + result = subprocess.run(cmd, capture_output=True, text=True, check=False) + except OSError as exc: + msg = f"ffprobe not available for {path}: {exc}" + raise RuntimeError(msg) from exc + if result.returncode != 0 or not result.stdout.strip(): + msg = f"ffprobe failed for {path}: {result.stderr[-300:]}" + raise RuntimeError(msg) + try: + return float(result.stdout.strip()) + except ValueError as exc: + msg = f"ffprobe returned non-numeric duration for {path}" + raise RuntimeError(msg) from exc + + +def extract_segment_wav( + src: Path, + dst: Path, + start: float, + end: float, + *, + padding: float = 0.0, + max_duration: float | None = None, +) -> float | None: + """Extract audio for MFA, optionally padded before/after the manifest interval. + + Returns the absolute extract start time in *src*, or None on failure. + The clip spans [max(0, start-padding), min(max_duration, end+padding)] when + *max_duration* is set, otherwise [max(0, start-padding), end+padding]. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + extract_start = max(0.0, start - padding) + extract_end = end + padding + if max_duration is not None: + extract_end = min(max_duration, extract_end) + duration = max(extract_end - extract_start, 0.01) + cmd = [ + ffmpeg_executable(), + "-nostdin", + "-y", + "-ss", + f"{extract_start:.6f}", + "-i", + str(src), + "-t", + f"{duration:.6f}", + "-ac", + "1", + "-acodec", + "pcm_s16le", + str(dst), + ] + try: + result = subprocess.run( + cmd, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + timeout=FFMPEG_TIMEOUT_S, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + logger.warning("ffmpeg segment extract failed to start for %s: %s", src, exc) + return None + if result.returncode != 0: + logger.warning( + "ffmpeg segment extract failed (%s [%.3f-%.3f] pad=%.3f): %s", + src, + extract_start, + extract_end, + padding, + result.stderr[-300:], + ) + return None + return extract_start + + +def map_segment_words_to_recording( + words: list[tuple[float, float, str]], + *, + extract_start: float, + extract_end: float, +) -> list[tuple[float, float, str]]: + """Map MFA word times from a padded segment clip back to recording time. + + All MFA word intervals are kept as aligned (only bounded by the clip MFA ran on). + """ + mapped: list[tuple[float, float, str]] = [] + for start, end, word in words: + abs_start = start + extract_start + abs_end = end + extract_start + if abs_end <= extract_start or abs_start >= extract_end: + continue + mapped.append((abs_start, abs_end, word)) + return mapped + + +def parse_textgrid_words(tg_path: Path) -> list[tuple[float, float, str]]: + import textgrid + + try: + tg = textgrid.TextGrid.fromFile(str(tg_path)) + tier = tg.getFirst("words") + except Exception as exc: + msg = f"failed to parse TextGrid {tg_path}: {exc}" + raise ValueError(msg) from exc + + words: list[tuple[float, float, str]] = [] + for iv in tier.intervals: + mark = (iv.mark or "").strip() + if mark and mark not in SILENCE_TOKENS: + words.append((iv.minTime, iv.maxTime, mark)) + return words + + +def safe_parse_textgrid_words(tg_path: Path) -> list[tuple[float, float, str]]: + try: + return parse_textgrid_words(tg_path) + except ImportError as exc: + msg = ( + "textgrid package is required to parse MFA TextGrids " + f"(pip install textgrid): {exc}" + ) + raise PipelineError( + msg + ) from exc + except Exception as exc: + log_exception(f"TextGrid parse failed for {tg_path}", exc) + return [] + + +def write_textgrid( + words: list[tuple[float, float, str]], + output_path: Path, + *, + xmin: float = 0.0, + xmax: float | None = None, +) -> None: + if xmax is None: + xmax = words[-1][1] + 0.01 if words else xmin + 0.01 + + intervals: list[tuple[float, float, str]] = [] + prev_end = xmin + for start, end, word in sorted(words, key=lambda x: x[0]): + if start > prev_end + 0.001: + intervals.append((prev_end, start, "")) + intervals.append((start, end, word)) + prev_end = end + if prev_end < xmax: + intervals.append((prev_end, xmax, "")) + + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", encoding="utf-8") as f: + f.write('File type = "ooTextFile"\n') + f.write('Object class = "TextGrid"\n\n') + f.write(f"xmin = {xmin}\n") + f.write(f"xmax = {xmax}\n") + f.write("tiers? \n") + f.write("size = 1\n") + f.write("item []:\n") + f.write(" item [1]:\n") + f.write(' class = "IntervalTier"\n') + f.write(' name = "words"\n') + f.write(f" xmin = {xmin}\n") + f.write(f" xmax = {xmax}\n") + f.write(f" intervals: size = {len(intervals)}\n") + for i, (s, e, text) in enumerate(intervals, 1): + safe = text.replace('"', '""') + f.write(f" intervals [{i}]:\n") + f.write(f" xmin = {s}\n") + f.write(f" xmax = {e}\n") + f.write(f' text = "{safe}"\n') + + +def merge_speech_intervals( + intervals: list[tuple[float, float]], + merge_gap: float, +) -> list[tuple[float, float]]: + merged: list[tuple[float, float]] = [] + for start, end in sorted(intervals): + if end <= start: + continue + if merged and (start - merged[-1][1]) <= merge_gap: + merged[-1] = (merged[-1][0], max(merged[-1][1], end)) + else: + merged.append((start, end)) + return merged + + +def merge_tagged_speech_intervals( + intervals: list[tuple[float, float, str]], + merge_gap: float, +) -> list[tuple[float, float, str]]: + merged: list[tuple[float, float, str]] = [] + for start, end, label in sorted(intervals): + if end <= start: + continue + if merged and (start - merged[-1][1]) <= merge_gap: + prev_start, prev_end, prev_label = merged[-1] + merged[-1] = (prev_start, max(prev_end, end), prev_label) + else: + merged.append((start, end, label)) + return merged + + +def textgrid_to_rttm_lines( + tg_path: Path, + *, + speaker: str, + merge_gap: float = 0.2, +) -> list[str]: + file_id = tg_path.stem + try: + intervals = [ + (start, end) + for start, end, _ in parse_textgrid_words(tg_path) + ] + except Exception as exc: + log_exception(f"RTTM conversion failed for {tg_path}", exc) + return [] + + merged = merge_speech_intervals(intervals, merge_gap) + + lines = [] + for start, end in merged: + dur = end - start + lines.append( + f"SPEAKER {file_id} 1 {start:.6f} {dur:.6f} {speaker} " + ) + return lines + + +def discover_sessions(data_root: Path) -> list[Path]: + sessions: list[Path] = [] + for path in data_root.iterdir(): + if path.is_symlink(): + # Cluster data_links are already filtered by the link stage. Avoid + # following every symlink to Lustre before stage 0 can start writing. + sessions.append(path) + elif path.is_dir() and (path / "machine_generated_transcript.json").is_file(): + sessions.append(path) + return sorted(sessions) + + +def load_speaker_count_tsv(path: Path) -> dict[str, int]: + """Load `` `` lines from a speaker-count TSV.""" + counts: dict[str, int] = {} + if not path.is_file(): + return counts + for raw_line in path.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + parts = line.split() + if len(parts) < 2: + continue + try: + counts[parts[1]] = int(parts[0]) + except ValueError: + continue + return counts + + +def load_session_id_list(path: Path) -> list[str]: + """Load one session id per line (comments and blanks ignored).""" + if not path.is_file(): + return [] + ids: list[str] = [] + for raw_line in path.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + ids.append(line.split()[0]) + return ids + + +def filter_sessions_by_ids(sessions: list[Path], session_ids: list[str]) -> list[Path]: + wanted = set(session_ids) + if not wanted: + return sessions + by_name = {session.name: session for session in sessions} + missing = sorted(wanted - set(by_name)) + if missing: + logger.warning("sessions-file: %d id(s) not found under data root (first: %s)", len(missing), missing[0]) + return [by_name[sid] for sid in session_ids if sid in by_name] + + +def order_sessions_by_speaker_priority( + sessions: list[Path], + speaker_counts: dict[str, int], + *, + min_priority_speakers: int, +) -> list[Path]: + """Put sessions with at least *min_priority_speakers* first (higher counts earlier).""" + if min_priority_speakers <= 1 or not speaker_counts: + return sessions + + def sort_key(session_dir: Path) -> tuple[int, int, str]: + count = speaker_counts.get(session_dir.name, 0) + priority_bucket = 0 if count >= min_priority_speakers else 1 + return (priority_bucket, -count, session_dir.name) + + return sorted(sessions, key=sort_key) + + +def recording_id(speaker_id: str, session_id: str) -> str: + return f"{speaker_id}_{session_id}_postprocessed" + + +def mixed_speaker_audio_path(audio_mixed_dir: Path, speaker_id: str, session_id: str) -> Path: + """Persistent pause-noise track used as one input to the session mix.""" + return audio_mixed_dir / "speakers" / f"{recording_id(speaker_id, session_id)}.opus" + + +def recording_textgrid_path( + textgrid_dir: Path, + recording_id: str, + *, + variant: str = "ordinary", +) -> Path: + suffix = {"ordinary": "", "fastmss": "_fastmss", "fb": "_fb"}.get(variant, "") + return textgrid_dir / f"{recording_id}{suffix}.TextGrid" + + +def recording_textgrid_paths(textgrid_dir: Path, recording_id: str) -> list[Path]: + ordinary = recording_textgrid_path(textgrid_dir, recording_id, variant="ordinary") + fb_path = recording_textgrid_path(textgrid_dir, recording_id, variant="fb") + if ordinary.is_file() and fb_path.is_file(): + return [ordinary, fb_path] + if ordinary.is_file(): + return [ordinary] + if fb_path.is_file(): + return [fb_path] + return [ordinary] + + +def fastmss_textgrid_path(textgrid_dir: Path, recording_id: str) -> Path: + return recording_textgrid_path(textgrid_dir, recording_id, variant="fastmss") + + +def session_textgrid_path( + textgrid_dir: Path, + session_id: str, + *, + variant: str = "ordinary", +) -> Path: + suffix = {"ordinary": "", "fastmss": "_fastmss"}.get(variant, "") + return textgrid_dir / f"{session_id}{suffix}.TextGrid" + + +def interval_overlaps(start: float, end: float, intervals: list[tuple[float, float]]) -> bool: + return any(start < interval_end and end > interval_start for interval_start, interval_end in intervals) + + +def speech_intervals_from_textgrid(tg_path: Path) -> list[tuple[float, float]]: + return [(start, end) for start, end, _ in parse_textgrid_words(tg_path)] + + +def fb_intervals_for_recording(textgrid_dir: Path, recording_id: str) -> list[tuple[float, float]]: + fb_path = recording_textgrid_path(textgrid_dir, recording_id, variant="fb") + if fb_path.is_file(): + return speech_intervals_from_textgrid(fb_path) + ordinary = recording_textgrid_path(textgrid_dir, recording_id, variant="ordinary") + fastmss = recording_textgrid_path(textgrid_dir, recording_id, variant="fastmss") + if ordinary.is_file() and fastmss.is_file(): + ordinary_words = parse_textgrid_words(ordinary) + fastmss_words = parse_textgrid_words(fastmss) + if len(ordinary_words) > len(fastmss_words): + fastmss_intervals = [(s, e) for s, e, _ in fastmss_words] + return [ + (start, end) + for start, end, word in ordinary_words + if word == "speech" and not interval_overlaps(start, end, fastmss_intervals) + ] + return [] + + +def alignment_items_from_textgrid(tg_path: Path) -> list: + from lhotse.supervision import AlignmentItem + + try: + words = parse_textgrid_words(tg_path) + except Exception as exc: + log_exception(f"alignment extraction failed for {tg_path}", exc) + return [] + + items = [] + for start, end, word in words: + items.append( + AlignmentItem( + symbol=word, + start=round(start, 6), + duration=round(end - start, 6), + ) + ) + return items + + +def alignment_items_for_lhotse( + main_tg_path: Path, + *, + fb_intervals: list[tuple[float, float]] | None = None, +) -> list: + if fb_intervals is None: + fb_intervals = fb_intervals_for_recording(main_tg_path.parent, main_tg_path.stem) + + from lhotse.supervision import AlignmentItem + + try: + words = parse_textgrid_words(main_tg_path) + except Exception as exc: + log_exception(f"alignment extraction failed for {main_tg_path}", exc) + return [] + + items = [] + for start, end, word in words: + if fb_intervals and interval_overlaps(start, end, fb_intervals): + continue + items.append( + AlignmentItem( + symbol=word, + start=round(start, 6), + duration=round(end - start, 6), + ) + ) + return items + + +def write_lhotse_concatenated_textgrid( + main_tg_path: Path, + output_path: Path, + *, + fb_intervals: list[tuple[float, float]] | None = None, + xmax: float | None = None, +) -> None: + if fb_intervals is None: + fb_intervals = fb_intervals_for_recording(main_tg_path.parent, main_tg_path.stem) + + words = parse_textgrid_words(main_tg_path) + if fb_intervals: + words = [ + (start, end, word) + for start, end, word in words + if not interval_overlaps(start, end, fb_intervals) + ] + write_textgrid(words, output_path, xmin=0.0, xmax=xmax) + + +def write_rttm(path: Path, lines: list[str]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8") + + +def speech_rttm_line( + file_id: str, + speaker: str, + start: float, + end: float, + *, + label: str = "speech", +) -> str: + """RTTM speech interval; *label* is written to the subtype field.""" + dur = max(end - start, 0.0) + return ( + f"SPEAKER {file_id} 1 {start:.6f} {dur:.6f} {label} {speaker} " + ) + + +def build_speech_rttm_lines( + file_id: str, + speaker: str, + intervals: list[tuple[float, float]], + *, + label: str = "speech", + merge_gap: float = 0.0, +) -> list[str]: + merged: list[tuple[float, float]] = [] + for start, end in sorted(intervals): + if end <= start: + continue + if merged and (start - merged[-1][1]) <= merge_gap: + merged[-1] = (merged[-1][0], max(merged[-1][1], end)) + else: + merged.append((start, end)) + return [ + speech_rttm_line(file_id, speaker, start, end, label=label) + for start, end in merged + ] + + +def normalization_log_entry( + row: dict, + *, + text_raw: str | None = None, + text_norm: str, + num2words_lang: str, + error: str = "", +) -> dict: + text = text_raw if text_raw is not None else (row.get("text_raw") or row.get("text") or "") + return { + "session_id": row.get("session_id"), + "speaker_id": row.get("speaker_id"), + "recording_id": row.get("recording_id"), + "segment_index": row.get("segment_index"), + "start": float(row["start"]), + "end": float(row["end"]), + "duration": float(row.get("duration", float(row["end"]) - float(row["start"]))), + "text": text, + "text_norm": text_norm, + "changed": text != text_norm, + "num2words_lang": num2words_lang, + "error": error, + } + + +def segment_fallback_log_entry( + seg: dict, + recording_id: str, + *, + reason: str, + detail: str = "", +) -> dict: + return { + "recording_id": recording_id, + "session_id": seg.get("session_id"), + "speaker_id": seg.get("speaker_id"), + "segment_index": seg.get("segment_index"), + "start": float(seg["start"]), + "end": float(seg["end"]), + "duration": float(seg.get("duration", float(seg["end"]) - float(seg["start"]))), + "text_norm": seg.get("text_norm", ""), + "reason": reason, + "detail": detail, + "fallback": "manifest_boundaries", + "rttm_label": "speech", + } + + +def merge_session_rttm( + rttm_paths: list[Path], + session_id: str, + output_path: Path, +) -> int: + parsed: list[tuple[float, str]] = [] + for path in rttm_paths: + if not path.is_file(): + logger.warning("%s: missing per-recording RTTM %s", session_id, path) + continue + try: + content = path.read_text(encoding="utf-8") + except OSError as exc: + log_exception(f"cannot read RTTM {path}", exc) + continue + for raw_line in content.splitlines(): + line = raw_line.strip() + if not line or line.startswith(";"): + continue + parts = line.split() + if len(parts) < 8 or parts[0] != "SPEAKER": + continue + try: + start = float(parts[3]) + except ValueError: + logger.warning("%s: invalid RTTM start time in %s: %s", session_id, path.name, line) + continue + parts[1] = session_id + parsed.append((start, " ".join(parts))) + + parsed.sort(key=lambda x: x[0]) + lines = [line for _, line in parsed] + write_rttm(output_path, lines) + return len(lines) + + +def load_norm_manifest_rows( + manifests_dir: Path, + *, + sessions: list[str] | None = None, +) -> tuple[list[dict], int]: + rows: list[dict] = [] + manifest_errors = 0 + wanted = set(sessions) if sessions else None + for path in sorted(manifests_dir.glob("*_norm.jsonl")): + if path.name == "all_norm.jsonl": + continue + try: + file_rows = load_jsonl(path) + except Exception as exc: + manifest_errors += 1 + log_exception(f"cannot load manifest {path}", exc) + continue + if wanted is not None: + file_rows = [r for r in file_rows if r.get("session_id") in wanted] + rows.extend(file_rows) + return rows, manifest_errors + + +def group_segments_by_recording(rows: list[dict]) -> dict[str, list[dict]]: + from collections import defaultdict + + grouped: dict[str, list[dict]] = defaultdict(list) + for row in rows: + grouped[row["recording_id"]].append(row) + for segments in grouped.values(): + segments.sort(key=lambda r: (r["start"], r["segment_index"])) + return grouped + + +def group_segments_by_session(rows: list[dict]) -> dict[str, list[dict]]: + from collections import defaultdict + + grouped: dict[str, list[dict]] = defaultdict(list) + for row in rows: + grouped[row["session_id"]].append(row) + for segments in grouped.values(): + segments.sort(key=lambda r: (r["start"], r["segment_index"])) + return grouped + + +def group_recordings_by_session(rows: list[dict]) -> dict[str, list[dict]]: + from collections import defaultdict + + grouped: dict[str, list[dict]] = defaultdict(list) + seen: dict[str, set[str]] = defaultdict(set) + for row in rows: + session_id = row["session_id"] + rec_id = row["recording_id"] + if rec_id in seen[session_id]: + continue + seen[session_id].add(rec_id) + grouped[session_id].append( + { + "recording_id": rec_id, + "speaker_id": row["speaker_id"], + "audio_path": Path(row["audio_filepath_16k"]), + } + ) + for recordings in grouped.values(): + recordings.sort(key=lambda e: e["recording_id"]) + return grouped + + +def load_fallback_intervals(fallback_log: Path, recording_id: str) -> list[tuple[float, float]]: + if not fallback_log.is_file(): + return [] + intervals: list[tuple[float, float]] = [] + try: + for row in load_jsonl(fallback_log): + if row.get("recording_id") != recording_id: + continue + intervals.append((float(row["start"]), float(row["end"]))) + except Exception as exc: + log_exception(f"cannot read fallback log for {recording_id}", exc) + return intervals + + +def alignment_items_from_words(words: list[tuple[float, float, str]]) -> list: + from lhotse.supervision import AlignmentItem + + return [ + AlignmentItem( + symbol=word, + start=round(start, 6), + duration=round(end - start, 6), + ) + for start, end, word in words + ] + + +def words_to_json(words: list[tuple[float, float, str]]) -> list[list]: + return [[start, end, word] for start, end, word in words] + + +def words_from_json(raw: list) -> list[tuple[float, float, str]]: + return [(float(s), float(e), str(w)) for s, e, w in raw] + + +def tagged_words_to_json(words: list[tuple[float, float, str, str]]) -> list[list]: + return [[start, end, word, speaker_id] for start, end, word, speaker_id in words] + + +def tagged_words_from_json(raw: list) -> list[tuple[float, float, str, str]]: + out: list[tuple[float, float, str, str]] = [] + for item in raw: + if len(item) == 4: + s, e, w, spk = item + out.append((float(s), float(e), str(w), str(spk))) + else: + s, e, w = item + out.append((float(s), float(e), str(w), "")) + return out + + +def alignment_record( + recording_id: str, + segments: list[dict], + *, + merged_words: list[tuple[float, float, str]], + fb_words: list[tuple[float, float, str]], + audio_duration: float, +) -> dict: + return { + "recording_id": recording_id, + "speaker_id": segments[0]["speaker_id"], + "session_id": segments[0]["session_id"], + "audio_filepath_16k": segments[0]["audio_filepath_16k"], + "audio_duration": audio_duration, + "merged_words": words_to_json(merged_words), + "fb_words": words_to_json(fb_words), + } + + +def session_alignment_record( + session_id: str, + *, + merged_words: list[tuple[float, float, str, str]], + fb_words: list[tuple[float, float, str, str]], + audio_duration: float, + recordings: list[dict], +) -> dict: + return { + "session_id": session_id, + "audio_duration": audio_duration, + "merged_words": tagged_words_to_json(merged_words), + "fb_words": tagged_words_to_json(fb_words), + "recordings": recordings, + } + + +def append_alignment_record( + path: Path, + record: dict, + *, + lock: threading.Lock | None = None, +) -> None: + append_jsonl(path, record, lock=lock) + + +def load_alignment_ids(path: Path) -> set[str]: + if not path.is_file(): + return set() + ids: set[str] = set() + for row in load_jsonl(path): + if "session_id" in row: + ids.add(row["session_id"]) + elif "recording_id" in row: + ids.add(row["recording_id"]) + return ids + + +def load_alignments_by_session(path: Path) -> dict[str, dict]: + by_id: dict[str, dict] = {} + if not path.is_file(): + return by_id + for row in load_jsonl(path): + if "session_id" in row: + by_id[row["session_id"]] = row + return by_id + + +def load_alignments_by_recording(path: Path) -> dict[str, dict]: + by_id: dict[str, dict] = {} + if not path.is_file(): + return by_id + for row in load_jsonl(path): + if "recordings" in row: + for rec in row["recordings"]: + by_id[rec["recording_id"]] = rec + elif "recording_id" in row: + by_id[row["recording_id"]] = row + return by_id + + +def build_rttm_lines_from_words( + recording_id: str, + speaker_id: str, + merged_words: list[tuple[float, float, str]], + fb_words: list[tuple[float, float, str]], + *, + merge_gap: float = 0.2, +) -> list[str]: + tagged: list[tuple[float, float, str]] = [] + for start, end, _ in merged_words: + tagged.append((start, end, "")) + for start, end, _ in fb_words: + tagged.append((start, end, "speech")) + merged = merge_tagged_speech_intervals(tagged, merge_gap) + return [ + speech_rttm_line(recording_id, speaker_id, start, end, label=label) + for start, end, label in merged + ] + + +def build_session_rttm_lines_from_words( + session_id: str, + merged_words: list[tuple[float, float, str, str]], + fb_words: list[tuple[float, float, str, str]], + *, + merge_gap: float = 0.2, +) -> list[str]: + from collections import defaultdict + + by_speaker: dict[str, list[tuple[float, float, str]]] = defaultdict(list) + for start, end, _, speaker_id in merged_words: + by_speaker[speaker_id].append((start, end, "")) + for start, end, _, speaker_id in fb_words: + by_speaker[speaker_id].append((start, end, "speech")) + + lines: list[tuple[float, str]] = [] + for speaker_id, tagged in by_speaker.items(): + merged = merge_tagged_speech_intervals(tagged, merge_gap) + for start, end, label in merged: + line = speech_rttm_line(session_id, speaker_id, start, end, label=label) + lines.append((start, line)) + lines.sort(key=lambda x: x[0]) + return [line for _, line in lines] + + +def merge_session_rttm_from_line_lists( + session_id: str, + line_lists: list[list[str]], +) -> list[str]: + parsed: list[tuple[float, str]] = [] + for lines in line_lists: + for raw_line in lines: + line = raw_line.strip() + if not line or line.startswith(";"): + continue + parts = line.split() + if len(parts) < 8 or parts[0] != "SPEAKER": + continue + try: + start = float(parts[3]) + except ValueError: + continue + parts[1] = session_id + parsed.append((start, " ".join(parts))) + parsed.sort(key=lambda x: x[0]) + return [line for _, line in parsed] + + +def build_recording_rttm_lines( + recording_id: str, + speaker_id: str, + tg_path: Path, + *, + fallback_log: Path | None = None, + merge_gap: float = 0.2, +) -> list[str]: + textgrid_dir = tg_path.parent + tagged: list[tuple[float, float, str]] = [] + for path in recording_textgrid_paths(textgrid_dir, recording_id): + try: + for start, end, word in parse_textgrid_words(path): + label = "speech" if word == "speech" else "" + tagged.append((start, end, label)) + except Exception as exc: + log_exception(f"RTTM conversion failed for {path}", exc) + + fb_path = recording_textgrid_path(textgrid_dir, recording_id, variant="fb") + if fallback_log is not None and not fb_path.is_file(): + for start, end in load_fallback_intervals(fallback_log, recording_id): + tagged.append((start, end, "speech")) + + merged = merge_tagged_speech_intervals(tagged, merge_gap) + return [ + speech_rttm_line(recording_id, speaker_id, start, end, label=label) + for start, end, label in merged + ] + + +def pad_speech_intervals( + speech_intervals: list[tuple[float, float]], + pad: float, + duration: float, +) -> list[tuple[float, float]]: + """Grow each speech interval by *pad* seconds on both sides, clamped to [0, duration]. + + Overlaps created by padding are merged. Used to keep a margin of untouched + audio around speech boundaries so pause noise never abuts speech. + """ + if pad <= 0: + return merge_speech_intervals(speech_intervals, 0.0) + padded = [ + (max(0.0, start - pad), min(duration, end + pad)) + for start, end in speech_intervals + ] + return merge_speech_intervals(padded, 0.0) + + +def invert_intervals( + speech_intervals: list[tuple[float, float]], + duration: float, +) -> list[tuple[float, float]]: + """Return gaps between *speech_intervals* over [0, *duration*] (pause regions).""" + pauses: list[tuple[float, float]] = [] + cursor = 0.0 + for start, end in sorted(speech_intervals): + if start > cursor + 1e-6: + pauses.append((cursor, start)) + cursor = max(cursor, end) + if cursor < duration - 1e-6: + pauses.append((cursor, duration)) + return pauses + + +def parse_rttm_speech_intervals( + lines: list[str], + *, + merge_gap: float = 0.2, +) -> list[tuple[float, float]]: + """Speech intervals from RTTM lines ( and speech subtype labels).""" + raw: list[tuple[float, float]] = [] + for raw_line in lines: + line = raw_line.strip() + if not line or line.startswith(";"): + continue + parts = line.split() + if len(parts) < 8 or parts[0] != "SPEAKER": + continue + label = parts[6] + if label not in {"", "speech"}: + continue + try: + start = float(parts[3]) + dur = float(parts[4]) + except ValueError: + continue + if dur > 0: + raw.append((start, start + dur)) + return merge_speech_intervals(raw, merge_gap) + + +def parse_session_rttm_by_speaker( + lines: list[str], + *, + merge_gap: float = 0.2, +) -> dict[str, list[tuple[float, float]]]: + """Speech intervals per speaker from a session-level RTTM (stage 4 output).""" + from collections import defaultdict + + by_speaker: dict[str, list[tuple[float, float]]] = defaultdict(list) + for raw_line in lines: + line = raw_line.strip() + if not line or line.startswith(";"): + continue + parts = line.split() + if len(parts) < 8 or parts[0] != "SPEAKER": + continue + label = parts[6] + if label not in {"", "speech"}: + continue + speaker_id = parts[7] + try: + start = float(parts[3]) + dur = float(parts[4]) + except ValueError: + continue + if dur > 0: + by_speaker[speaker_id].append((start, start + dur)) + + return { + speaker_id: merge_speech_intervals(intervals, merge_gap) + for speaker_id, intervals in by_speaker.items() + } + + +def session_rttm_path(audio_mixed_dir: Path, session_id: str) -> Path: + return audio_mixed_dir / f"{session_id}.rttm" + + +def load_session_rttm_by_speaker( + rttm_path: Path, + *, + merge_gap: float = 0.2, +) -> dict[str, list[tuple[float, float]]]: + if not rttm_path.is_file(): + return {} + lines = rttm_path.read_text(encoding="utf-8").splitlines() + return parse_session_rttm_by_speaker(lines, merge_gap=merge_gap) + + +def speech_intervals_from_recording_alignment( + rec_row: dict, + *, + merge_gap: float = 0.2, +) -> list[tuple[float, float]]: + merged = words_from_json(rec_row.get("merged_words", [])) + fb = words_from_json(rec_row.get("fb_words", [])) + raw = [(start, end) for start, end, _ in merged] + [(start, end) for start, end, _ in fb] + return merge_speech_intervals(raw, merge_gap) + + +def decode_audio_mono_f32(path: Path, *, target_sr: int = 16000) -> tuple: + import numpy as np + + cmd = [ + ffmpeg_executable(), + "-nostdin", + "-i", + str(path), + "-ar", + str(target_sr), + "-ac", + "1", + "-f", + "f32le", + "pipe:1", + ] + try: + result = subprocess.run( + cmd, + stdin=subprocess.DEVNULL, + capture_output=True, + check=False, + timeout=FFMPEG_TIMEOUT_S, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + logger.warning("ffmpeg decode failed to start for %s: %s", path, exc) + raise + if result.returncode != 0: + msg = result.stderr.decode(errors="replace")[-400:] + msg_0 = f"ffmpeg decode failed for {path}: {msg}" + raise RuntimeError(msg_0) + audio = np.frombuffer(result.stdout, dtype=np.float32) + return audio, target_sr + + +def encode_audio_mono_f32_to_opus( + audio, + dst: Path, + *, + sample_rate: int = 16000, + opus_bitrate: str = "32k", +) -> bool: + dst.parent.mkdir(parents=True, exist_ok=True) + cmd = [ + ffmpeg_executable(), + "-y", + "-f", + "f32le", + "-ar", + str(sample_rate), + "-ac", + "1", + "-i", + "pipe:0", + "-c:a", + "libopus", + "-b:a", + opus_bitrate, + "-application", + "voip", + "-vbr", + "on", + str(dst), + ] + try: + result = subprocess.run( + cmd, + input=audio.tobytes(), + capture_output=True, + check=False, + timeout=FFMPEG_TIMEOUT_S, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + logger.warning("ffmpeg opus encode failed to start for %s: %s", dst, exc) + return False + if result.returncode != 0: + logger.warning( + "ffmpeg opus encode failed for %s: %s", + dst, + result.stderr.decode(errors="replace")[-400:], + ) + return False + return True + + +def apply_white_noise_in_pause_intervals( + src: Path, + dst: Path, + pause_intervals: list[tuple[float, float]], + *, + target_sr: int = 16000, + noise_level: float = 0.0002, + opus_bitrate: str = "32k", + seed: int | None = None, + preserve_speech: bool = True, + stitch_ms: float = 5.0, +) -> bool: + """Replace *pause_intervals* with white noise; RTTM speech samples stay intact. + + Speech regions are never modified. Only samples inside each pause interval are + written. When *preserve_speech* is True and *stitch_ms* > 0, a linear + crossfade is applied **inside** the pause at both ends: original pause audio + is blended into white noise at the pause start and blended back out before the + pause end (still within the pause interval). + + When *preserve_speech* is False, the pause interior is filled with noise + with no crossfade. + """ + import numpy as np + + if not pause_intervals: + import shutil + + try: + shutil.copy2(src, dst) + return True + except OSError as exc: + log_exception(f"cannot copy audio to {dst}", exc) + return False + + try: + audio, sr = decode_audio_mono_f32(src, target_sr=target_sr) + except RuntimeError as exc: + log_exception(f"decode failed for {src}", exc) + return False + + audio = np.array(audio, dtype=np.float32, copy=True) + n = len(audio) + rng = np.random.default_rng(seed) + stitch = max(0, round(stitch_ms * sr / 1000.0)) if preserve_speech else 0 + + for start, end in pause_intervals: + i0 = max(0, int(start * sr)) + i1 = min(n, int(end * sr)) + length = i1 - i0 + if length <= 0: + continue + + orig_pause = audio[i0:i1].copy() + noise = rng.standard_normal(length, dtype=np.float32) * noise_level + + if stitch > 0: + fade = min(stitch, length // 2) + if fade > 0: + ramp_in = np.linspace(0.0, 1.0, fade, dtype=np.float32) + ramp_out = np.linspace(1.0, 0.0, fade, dtype=np.float32) + noise[:fade] = orig_pause[:fade] * (1.0 - ramp_in) + noise[:fade] * ramp_in + noise[-fade:] = orig_pause[-fade:] * ramp_out + noise[-fade:] * (1.0 - ramp_out) + + audio[i0:i1] = noise + + return encode_audio_mono_f32_to_opus(audio, dst, sample_rate=sr, opus_bitrate=opus_bitrate) + + +def prepare_speaker_audio_for_session_mix( + audio_path: Path, + dst: Path, + *, + speech_intervals: list[tuple[float, float]], + audio_duration: float | None = None, + opus_bitrate: str = "32k", + noise_level: float = 0.0002, + seed: int | None = None, + preserve_speech: bool = True, + stitch_ms: float = 5.0, + boundary_indent: float = 0.5, +) -> bool: + """Fill non-speech (pause) regions with white noise before session mixing. + + *boundary_indent* keeps that many seconds of original audio on each side of a + speech interval untouched (pause noise starts 0.5s after speech ends and stops + 0.5s before speech begins by default). + """ + if audio_duration is None: + try: + audio_duration = ffprobe_duration(audio_path) + except RuntimeError: + audio_duration = max((end for _, end in speech_intervals), default=0.0) + 0.01 + + padded_speech = pad_speech_intervals(speech_intervals, boundary_indent, audio_duration) + pause_intervals = invert_intervals(padded_speech, audio_duration) + return apply_white_noise_in_pause_intervals( + audio_path, + dst, + pause_intervals, + opus_bitrate=opus_bitrate, + noise_level=noise_level, + seed=seed, + preserve_speech=preserve_speech, + stitch_ms=stitch_ms, + ) + + +def session_mixed_audio_path(audio_mixed_dir: Path, session_id: str) -> Path: + return audio_mixed_dir / f"{session_id}.opus" + + +def mix_audio_files(audio_paths: list[Path], output_path: Path, *, opus_bitrate: str = "32k") -> bool: + existing = [p for p in audio_paths if p.is_file()] + if not existing: + return False + output_path.parent.mkdir(parents=True, exist_ok=True) + if len(existing) == 1: + import shutil + + try: + shutil.copy2(existing[0], output_path) + return True + except OSError as exc: + log_exception(f"cannot copy mixed audio to {output_path}", exc) + return False + + cmd = [ffmpeg_executable(), "-nostdin", "-y"] + for path in existing: + cmd.extend(["-i", str(path)]) + n = len(existing) + if output_path.suffix.lower() == ".opus": + cmd.extend( + [ + "-filter_complex", + f"amix=inputs={n}:duration=longest:dropout_transition=0", + "-ac", + "1", + "-c:a", + "libopus", + "-b:a", + opus_bitrate, + "-application", + "voip", + "-vbr", + "on", + str(output_path), + ] + ) + else: + cmd.extend( + [ + "-filter_complex", + f"amix=inputs={n}:duration=longest:dropout_transition=0", + "-ac", + "1", + "-acodec", + "pcm_s16le", + str(output_path), + ] + ) + try: + result = subprocess.run( + cmd, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + timeout=FFMPEG_TIMEOUT_S, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + logger.warning("ffmpeg mix failed to start for %s: %s", output_path.name, exc) + return False + if result.returncode != 0: + logger.warning("ffmpeg mix failed for %s: %s", output_path.name, result.stderr[-400:]) + return False + return True + + +def run_main(main_fn) -> None: + """Entry-point wrapper: log tracebacks and return non-zero exit codes.""" + try: + raise SystemExit(main_fn()) + except PipelineError as exc: + logger.exception("Pipeline failed") + raise SystemExit(1) from exc + except KeyboardInterrupt: + logger.exception("Interrupted") + raise SystemExit(130) from None + except Exception as exc: + logger.exception("Unhandled error") + raise SystemExit(1) from exc diff --git a/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_manifest.py b/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_manifest.py new file mode 100644 index 0000000000..8dbc5f3b95 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_manifest.py @@ -0,0 +1,255 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Build normalized MFA rows directly from one raw session, without persisted manifests.""" + +from __future__ import annotations + +import json +import logging +import re +import unicodedata +from typing import TYPE_CHECKING + +from david_ai_common import ( + PipelineError, + log_exception, + recording_id, +) + +if TYPE_CHECKING: + from pathlib import Path + +logger = logging.getLogger(__name__) + +_DIGIT_DECADE_RE = re.compile(r"(? str: + """Insert spaces around punctuation, keeping apostrophes and hyphens word-internal.""" + text = re.sub(r"([^\w\s'-])", r" \1 ", text) + return re.sub(r"\s+", " ", text).strip() + + +def strip_digit_grouping_commas(text: str) -> str: + return _DIGIT_GROUP_COMMA_RE.sub("", text) + + +def _split_verbalized_num2words(spoken: str) -> list[str]: + spoken = _NUM2WORDS_PUNCT_RE.sub(" ", spoken) + return [word for word in spoken.split() if word] + + +def verbalize_digit_string(num_str: str, *, num2words_lang: str) -> str: + from num2words import num2words + + spoken = num2words(int(num_str), lang=num2words_lang) + return " ".join(_split_verbalized_num2words(spoken.casefold())) + + +def verbalize_decade(num_str: str, *, num2words_lang: str) -> str: + n = int(num_str) + if num2words_lang == "en": + if 0 < n < 100 and n % 10 == 0: + base = verbalize_digit_string(num_str, num2words_lang=num2words_lang) + return f"{base[:-1]}ies" if base.endswith("y") else f"{base}s" + if 1000 <= n <= 2090 and n % 10 == 0: + head = n // 100 + decade = n % 100 + if decade: + return ( + f"{verbalize_digit_string(str(head), num2words_lang=num2words_lang)} " + f"{verbalize_decade(str(decade), num2words_lang=num2words_lang)}" + ) + return f"{verbalize_digit_string(num_str, num2words_lang=num2words_lang)} s" + + +def preprocess_spoken_numbers(text: str, *, num2words_lang: str) -> str: + lang = (num2words_lang or "").strip() + if not lang: + return text + + def _decade(match: re.Match[str]) -> str: + return f" {verbalize_decade(match.group(1), num2words_lang=lang)} " + + def _feet_inches(match: re.Match[str]) -> str: + feet = verbalize_digit_string(match.group(1), num2words_lang=lang) + inches = verbalize_digit_string(match.group(2), num2words_lang=lang) + return f"{feet} {inches}" + + def _hyphen_prefix(match: re.Match[str]) -> str: + return f"{verbalize_digit_string(match.group(1), num2words_lang=lang)}{match.group(2)}" + + def _general(match: re.Match[str]) -> str: + return f" {verbalize_digit_string(match.group(0), num2words_lang=lang)} " + + text = _DIGIT_DECADE_RE.sub(_decade, text) + text = _DIGIT_FEET_INCHES_RE.sub(_feet_inches, text) + text = _DIGIT_HYPHEN_PREFIX_RE.sub(_hyphen_prefix, text) + text = _DIGIT_GENERAL_RE.sub(_general, text) + return re.sub(r"\s+", " ", text).strip() + + +def normalize_text(text: str, *, num2words_lang: str = "en") -> str: + """Normalize English transcript text using only tutorial-local helpers.""" + lang = (num2words_lang or "").strip() + try: + prepared = separate_gluing_punctuation(strip_digit_grouping_commas(text)) + prepared = preprocess_spoken_numbers(prepared, num2words_lang=lang) if lang else prepared + prepared = "".join( + character + for character in unicodedata.normalize("NFC", prepared).translate(_SMART_QUOTE_TRANSLATION) + if not unicodedata.category(character).startswith("C") + ).casefold() + rebuilt: list[str] = [] + allowed = _ENGLISH_ALPHABET | _PERMITTED_SYMBOLS + for character in prepared: + category = unicodedata.category(character) + if character in allowed: + rebuilt.append(character) + elif character.isspace() or category.startswith(("Z", "P", "S")): + rebuilt.append(" ") + else: + rebuilt.append(character) + + normalized: list[str] = [] + for token in "".join(rebuilt).split(): + if all(character in allowed for character in token): + normalized.append(token) + continue + folded = "".join( + character + for character in unicodedata.normalize("NFKD", token) + if unicodedata.category(character) != "Mn" + ) + normalized.append(folded if folded and all(character in allowed for character in folded) else "spn") + return " ".join(normalized) + except Exception as exc: + msg = f"normalization failed for text snippet: {text[:80]!r}" + raise ValueError(msg) from exc + + +def resolve_speaker_audio_path(session_dir: Path, speaker_id: str) -> Path: + """Resolve one speaker WAV using the supported filename priority.""" + candidates = ( + session_dir / f"{speaker_id}_postprocess.wav", + session_dir / f"{speaker_id}_postprocessed.wav", + session_dir / f"{speaker_id}.wav", + session_dir / f"{speaker_id}_preprocessed.wav", + ) + for candidate in candidates: + if candidate.is_file(): + return candidate.resolve() + expected = ", ".join(path.name for path in candidates) + msg = f"no speaker audio for {speaker_id}; tried: {expected}" + raise FileNotFoundError(msg) + + +def build_session_rows( + session_dir: Path, + *, + num2words_lang: str = "en", +) -> list[dict]: + """Read one raw session and create normalized rows entirely in memory.""" + session_id = session_dir.name + transcript_path = session_dir / "machine_generated_transcript.json" + if not transcript_path.is_file(): + msg = f"missing transcript: {transcript_path}" + raise FileNotFoundError(msg) + try: + with transcript_path.open(encoding="utf-8") as stream: + payload = json.load(stream) + except json.JSONDecodeError as exc: + msg = f"invalid JSON in {transcript_path}: {exc}" + raise ValueError(msg) from exc + except OSError as exc: + msg = f"cannot read {transcript_path}: {exc}" + raise PipelineError(msg) from exc + + segments = payload.get("transcript") if isinstance(payload, dict) else None + if not isinstance(segments, list): + msg = f"expected transcript list in {transcript_path}" + raise TypeError(msg) + + speaker_ids = { + str(segment["speaker"]) + for segment in segments + if isinstance(segment, dict) and segment.get("speaker") + } + norm_rows: list[dict] = [] + for speaker_id in sorted(speaker_ids): + audio_path = resolve_speaker_audio_path(session_dir, speaker_id) + rec_id = recording_id(speaker_id, session_id) + speaker_segments = [ + segment + for segment in segments + if isinstance(segment, dict) and segment.get("speaker") == speaker_id + ] + + for index, segment in enumerate(speaker_segments): + text_raw = (segment.get("text") or "").strip() + try: + start = float(segment["start"]) + end = float(segment["end"]) + except (KeyError, TypeError, ValueError) as exc: + logger.warning("%s/%s segment %d: invalid boundaries: %s", session_id, speaker_id, index, exc) + continue + if end <= start: + continue + + text_norm = "" + try: + text_norm = normalize_text(text_raw, num2words_lang=num2words_lang) if text_raw else "" + except Exception as exc: + log_exception(f"{session_id}/{speaker_id} segment {index} normalization", exc) + + row = { + "session_id": session_id, + "speaker_id": speaker_id, + "recording_id": rec_id, + "segment_index": index, + "start": start, + "end": end, + "duration": round(end - start, 6), + "text": text_norm, + "text_raw": text_raw, + "text_norm": text_norm, + "audio_filepath": str(audio_path.resolve()), + "audio_filepath_16k": str(audio_path.resolve()), + } + norm_rows.append(row) + return norm_rows diff --git a/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_mfa_align.py b/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_mfa_align.py new file mode 100644 index 0000000000..cd5f6245c7 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_mfa_align.py @@ -0,0 +1,568 @@ +"""Ephemeral MFA alignment used only by the on-the-fly RAM E2E pipeline.""" + +from __future__ import annotations + +import logging +import os +import shutil +import subprocess +import tempfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING + +from david_ai_common import ( + PipelineError, + append_jsonl, + append_mfa_g2p_args, + extract_segment_wav, + ffprobe_duration, + group_segments_by_recording, + log_exception, + map_segment_words_to_recording, + mfa_subprocess_env, + resolve_mfa_acoustic_model, + resolve_mfa_g2p_model, + run_thread_pool, + safe_parse_textgrid_words, + segment_fallback_log_entry, + session_textgrid_path, + words_to_json, + write_textgrid, +) + +if TYPE_CHECKING: + import threading + +logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO) +logger = logging.getLogger(__name__) + +# Guard against a wedged `mfa align` (e.g. SQLite/pynini lock contention seen when +# many run in a worker pool) blocking its caller forever and hanging the whole shard, +# mirroring the FFMPEG_TIMEOUT_S protection already used for ffmpeg subprocesses. +MFA_ALIGN_TIMEOUT_S = 3600 + + +def _segment_extract_workers(num_segments: int) -> int: + raw = os.environ.get("SEG_EXTRACT_WORKERS", "").strip() + if raw: + try: + return max(1, min(int(raw), num_segments)) + except ValueError: + pass + return max(1, min(num_segments, 8)) + + +def _export_segment_for_mfa( + seg: dict, + *, + audio_path: Path, + corpus_dir: Path, + segment_padding: float, + audio_duration: float, +) -> tuple[dict, Path, float] | None: + seg_idx = int(seg["segment_index"]) + seg_wav = corpus_dir / f"seg_{seg_idx:05d}.wav" + seg_txt = corpus_dir / f"seg_{seg_idx:05d}.txt" + seg_start = float(seg["start"]) + seg_end = float(seg["end"]) + extract_start = extract_segment_wav( + audio_path, + seg_wav, + seg_start, + seg_end, + padding=segment_padding, + max_duration=audio_duration, + ) + if extract_start is None: + return None + seg_txt.write_text(seg["text_norm"].strip(), encoding="utf-8") + return seg, seg_wav, extract_start + + +@dataclass +class RecordingAlignResult: + ok: bool + mfa_segments: int = 0 + fallback_segments: int = 0 + fallback_entries: list[dict] = field(default_factory=list) + merged_words: list[tuple[float, float, str]] = field(default_factory=list) + fb_words: list[tuple[float, float, str]] = field(default_factory=list) + audio_duration: float = 0.0 + + +@dataclass +class SessionAlignResult: + ok: bool + mfa_segments: int = 0 + fallback_segments: int = 0 + merged_words: list[tuple[float, float, str, str]] = field(default_factory=list) + fb_words: list[tuple[float, float, str, str]] = field(default_factory=list) + audio_duration: float = 0.0 + recordings: list[dict] = field(default_factory=list) + + +def _apply_segment_fallback( + seg: dict, + recording_id: str, + *, + reason: str, + detail: str, + fb_words: list[tuple[float, float, str]], + fallback_entries: list[dict], +) -> None: + start = float(seg["start"]) + end = float(seg["end"]) + fb_words.append((start, end, "speech")) + fallback_entries.append( + segment_fallback_log_entry( + seg, + recording_id, + reason=reason, + detail=detail, + ) + ) + logger.warning( + "%s segment %s: MFA failed (%s); using manifest speech [%.3f, %.3f]", + recording_id, + seg.get("segment_index"), + reason, + start, + end, + ) + + +def _segment_miss_or_fallback( + seg: dict, + recording_id: str, + *, + reason: str, + detail: str, + fb_words: list[tuple[float, float, str]], + fallback_entries: list[dict], + use_fallback: bool, +) -> None: + if use_fallback: + _apply_segment_fallback( + seg, + recording_id, + reason=reason, + detail=detail, + fb_words=fb_words, + fallback_entries=fallback_entries, + ) + + +def align_recording( + recording_id: str, + segments: list[dict], + *, + mfa_dict: Path, + mfa_acoustic: str, + temp_parent: Path, + num_jobs: int, + fallback_log: Path, + segment_padding: float, + fallback_log_lock: threading.Lock | None = None, + worker_mfa_root: Path | None = None, + worker_acoustic: str | None = None, + worker_g2p: str | None = None, + mfa_g2p: str | None = None, + keep_temp: bool = False, + use_fallback: bool = True, +) -> RecordingAlignResult: + temp_parent.mkdir(parents=True, exist_ok=True) + try: + if keep_temp: + temp_root = temp_parent / f"align_{recording_id}" + if temp_root.exists(): + shutil.rmtree(temp_root) + temp_root.mkdir(parents=True, exist_ok=True) + return _align_recording_impl( + recording_id, + segments, + mfa_dict=mfa_dict, + mfa_acoustic=mfa_acoustic, + temp_root=temp_root, + num_jobs=num_jobs, + fallback_log=fallback_log, + segment_padding=segment_padding, + fallback_log_lock=fallback_log_lock, + worker_mfa_root=worker_mfa_root, + worker_acoustic=worker_acoustic, + worker_g2p=worker_g2p, + mfa_g2p=mfa_g2p, + cleanup_temp=False, + use_fallback=use_fallback, + ) + with tempfile.TemporaryDirectory( + prefix=f"mfa_{recording_id}_", + dir=temp_parent, + ) as td: + return _align_recording_impl( + recording_id, + segments, + mfa_dict=mfa_dict, + mfa_acoustic=mfa_acoustic, + temp_root=Path(td), + num_jobs=num_jobs, + fallback_log=fallback_log, + segment_padding=segment_padding, + fallback_log_lock=fallback_log_lock, + worker_mfa_root=worker_mfa_root, + worker_acoustic=worker_acoustic, + worker_g2p=worker_g2p, + mfa_g2p=mfa_g2p, + cleanup_temp=False, + use_fallback=use_fallback, + ) + except Exception as exc: + log_exception(f"MFA alignment failed for {recording_id}", exc) + return RecordingAlignResult(ok=False) + + +def _align_recording_impl( + recording_id: str, + segments: list[dict], + *, + mfa_dict: Path, + mfa_acoustic: str, + temp_root: Path, + num_jobs: int, + fallback_log: Path, + segment_padding: float, + fallback_log_lock: threading.Lock | None = None, + worker_mfa_root: Path | None = None, + worker_acoustic: str | None = None, + worker_g2p: str | None = None, + mfa_g2p: str | None = None, + cleanup_temp: bool = True, + use_fallback: bool = True, +) -> RecordingAlignResult: + audio_path = Path(segments[0]["audio_filepath_16k"]) + speaker_id = segments[0]["speaker_id"] + if not audio_path.is_file(): + logger.warning("%s: missing 16k audio %s", recording_id, audio_path) + return RecordingAlignResult(ok=False) + + usable = [s for s in segments if (s.get("text_norm") or "").strip()] + if not usable: + logger.warning("%s: no segments with normalized text", recording_id) + return RecordingAlignResult(ok=False) + + try: + audio_duration = ffprobe_duration(audio_path) + except RuntimeError: + audio_duration = max(float(s["end"]) for s in usable) + 0.05 + + corpus_name = f"corpus_{recording_id}" + corpus_dir = temp_root / corpus_name / speaker_id + aligned_dir = temp_root / "aligned" + corpus_dir.mkdir(parents=True, exist_ok=True) + aligned_dir.mkdir(parents=True, exist_ok=True) + + merged_words: list[tuple[float, float, str]] = [] + fb_words: list[tuple[float, float, str]] = [] + fallback_entries: list[dict] = [] + mfa_segments = 0 + + seg_meta: list[tuple[dict, Path, float]] = [] + + def _export_one(seg: dict) -> tuple[dict, Path, float] | None: + try: + return _export_segment_for_mfa( + seg, + audio_path=audio_path, + corpus_dir=corpus_dir, + segment_padding=segment_padding, + audio_duration=audio_duration, + ) + except OSError as exc: + log_exception(f"{recording_id} segment {seg.get('segment_index')} export", exc) + return None + + extract_results = run_thread_pool( + usable, + _export_one, + workers=_segment_extract_workers(len(usable)), + ) + for seg, exported in zip(usable, extract_results, strict=True): + if exported is None: + _segment_miss_or_fallback( + seg, + recording_id, + reason="segment_export_failed", + detail="ffmpeg extract failed", + fb_words=fb_words, + fallback_entries=fallback_entries, + use_fallback=use_fallback, + ) + continue + seg_meta.append(exported) + + mfa_failed_globally = False + if seg_meta: + mfa_root = worker_mfa_root or (temp_root / "mfa_root") + if worker_mfa_root is None: + mfa_root.mkdir(parents=True, exist_ok=True) + acoustic_arg = worker_acoustic or resolve_mfa_acoustic_model(mfa_acoustic) + align_cmd = [ + "mfa", + "align", + str(corpus_dir.parent), + str(mfa_dict), + acoustic_arg, + str(aligned_dir), + ] + align_cmd.append("--clean" if worker_mfa_root is None else "--no_clean") + align_cmd.extend( + [ + "--use_mp", + "-j", + str(num_jobs), + "--beam", + "100", + "--retry_beam", + "400", + "--output_format", + "long_textgrid", + "--uses_speaker_adaptation", + "false", + "-t", + str(mfa_root), + ] + ) + g2p_arg = worker_g2p + if g2p_arg is None and mfa_g2p: + try: + g2p_arg = str(resolve_mfa_g2p_model(mfa_g2p)) + except FileNotFoundError: + logger.warning("%s: MFA G2P model not found for %r", recording_id, mfa_g2p) + append_mfa_g2p_args(align_cmd, g2p_path=g2p_arg) + logger.info("%s: running MFA on %d segments", recording_id, len(seg_meta)) + mfa_env = mfa_subprocess_env(temp_root=temp_root, mfa_root=mfa_root) + try: + result = subprocess.run( + align_cmd, + capture_output=True, + text=True, + env=mfa_env, + timeout=MFA_ALIGN_TIMEOUT_S, + check=False, + ) + except subprocess.TimeoutExpired as exc: + logger.exception("%s: mfa align timed out after %ds", recording_id, MFA_ALIGN_TIMEOUT_S) + mfa_failed_globally = True + detail = f"mfa align timed out after {MFA_ALIGN_TIMEOUT_S}s: {exc}" + except OSError as exc: + logger.exception("%s: mfa align failed to start", recording_id) + mfa_failed_globally = True + detail = str(exc) + else: + detail = result.stderr[-1200:] if result.returncode != 0 else "" + if result.returncode != 0: + logger.error( + "%s: mfa align failed (exit %d): %s", + recording_id, + result.returncode, + detail, + ) + mfa_failed_globally = True + + for seg, seg_wav, extract_start in seg_meta: + if mfa_failed_globally: + _segment_miss_or_fallback( + seg, + recording_id, + reason="mfa_align_failed", + detail=detail, + fb_words=fb_words, + fallback_entries=fallback_entries, + use_fallback=use_fallback, + ) + continue + + tg_path = aligned_dir / speaker_id / f"{seg_wav.stem}.TextGrid" + seg_end = float(seg["end"]) + extract_end = min(audio_duration, seg_end + segment_padding) + words = safe_parse_textgrid_words(tg_path) if tg_path.is_file() else [] + mapped_words = map_segment_words_to_recording( + words, + extract_start=extract_start, + extract_end=extract_end, + ) + if not mapped_words: + reason = "missing_textgrid" if not tg_path.is_file() else "empty_alignment" + _segment_miss_or_fallback( + seg, + recording_id, + reason=reason, + detail=tg_path.name, + fb_words=fb_words, + fallback_entries=fallback_entries, + use_fallback=use_fallback, + ) + continue + + merged_words.extend(mapped_words) + mfa_segments += 1 + + if not merged_words and not fb_words: + logger.warning("%s: no segment output produced", recording_id) + return RecordingAlignResult(ok=False) + + merged_words.sort(key=lambda x: x[0]) + fb_words.sort(key=lambda x: x[0]) + + if use_fallback: + for entry in fallback_entries: + try: + append_jsonl(fallback_log, entry, lock=fallback_log_lock) + except PipelineError as exc: + log_exception(f"cannot write MFA fallback log for {recording_id}", exc) + + if cleanup_temp and temp_root.exists(): + shutil.rmtree(temp_root, ignore_errors=True) + if worker_mfa_root is not None: + stale_db = worker_mfa_root / corpus_name + if stale_db.exists(): + shutil.rmtree(stale_db, ignore_errors=True) + + logger.info( + "%s: aligned %d MFA words, %d fallback segments", + recording_id, + len(merged_words), + len(fallback_entries), + ) + return RecordingAlignResult( + ok=True, + mfa_segments=mfa_segments, + fallback_segments=len(fallback_entries), + fallback_entries=fallback_entries, + merged_words=merged_words, + fb_words=fb_words, + audio_duration=audio_duration, + ) + + +def align_session( + session_id: str, + segments: list[dict], + *, + mfa_dict: Path, + mfa_acoustic: str, + textgrid_dir: Path, + temp_parent: Path, + num_jobs: int, + fallback_log: Path, + segment_padding: float, + fallback_log_lock: threading.Lock | None = None, + worker_mfa_root: Path | None = None, + worker_acoustic: str | None = None, + worker_g2p: str | None = None, + mfa_g2p: str | None = None, + keep_temp: bool = False, + use_fallback: bool = True, + write_textgrids: bool = True, +) -> SessionAlignResult: + by_recording = group_segments_by_recording(segments) + session_merged: list[tuple[float, float, str, str]] = [] + session_fb: list[tuple[float, float, str, str]] = [] + recording_rows: list[dict] = [] + session_duration = 0.0 + mfa_segments = 0 + fallback_segments = 0 + + for rec_id, rec_segments in sorted(by_recording.items()): + result = align_recording( + rec_id, + rec_segments, + mfa_dict=mfa_dict, + mfa_acoustic=mfa_acoustic, + temp_parent=temp_parent, + num_jobs=num_jobs, + fallback_log=fallback_log, + segment_padding=segment_padding, + fallback_log_lock=fallback_log_lock, + worker_mfa_root=worker_mfa_root, + worker_acoustic=worker_acoustic, + worker_g2p=worker_g2p, + mfa_g2p=mfa_g2p, + keep_temp=keep_temp, + use_fallback=use_fallback, + ) + if not result.ok: + logger.warning("%s: speaker recording %s failed", session_id, rec_id) + continue + + speaker_id = rec_segments[0]["speaker_id"] + for start, end, word in result.merged_words: + session_merged.append((start, end, word, speaker_id)) + for start, end, word in result.fb_words: + session_fb.append((start, end, word, speaker_id)) + session_duration = max(session_duration, result.audio_duration) + mfa_segments += result.mfa_segments + fallback_segments += result.fallback_segments + recording_rows.append( + { + "recording_id": rec_id, + "speaker_id": speaker_id, + "session_id": session_id, + "audio_filepath_16k": rec_segments[0]["audio_filepath_16k"], + "audio_duration": result.audio_duration, + "merged_words": words_to_json(result.merged_words), + "fb_words": words_to_json(result.fb_words), + } + ) + + if not session_merged and not session_fb: + logger.warning("%s: no session alignment output", session_id) + return SessionAlignResult(ok=False) + + session_merged.sort(key=lambda x: x[0]) + session_fb.sort(key=lambda x: x[0]) + max_seg_end = max((end for _, end, _, _ in session_merged + session_fb), default=0.0) + xmax = max(session_duration, max_seg_end) + 0.01 + + fastmss_words = [(s, e, w) for s, e, w, _ in session_merged] + ordinary_words = sorted( + [(s, e, w) for s, e, w, _ in session_merged] + [(s, e, w) for s, e, w, _ in session_fb], + key=lambda x: x[0], + ) + if write_textgrids: + write_textgrid( + fastmss_words, + session_textgrid_path(textgrid_dir, session_id, variant="fastmss"), + xmin=0.0, + xmax=xmax, + ) + write_textgrid( + ordinary_words, + session_textgrid_path(textgrid_dir, session_id, variant="ordinary"), + xmin=0.0, + xmax=xmax, + ) + logger.info( + "%s: session TextGrids (%d MFA words, %d fallback, %d speakers)", + session_id, + len(session_merged), + len(session_fb), + len(recording_rows), + ) + else: + logger.info( + "%s: MFA alignment (%d words, %d fallback, %d speakers; TextGrids skipped)", + session_id, + len(session_merged), + len(session_fb), + len(recording_rows), + ) + return SessionAlignResult( + ok=True, + mfa_segments=mfa_segments, + fallback_segments=fallback_segments, + merged_words=session_merged, + fb_words=session_fb, + audio_duration=xmax, + recordings=recording_rows, + ) diff --git a/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_ram_lhotse.py b/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_ram_lhotse.py new file mode 100644 index 0000000000..181a83319e --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_ram_lhotse.py @@ -0,0 +1,92 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TextGrid writers used by the on-the-fly RAM E2E pipeline.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from david_ai_common import ( + fastmss_textgrid_path, + recording_textgrid_path, + session_textgrid_path, + words_from_json, + write_textgrid, +) + +if TYPE_CHECKING: + from pathlib import Path + + from david_ai_mfa_align import SessionAlignResult + + +def write_recording_textgrids(align_result: SessionAlignResult, textgrid_dir: Path) -> None: + """Write ordinary and FastMSS TextGrids for every speaker recording.""" + textgrid_dir.mkdir(parents=True, exist_ok=True) + for rec_row in align_result.recordings: + rec_id = rec_row["recording_id"] + merged = words_from_json(rec_row["merged_words"]) + fallback = words_from_json(rec_row["fb_words"]) + duration = float(rec_row.get("audio_duration", 0.0)) + max_word_end = max((end for _, end, _ in merged + fallback), default=0.0) + xmax = max(duration, max_word_end) + 0.01 + + write_textgrid( + merged, + fastmss_textgrid_path(textgrid_dir, rec_id), + xmin=0.0, + xmax=xmax, + ) + write_textgrid( + sorted(merged + fallback, key=lambda word: word[0]), + recording_textgrid_path(textgrid_dir, rec_id), + xmin=0.0, + xmax=xmax, + ) + + +def write_session_textgrids(align_result: SessionAlignResult, textgrid_dir: Path) -> None: + """Write ordinary and FastMSS TextGrids for the mixed session timeline.""" + session_id = align_result.recordings[0]["session_id"] if align_result.recordings else "" + if not session_id: + return + + fastmss_words = [(start, end, word) for start, end, word, _ in align_result.merged_words] + ordinary_words = sorted( + fastmss_words + [(start, end, word) for start, end, word, _ in align_result.fb_words], + key=lambda word: word[0], + ) + max_word_end = max((end for _, end, _ in ordinary_words), default=0.0) + xmax = max(float(align_result.audio_duration), max_word_end) + 0.01 + textgrid_dir.mkdir(parents=True, exist_ok=True) + + write_textgrid( + fastmss_words, + session_textgrid_path(textgrid_dir, session_id, variant="fastmss"), + xmin=0.0, + xmax=xmax, + ) + write_textgrid( + ordinary_words, + session_textgrid_path(textgrid_dir, session_id, variant="ordinary"), + xmin=0.0, + xmax=xmax, + ) + + +def write_all_textgrids(align_result: SessionAlignResult, textgrid_dir: Path) -> None: + """Persist session and per-recording TextGrids in both required variants.""" + write_session_textgrids(align_result, textgrid_dir) + write_recording_textgrids(align_result, textgrid_dir) diff --git a/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_ram_session.py b/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_ram_session.py new file mode 100644 index 0000000000..9e177520c9 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/opus/david_ai_ram_session.py @@ -0,0 +1,407 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-session worker for the strict on-the-fly RAM E2E pipeline.""" + +from __future__ import annotations + +import os +import shutil +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from david_ai_common import ( + PipelineError, + build_session_rttm_lines_from_words, + fastmss_textgrid_path, + ffprobe_duration, + group_recordings_by_session, + group_segments_by_recording, + log_exception, + mix_audio_files, + mixed_speaker_audio_path, + prepare_speaker_audio_for_session_mix, + recording_id, + recording_textgrid_path, + run_thread_pool, + session_mixed_audio_path, + session_rttm_path, + session_textgrid_path, + setup_mfa_worker_root, + write_rttm, +) +from david_ai_manifest import build_session_rows +from david_ai_mfa_align import align_session +from david_ai_ram_lhotse import write_all_textgrids + +if TYPE_CHECKING: + from pathlib import Path + +_PROCESS_MFA: dict | None = None + + +def _lazy_mfa_worker( + *, + ram_dir: Path, + mfa_dict: Path, + mfa_acoustic: str, + mfa_g2p: str, +) -> tuple[Path, Path, str, str | None, Path]: + """Initialize one ephemeral MFA model/database root per process.""" + global _PROCESS_MFA + if _PROCESS_MFA is None: + worker_dir = ram_dir / "mfa_workers" / f"worker_{os.getpid()}" + worker_dir.mkdir(parents=True, exist_ok=True) + mfa_root, local_dict, acoustic_arg, g2p_arg = setup_mfa_worker_root( + worker_dir, + mfa_dict=mfa_dict, + mfa_acoustic=mfa_acoustic, + mfa_g2p=mfa_g2p, + ) + _PROCESS_MFA = { + "mfa_root": mfa_root, + "local_dict": local_dict, + "acoustic_arg": acoustic_arg, + "g2p_arg": g2p_arg, + "temp_parent": worker_dir / "align_temp", + } + cfg = _PROCESS_MFA + return cfg["mfa_root"], cfg["local_dict"], cfg["acoustic_arg"], cfg["g2p_arg"], cfg["temp_parent"] + + +@dataclass +class SessionRamResult: + session_id: str + ok: bool + error: str = "" + + +def session_done_path(work_dir: Path, session_id: str) -> Path: + return work_dir / ".done" / "sessions" / f"{session_id}.done" + + +def is_session_done(work_dir: Path, session_id: str) -> bool: + """Return whether a previous run validated every required session output.""" + return session_done_path(work_dir, session_id).is_file() + + +def _clear_session_done(work_dir: Path, session_id: str) -> None: + session_done_path(work_dir, session_id).unlink(missing_ok=True) + + +def _nonempty(path: Path) -> bool: + return path.is_file() and path.stat().st_size > 0 + + +def _validate_session_outputs( + session_id: str, + norm_rows: list[dict], + *, + audio_mixed_dir: Path, + textgrid_dir: Path, +) -> None: + """Require every declared deliverable before writing the session success flag.""" + required = [ + session_mixed_audio_path(audio_mixed_dir, session_id), + session_rttm_path(audio_mixed_dir, session_id), + session_textgrid_path(textgrid_dir, session_id, variant="ordinary"), + session_textgrid_path(textgrid_dir, session_id, variant="fastmss"), + ] + speaker_ids = sorted({row["speaker_id"] for row in norm_rows}) + for speaker_id in speaker_ids: + rec_id = recording_id(speaker_id, session_id) + required.extend( + [ + mixed_speaker_audio_path(audio_mixed_dir, speaker_id, session_id), + recording_textgrid_path(textgrid_dir, rec_id, variant="ordinary"), + fastmss_textgrid_path(textgrid_dir, rec_id), + ] + ) + missing = [str(path) for path in required if not _nonempty(path)] + if missing: + msg = f"missing or empty session outputs: {missing}" + raise PipelineError(msg) + + +def _mark_session_done(work_dir: Path, session_id: str) -> None: + path = session_done_path(work_dir, session_id) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("ok\n", encoding="utf-8") + + +def _finalize_session_success( + session_id: str, + norm_rows: list[dict], + *, + work_dir: Path, + audio_mixed_dir: Path, + textgrid_dir: Path, +) -> None: + _validate_session_outputs( + session_id, + norm_rows, + audio_mixed_dir=audio_mixed_dir, + textgrid_dir=textgrid_dir, + ) + _mark_session_done(work_dir, session_id) + + +def _mix_prep_workers(num_speakers: int) -> int: + raw = os.environ.get("MIX_PREP_WORKERS", "").strip() + if raw: + try: + return max(1, min(int(raw), num_speakers)) + except ValueError: + pass + return max(1, num_speakers) + + +def _manifest_speech_intervals_by_recording(norm_rows: list[dict]) -> dict[str, list[tuple[float, float]]]: + """Return exact original manifest boundaries, before the 0.5-second protection offset.""" + return { + rec_id: [(float(row["start"]), float(row["end"])) for row in rows] + for rec_id, rows in group_segments_by_recording(norm_rows).items() + } + + +def _prepare_speaker_tracks_for_mix( + entries: list[dict], + *, + session_id: str, + session_scratch: Path, + audio_mixed_dir: Path, + manifest_speech: dict[str, list[tuple[float, float]]], + rec_durations: dict[str, float], + opus_bitrate: str, + noise_level: float, + stitch_ms: float, + boundary_offset: float, +) -> list[tuple[Path, Path]]: + """Replace pauses using original manifest masks and prepare persistent track destinations.""" + specs: list[tuple[Path, Path, Path, list[tuple[float, float]], float, int, str]] = [] + for entry in entries: + rec_id = entry["recording_id"] + speaker_id = entry["speaker_id"] + src = entry["audio_path"] + if not src.is_file(): + msg = f"missing source audio {src}" + raise FileNotFoundError(msg) + + speech = manifest_speech.get(rec_id) + if not speech: + msg = f"no original manifest boundaries for {rec_id}" + raise PipelineError(msg) + duration = rec_durations.get(rec_id, 0.0) + if duration <= 0: + try: + duration = ffprobe_duration(src) + except RuntimeError: + duration = max(end for _, end in speech) + 0.01 + + local_dst = session_scratch / f"{rec_id}.opus" + persistent_dst = mixed_speaker_audio_path(audio_mixed_dir, speaker_id, session_id) + seed = hash((session_id, rec_id)) & 0xFFFFFFFF + specs.append((src, local_dst, persistent_dst, speech, duration, seed, rec_id)) + + def _prepare_one( + spec: tuple[Path, Path, Path, list[tuple[float, float]], float, int, str], + ) -> tuple[Path, Path]: + src, local_dst, persistent_dst, speech, duration, seed, rec_id = spec + if not prepare_speaker_audio_for_session_mix( + src, + local_dst, + speech_intervals=speech, + audio_duration=duration, + opus_bitrate=opus_bitrate, + noise_level=noise_level, + seed=seed, + preserve_speech=True, + stitch_ms=stitch_ms, + boundary_indent=boundary_offset, + ): + msg = f"pause noise prep failed for {rec_id}" + raise PipelineError(msg) + return local_dst, persistent_dst + + return run_thread_pool(specs, _prepare_one, workers=_mix_prep_workers(len(specs))) + + +def _publish_audio(local_path: Path, output_path: Path) -> None: + """Publish completed local audio without exposing a partial final path.""" + output_path.parent.mkdir(parents=True, exist_ok=True) + temp_path = output_path.with_name(f".{output_path.name}.{os.getpid()}.tmp") + try: + shutil.copyfile(local_path, temp_path) + os.replace(temp_path, output_path) + finally: + if temp_path.is_file(): + temp_path.unlink() + + +def _mix_session_from_manifest( + session_id: str, + norm_rows: list[dict], + *, + audio_mixed_dir: Path, + session_ram: Path, + opus_bitrate: str, + noise_level: float, + stitch_ms: float, + boundary_offset: float, + rec_durations: dict[str, float], +) -> None: + entries = group_recordings_by_session(norm_rows).get(session_id, []) + if not entries: + msg = "no speaker recordings to mix" + raise PipelineError(msg) + + session_scratch = session_ram / "mix" + if session_scratch.exists(): + shutil.rmtree(session_scratch, ignore_errors=True) + session_scratch.mkdir(parents=True, exist_ok=True) + + prepared_tracks = _prepare_speaker_tracks_for_mix( + entries, + session_id=session_id, + session_scratch=session_scratch, + audio_mixed_dir=audio_mixed_dir, + manifest_speech=_manifest_speech_intervals_by_recording(norm_rows), + rec_durations=rec_durations, + opus_bitrate=opus_bitrate, + noise_level=noise_level, + stitch_ms=stitch_ms, + boundary_offset=boundary_offset, + ) + local_mixed = session_scratch / f"{session_id}.opus" + if not mix_audio_files([local for local, _ in prepared_tracks], local_mixed, opus_bitrate=opus_bitrate): + msg = "session mix failed" + raise PipelineError(msg) + + for local_path, persistent_path in prepared_tracks: + _publish_audio(local_path, persistent_path) + _publish_audio(local_mixed, session_mixed_audio_path(audio_mixed_dir, session_id)) + + +def process_session_ram( + session_dir: Path, + *, + work_dir: Path, + audio_mixed_dir: Path, + textgrid_dir: Path, + mfa_dict: Path, + mfa_acoustic: str, + mfa_g2p: str, + ram_dir: Path, + num2words_lang: str = "en", + mfa_num_jobs: int = 2, + segment_padding: float = 0.5, + rttm_merge_gap: float = 0.2, + opus_bitrate: str = "32k", + noise_level: float = 0.0002, + stitch_ms: float = 5.0, + boundary_offset: float = 0.5, +) -> SessionRamResult: + """Run every E2E step from raw transcript/WAV, without reading persisted pipeline state.""" + session_id = session_dir.name + session_ram = ram_dir / "sessions" / session_id + _clear_session_done(work_dir, session_id) + + try: + norm_rows = build_session_rows( + session_dir, + num2words_lang=num2words_lang, + ) + if not norm_rows: + return SessionRamResult(session_id=session_id, ok=False, error="no manifest rows") + + if session_ram.exists(): + shutil.rmtree(session_ram, ignore_errors=True) + session_ram.mkdir(parents=True, exist_ok=True) + textgrid_dir.mkdir(parents=True, exist_ok=True) + audio_mixed_dir.mkdir(parents=True, exist_ok=True) + + worker_mfa_root, local_dict, acoustic_arg, g2p_arg, temp_parent = _lazy_mfa_worker( + ram_dir=ram_dir, + mfa_dict=mfa_dict, + mfa_acoustic=mfa_acoustic, + mfa_g2p=mfa_g2p, + ) + temp_parent.mkdir(parents=True, exist_ok=True) + align_result = align_session( + session_id, + norm_rows, + mfa_dict=local_dict, + mfa_acoustic=mfa_acoustic, + textgrid_dir=textgrid_dir, + temp_parent=temp_parent / session_id, + num_jobs=mfa_num_jobs, + fallback_log=session_ram / "fallback.jsonl", + segment_padding=segment_padding, + worker_mfa_root=worker_mfa_root, + worker_acoustic=acoustic_arg, + worker_g2p=g2p_arg, + mfa_g2p=mfa_g2p, + keep_temp=False, + use_fallback=True, + write_textgrids=False, + ) + if not align_result.ok: + return SessionRamResult(session_id=session_id, ok=False, error="MFA alignment failed") + if align_result.mfa_segments == 0: + return SessionRamResult( + session_id=session_id, + ok=False, + error="MFA produced zero aligned segments; refusing fallback-only completion", + ) + + rttm_lines = build_session_rttm_lines_from_words( + session_id, + align_result.merged_words, + align_result.fb_words, + merge_gap=rttm_merge_gap, + ) + if not rttm_lines: + return SessionRamResult(session_id=session_id, ok=False, error="empty session RTTM") + write_rttm(session_rttm_path(audio_mixed_dir, session_id), rttm_lines) + + write_all_textgrids(align_result, textgrid_dir) + + rec_durations = { + rec["recording_id"]: float(rec.get("audio_duration", 0.0)) + for rec in align_result.recordings + } + _mix_session_from_manifest( + session_id, + norm_rows, + audio_mixed_dir=audio_mixed_dir, + session_ram=session_ram, + opus_bitrate=opus_bitrate, + noise_level=noise_level, + stitch_ms=stitch_ms, + boundary_offset=boundary_offset, + rec_durations=rec_durations, + ) + _finalize_session_success( + session_id, + norm_rows, + work_dir=work_dir, + audio_mixed_dir=audio_mixed_dir, + textgrid_dir=textgrid_dir, + ) + return SessionRamResult(session_id=session_id, ok=True) + except Exception as exc: + log_exception(f"RAM session pipeline failed for {session_id}", exc) + return SessionRamResult(session_id=session_id, ok=False, error=str(exc)) + finally: + shutil.rmtree(session_ram, ignore_errors=True) diff --git a/tutorials/audio/david_ai_redelivered_mfa/opus/requirements.txt b/tutorials/audio/david_ai_redelivered_mfa/opus/requirements.txt new file mode 100644 index 0000000000..3c8d7e7822 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/opus/requirements.txt @@ -0,0 +1 @@ +-r ../requirements.txt diff --git a/tutorials/audio/david_ai_redelivered_mfa/opus/run_david_ai_mfa_ram_session.sh b/tutorials/audio/david_ai_redelivered_mfa/opus/run_david_ai_mfa_ram_session.sh new file mode 100755 index 0000000000..66975b7eb2 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/opus/run_david_ai_mfa_ram_session.sh @@ -0,0 +1,104 @@ +#!/bin/bash +# Strict on-the-fly David AI E2E pipeline. +# +# Every unfinished session starts from raw WAVs + machine_generated_transcript.json: +# normalize in memory -> MFA with base dictionary + runtime G2P -> RTTM -> +# manifest-mask pause-noise Opus -> mixed Opus -> ordinary and FastMSS TextGrids. +# No original per-speaker audio copy, persisted manifests, shared lexicon, +# or partial output cache is read. A validated session.done flag skips completed +# sessions and is written only after every required output passes validation. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CURATOR_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)" + +DATA_ROOT="${DATA_ROOT:?Set DATA_ROOT to the raw-session directory}" +WORK_DIR="${WORK_DIR:-$SCRIPT_DIR/workdir_e2e}" +TEXTGRID_DIR="${TEXTGRID_DIR:-$WORK_DIR/textgrids}" +AUDIO_MIXED_DIR="${AUDIO_MIXED_DIR:-$WORK_DIR/audio_mixed}" +LOG_DIR="${LOG_DIR:-$WORK_DIR/logs}" +RAM_DIR="${RAM_DIR:-/tmp/david_ai_ram_session_${SLURM_JOB_ID:-$$}_${SLURM_ARRAY_TASK_ID:-0}}" + +MFA_ROOT_DIR="${MFA_ROOT_DIR:-$HOME/MFA_models}" +MFA_DICT_NAME="${MFA_DICT_NAME:-english_us_arpa}" +MFA_G2P="${MFA_G2P:-english_us_arpa}" +MFA_ACOUSTIC="${MFA_ACOUSTIC:-english_us_arpa}" +MFA_NUM_JOBS="${MFA_NUM_JOBS:-2}" +WORKERS="${WORKERS:-4}" +SEGMENT_PADDING="${SEGMENT_PADDING:-0.5}" +RTTM_MERGE_GAP="${RTTM_MERGE_GAP:-0.2}" +OPUS_BITRATE="${OPUS_BITRATE:-32k}" +NUM2WORDS_LANG="${NUM2WORDS_LANG:-en}" +SESSIONS_FILE="${SESSIONS_FILE:-}" +SHARD_COUNT="${SHARD_COUNT:-1}" +SHARD_INDEX="${SHARD_INDEX:-0}" + +PYTHON="${PYTHON:-}" +if [[ -z "$PYTHON" && -x "$CURATOR_ROOT/.venv/bin/python" ]]; then + PYTHON="$CURATOR_ROOT/.venv/bin/python" +fi +PYTHON="${PYTHON:-python3}" +MFA_ENV="${MFA_ENV:-$HOME/miniconda3/envs/curator_pain_1}" +if [[ -x "$MFA_ENV/bin/mfa" ]]; then + if [[ -n "${FFMPEG_BIN:-}" && -x "$FFMPEG_BIN" ]]; then + export PATH="${PATH}:$MFA_ENV/bin" + else + export PATH="$MFA_ENV/bin:$PATH" + fi +fi +export MFA_ROOT_DIR + +if [[ -n "${FFMPEG_BIN:-}" && -x "$FFMPEG_BIN" ]]; then + : +elif ! command -v ffmpeg >/dev/null 2>&1; then + echo "ERROR: ffmpeg not on PATH" >&2 + exit 1 +fi +if ! command -v mfa >/dev/null 2>&1; then + echo "ERROR: mfa not on PATH (MFA_ENV=$MFA_ENV)" >&2 + exit 1 +fi +if [[ ! -d "$DATA_ROOT" ]]; then + echo "ERROR: data root does not exist: $DATA_ROOT" >&2 + exit 1 +fi + +mkdir -p "$LOG_DIR" "$TEXTGRID_DIR" "$AUDIO_MIXED_DIR" "$RAM_DIR" +RUN_ID="$(date +%Y%m%d_%H%M%S)_${SLURM_JOB_ID:-local}_${SLURM_ARRAY_TASK_ID:-0}" +LOG_FILE="$LOG_DIR/run_e2e_${RUN_ID}.log" +exec > >(tee -a "$LOG_FILE") 2>&1 + +CMD=( + "$PYTHON" "$SCRIPT_DIR/stage_ram_session_pipeline.py" + --data-root "$DATA_ROOT" + --work-dir "$WORK_DIR" + --audio-mixed-dir "$AUDIO_MIXED_DIR" + --textgrid-dir "$TEXTGRID_DIR" + --mfa-dict-name "$MFA_DICT_NAME" + --mfa-acoustic "$MFA_ACOUSTIC" + --mfa-g2p "$MFA_G2P" + --ram-dir "$RAM_DIR" + --num2words-lang "$NUM2WORDS_LANG" + --mfa-num-jobs "$MFA_NUM_JOBS" + --segment-padding "$SEGMENT_PADDING" + --rttm-merge-gap "$RTTM_MERGE_GAP" + --opus-bitrate "$OPUS_BITRATE" + --noise-level 0.0002 + --stitch-ms 5 + --boundary-offset 0.5 + --workers "$WORKERS" + --shard-count "$SHARD_COUNT" + --shard-index "$SHARD_INDEX" +) +[[ -n "$SESSIONS_FILE" ]] && CMD+=(--sessions-file "$SESSIONS_FILE") + +echo "[$(date '+%Y-%m-%d %H:%M:%S')] ON-THE-FLY E2E START" +echo "DATA_ROOT=$DATA_ROOT" +echo "WORK_DIR=$WORK_DIR" +echo "WORKERS=$WORKERS MFA_NUM_JOBS=$MFA_NUM_JOBS" +echo "MFA dictionary=$MFA_DICT_NAME runtime_g2p=$MFA_G2P" +echo "Pause mask=original manifest boundaries +/-0.5s, noise=0.0002, smoothing=5ms" +echo "Command: ${CMD[*]}" +"${CMD[@]}" +echo "[$(date '+%Y-%m-%d %H:%M:%S')] ON-THE-FLY E2E DONE" diff --git a/tutorials/audio/david_ai_redelivered_mfa/opus/stage_ram_session_pipeline.py b/tutorials/audio/david_ai_redelivered_mfa/opus/stage_ram_session_pipeline.py new file mode 100755 index 0000000000..b497015231 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/opus/stage_ram_session_pipeline.py @@ -0,0 +1,242 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Strict on-the-fly E2E: raw sessions -> MFA/G2P -> audio, RTTM, and TextGrids.""" + +from __future__ import annotations + +import argparse +import logging +import shutil +import tempfile +from concurrent.futures import ProcessPoolExecutor, as_completed +from dataclasses import dataclass +from pathlib import Path + +from david_ai_common import PipelineError, discover_sessions, resolve_mfa_dict, run_main +from david_ai_ram_session import SessionRamResult, is_session_done, process_session_ram + +logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO) +logger = logging.getLogger(__name__) + + +@dataclass +class RamSessionTask: + session_dir: str + work_dir: str + audio_mixed_dir: str + textgrid_dir: str + mfa_dict: str + mfa_acoustic: str + mfa_g2p: str + ram_dir: str + num2words_lang: str + mfa_num_jobs: int + segment_padding: float + rttm_merge_gap: float + opus_bitrate: str + noise_level: float + stitch_ms: float + boundary_offset: float + + +def _run_session_task(task: RamSessionTask) -> SessionRamResult: + return process_session_ram( + Path(task.session_dir), + work_dir=Path(task.work_dir), + audio_mixed_dir=Path(task.audio_mixed_dir), + textgrid_dir=Path(task.textgrid_dir), + mfa_dict=Path(task.mfa_dict), + mfa_acoustic=task.mfa_acoustic, + mfa_g2p=task.mfa_g2p, + ram_dir=Path(task.ram_dir), + num2words_lang=task.num2words_lang, + mfa_num_jobs=task.mfa_num_jobs, + segment_padding=task.segment_padding, + rttm_merge_gap=task.rttm_merge_gap, + opus_bitrate=task.opus_bitrate, + noise_level=task.noise_level, + stitch_ms=task.stitch_ms, + boundary_offset=task.boundary_offset, + ) + + +def sessions_without_done_flags(sessions: list[Path], work_dir: Path) -> list[Path]: + """Select only sessions that have not completed successfully.""" + return [session for session in sessions if not is_session_done(work_dir, session.name)] + + +def filter_sessions_from_file(sessions: list[Path], sessions_file: Path) -> list[Path]: + """Restrict discovered sessions to IDs listed one per line.""" + requested = { + line.strip() + for line in sessions_file.read_text(encoding="utf-8").splitlines() + if line.strip() and not line.lstrip().startswith("#") + } + available = {session.name for session in sessions} + missing = sorted(requested - available) + if missing: + logger.warning("%d requested session IDs were not found under DATA_ROOT", len(missing)) + return [session for session in sessions if session.name in requested] + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--data-root", type=Path, required=True) + ap.add_argument("--work-dir", type=Path, required=True) + ap.add_argument("--audio-mixed-dir", type=Path, default=None) + ap.add_argument("--textgrid-dir", type=Path, default=None) + ap.add_argument("--mfa-dict-name", default="english_us_arpa") + ap.add_argument("--mfa-acoustic", default="english_us_arpa") + ap.add_argument("--mfa-g2p", default="english_us_arpa") + ap.add_argument( + "--ram-dir", + type=Path, + default=Path(tempfile.gettempdir()) / "david_ai_ram_session", + ) + ap.add_argument("--num2words-lang", default="en") + ap.add_argument("--mfa-num-jobs", type=int, default=2) + ap.add_argument("--segment-padding", type=float, default=0.5) + ap.add_argument("--rttm-merge-gap", type=float, default=0.2) + ap.add_argument("--opus-bitrate", default="32k") + ap.add_argument("--noise-level", type=float, default=0.0002) + ap.add_argument("--stitch-ms", type=float, default=5.0) + ap.add_argument("--boundary-offset", type=float, default=0.5) + ap.add_argument("--workers", type=int, default=4) + ap.add_argument("--sessions-file", type=Path, default=None) + ap.add_argument("--shard-count", type=int, default=1) + ap.add_argument("--shard-index", type=int, default=0) + args = ap.parse_args() + + if args.shard_count < 1: + msg = f"--shard-count must be >= 1, got {args.shard_count}" + raise PipelineError(msg) + if not 0 <= args.shard_index < args.shard_count: + msg = f"--shard-index must be in [0, {args.shard_count}), got {args.shard_index}" + raise PipelineError( + msg + ) + + work_dir = args.work_dir.resolve() + data_root = args.data_root.resolve() + audio_mixed_dir = (args.audio_mixed_dir or work_dir / "audio_mixed").resolve() + textgrid_dir = (args.textgrid_dir or work_dir / "textgrids").resolve() + ram_dir = args.ram_dir.resolve() + for path in (audio_mixed_dir, textgrid_dir, ram_dir): + path.mkdir(parents=True, exist_ok=True) + + mfa_dict_path = resolve_mfa_dict(args.mfa_dict_name) + logger.info( + "Using base MFA dictionary %s with runtime G2P model %s", + mfa_dict_path, + args.mfa_g2p, + ) + + sessions = discover_sessions(data_root) + if args.sessions_file is not None: + sessions_file = args.sessions_file.resolve() + if not sessions_file.is_file(): + msg = f"sessions file does not exist: {sessions_file}" + raise PipelineError(msg) + sessions = filter_sessions_from_file(sessions, sessions_file) + logger.info("Restricted run to %d sessions from %s", len(sessions), sessions_file) + if not sessions: + msg = f"No sessions under {data_root}" + raise PipelineError(msg) + if args.shard_count > 1: + total = len(sessions) + sessions = [ + session + for index, session in enumerate(sessions) + if index % args.shard_count == args.shard_index + ] + logger.info( + "Shard %d/%d: processing %d of %d raw sessions", + args.shard_index, + args.shard_count, + len(sessions), + total, + ) + + pending_sessions = sessions_without_done_flags(sessions, work_dir) + skipped_sessions = len(sessions) - len(pending_sessions) + workers = max(1, args.workers) + logger.info( + "Resumable E2E START: sessions=%d pending=%d done=%d workers=%d mfa_jobs=%d ram_dir=%s", + len(sessions), + len(pending_sessions), + skipped_sessions, + workers, + args.mfa_num_jobs, + ram_dir, + ) + if not pending_sessions: + logger.info("All %d selected sessions already have validated done flags", len(sessions)) + return 0 + tasks = [ + RamSessionTask( + session_dir=str(session.resolve()), + work_dir=str(work_dir), + audio_mixed_dir=str(audio_mixed_dir), + textgrid_dir=str(textgrid_dir), + mfa_dict=str(mfa_dict_path), + mfa_acoustic=args.mfa_acoustic, + mfa_g2p=args.mfa_g2p, + ram_dir=str(ram_dir), + num2words_lang=args.num2words_lang, + mfa_num_jobs=args.mfa_num_jobs, + segment_padding=args.segment_padding, + rttm_merge_gap=args.rttm_merge_gap, + opus_bitrate=args.opus_bitrate, + noise_level=args.noise_level, + stitch_ms=args.stitch_ms, + boundary_offset=args.boundary_offset, + ) + for session in pending_sessions + ] + + ok = fail = completed = 0 + with ProcessPoolExecutor(max_workers=workers) as pool: + futures = [pool.submit(_run_session_task, task) for task in tasks] + for future in as_completed(futures): + result = future.result() + completed += 1 + if result.ok: + ok += 1 + else: + fail += 1 + logger.warning("%s failed: %s", result.session_id, result.error) + if completed % 50 == 0 or completed == len(futures): + logger.info( + "E2E progress: %d/%d (ok=%d fail=%d)", + completed, + len(futures), + ok, + fail, + ) + + shutil.rmtree(ram_dir, ignore_errors=True) + logger.info( + "Resumable E2E DONE: ok=%d fail=%d previously_done=%d workers=%d", + ok, + fail, + skipped_sessions, + workers, + ) + return 0 if fail == 0 and ok == len(pending_sessions) else 1 + + +if __name__ == "__main__": + run_main(main) diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/README.md b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/README.md new file mode 100644 index 0000000000..b9c9466909 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/README.md @@ -0,0 +1,230 @@ +# Parakeet segment WER and FastMSS manifest pipeline + +This Curator pipeline evaluates completed masked per-speaker WAVs with NVIDIA +Parakeet and creates filtered per-speaker training manifests. + +## Flow + +1. Read ground-truth segments from each session transcript manifest. +2. Resolve the matching `audio_16k_masked/.wav`. +3. Read recording-global words from `_fastmss.TextGrid`. +4. Extract each exact ground-truth segment into node-local scratch. + Clips shorter than 100 ms are context-padded to 100 ms for stable Parakeet + features; manifest start/end values remain unchanged. +5. Run `nvidia/parakeet-tdt-0.6b-v2` with batched GPU inference. +6. Normalize reference and hypothesis text identically. +7. Compute segment-level substitutions, deletions, insertions, and WER percent. +8. Delete the temporary segment WAV immediately after scoring. +9. Generate WER histograms, percentiles, a Tukey upper fence, and a proposed threshold. +10. Exclude segments above the applied threshold (100% by default). +11. Optionally exclude segments without FastMSS word alignment (enabled by default). +12. Write one filtered JSONL manifest per masked speaker recording. +13. Build exact-0%, 0–10%, and 0–100% Lhotse CutSets for the shard. + +## Requirements + +- Completed WAV MFA pipeline outputs: + - `audio_16k_masked/*.wav` + - `textgrids/*_fastmss.TextGrid` +- Raw session transcript manifests +- NVIDIA GPU +- `nvidia-ml-py` so Xenna can discover local GPUs +- NeMo ASR / Parakeet environment +- Curator installed from this repository + +Install the additional environment dependencies: + +```bash +python -m pip install -r parakeet_wer/requirements.txt +``` + +The Parakeet model may require approved outbound access on first use. Set +`MODEL_CACHE_DIR` to a persistent pre-populated model cache for offline runs. + +## Run one shard + +```bash +cd tutorials/audio/david_ai_redelivered_mfa/parakeet_wer + +DATA_ROOT=/path/to/raw/sessions \ +MFA_WORK_DIR=/path/to/completed/wav-mfa-workdir \ +OUTPUT_DIR=/path/to/parakeet-wer-output \ +SESSIONS_FILE=/path/to/session-list.txt \ +PARAKEET_MODEL=nvidia/parakeet-tdt-0.6b-v2 \ +ASR_BATCH_SIZE=16 \ +WER_THRESHOLD_PCT=100 \ +SHARD_COUNT=1 \ +SHARD_INDEX=0 \ +bash run_parakeet_wer.sh +``` + +For an array, set `SHARD_COUNT` to the total task count and +`SHARD_INDEX=$SLURM_ARRAY_TASK_ID`. Each shard writes to a separate output +directory. + +## Multi-node cluster run + +The cluster launcher submits one one-GPU array task per shard. It also submits +a dependent CPU job that runs only after the complete array succeeds, merging +the dataset-wide WER report and all three Lhotse CutSets: + +```bash +DATA_ROOT=/shared/path/to/raw/sessions \ +MFA_WORK_DIR=/shared/path/to/completed/wav-mfa-workdir \ +OUTPUT_DIR=/shared/path/to/parakeet-wer-output \ +ASR_ENV=/shared/path/to/parakeet-conda-env \ +MODEL_CACHE_DIR=/shared/path/to/model-cache \ +SESSIONS_FILE=/shared/path/to/session-list.txt \ +NUM_NODES=8 \ +MAX_CONCURRENT_NODES=8 \ +SLURM_ACCOUNT= \ +SLURM_PARTITION= \ +bash parakeet_wer/cluster/run_multinode.sh +``` + +All paths must be visible at the same absolute location on every node. The +launcher preloads the model once before submitting the array to prevent +concurrent cache writes. Set `PRELOAD_MODEL=0` only when the shared cache is +already complete. Optional Pyxis settings are `CONTAINER_IMAGE` and +`CONTAINER_MOUNTS`. The scripts never copy dataset outputs between hosts. + +## Local two-GPU run + +The two-GPU launcher pre-caches the model once and starts one Curator/Xenna +pipeline with both GPUs visible. The ASR stage requests one GPU per worker, so +Xenna schedules concurrent Parakeet workers inside one Ray cluster without +cross-cluster conflicts: + +```bash +DATA_ROOT=/path/to/raw/sessions \ +MFA_WORK_DIR=/path/to/completed/wav-mfa-workdir \ +OUTPUT_DIR=/path/to/parakeet-wer-output \ +SESSIONS_FILE=/path/to/session-list.txt \ +GPU_IDS=0,1 \ +ASR_BATCH_SIZE=16 \ +WER_THRESHOLD_PCT=100 \ +bash run_local_2gpu.sh +``` + +Logs: + +```text +/logs/local_2gpu.log +``` + +Both GPU workers share a pre-populated read-only model cache. Curator manages +their worker and GPU isolation. The launcher sets `ASR_WORKERS=2`, and each ASR +worker requests exactly one GPU. + +## Threshold selection + +The default applied threshold keeps segments with: + +```text +WER <= 100% +``` + +`wer_distribution.json` reports: + +- min, max, and mean WER +- P25, P50, P75, P90, P95, and P99 +- histogram counts +- Tukey upper fence: `P75 + 1.5 × (P75 - P25)` +- proposed threshold: + +```text +min(100%, max(25%, P95, Tukey upper fence)) +``` + +Review distributions across all shards before changing the production +threshold. To apply each shard's proposal automatically, set: + +```bash +USE_RECOMMENDED_THRESHOLD=1 +``` + +For consistent filtering across a dataset, prefer a single threshold selected +from the merged distribution. + +After all shards finish, merge their WER distributions: + +```bash +python analyze_wer_distribution.py \ + --output-dir /path/to/parakeet-wer-output \ + --applied-threshold-pct 100 +``` + +This writes `wer_distribution_merged.json` without loading transcript or audio +files again. + +## Outputs + +```text +/shards/shard_/ +├── segments_with_wer.jsonl +├── wer_distribution.json +├── manifests/ +│ └── .jsonl +└── lhotse/ + ├── wer_000/cuts.jsonl.gz + ├── wer_000_010/cuts.jsonl.gz + └── wer_000_100/cuts.jsonl.gz +``` + +Each filtered manifest row contains: + +```json +{ + "audio_filepath": "/path/to/masked-speaker.wav", + "offset": 1.25, + "duration": 2.5, + "text": "normalized ground truth", + "text_raw": "Ground truth", + "pred_text": "parakeet hypothesis", + "wer_pct": 12.5, + "session_id": "", + "speaker_id": "", + "recording_id": "", + "segment_index": 0, + "alignment": [ + {"symbol": "word", "start": 0.1, "duration": 0.3} + ], + "alignment_source": "fastmss_textgrid" +} +``` + +Alignment starts are relative to the segment offset. The audit JSONL also +retains absolute FastMSS word times and rejection reasons. + +## Lhotse WER variants + +Every pipeline shard automatically builds three nested Lhotse CutSets from its +completed audit JSONL. Disable this only for diagnostics with +`run_pipeline.py --no-build-lhotse`. + +To rebuild a shard manually: + +```bash +python build_lhotse_variants.py \ + --audit-jsonl /shards/shard_00000/segments_with_wer.jsonl \ + --output-dir /shards/shard_00000/lhotse +``` + +Outputs: + +```text +/shards/shard_00000/lhotse/ +├── wer_000/cuts.jsonl.gz # exact 0% WER +├── wer_000_010/cuts.jsonl.gz # 0–10% WER +└── wer_000_100/cuts.jsonl.gz # 0–100% WER +``` + +All variants require non-empty FastMSS word alignment. Each `MonoCut` references +the full masked speaker WAV using its segment offset/duration, and its +supervision stores relative word alignments plus Parakeet text and WER metadata. + +The dependent cluster merge job writes dataset-level variants to: + +```text +/lhotse_merged//cuts.jsonl.gz +``` diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/analyze_wer_distribution.py b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/analyze_wer_distribution.py new file mode 100755 index 0000000000..deb9d52e72 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/analyze_wer_distribution.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Merge shard WER audit files and propose one dataset-wide threshold.""" + +from __future__ import annotations + +import argparse +import json +from array import array +from pathlib import Path + +import numpy as np +from manifest import build_wer_distribution_from_values + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--applied-threshold-pct", type=float, default=100.0) + args = parser.parse_args() + + audit_paths = sorted(args.output_dir.glob("shards/shard_*/segments_with_wer.jsonl")) + if not audit_paths: + parser.error(f"no shard audit files found under {args.output_dir}") + + values = array("d") + total_segments = 0 + for audit_path in audit_paths: + with audit_path.open(encoding="utf-8") as stream: + for line in stream: + row = json.loads(line) + total_segments += 1 + if row.get("wer_pct") is not None: + values.append(float(row["wer_pct"])) + + finite = np.frombuffer(values, dtype=np.float64) + report = build_wer_distribution_from_values( + finite, + total_segments=total_segments, + applied_threshold_pct=args.applied_threshold_pct, + ) + report["shards"] = len(audit_paths) + report_path = args.output_dir / "wer_distribution_merged.json" + report_path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + print(f"Report: {report_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/build_lhotse_variants.py b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/build_lhotse_variants.py new file mode 100755 index 0000000000..901d31008d --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/build_lhotse_variants.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Build exact-0%, 0-10%, and 0-100% WER Lhotse CutSet variants.""" + +from __future__ import annotations + +import argparse +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from lhotse import CutSet, MonoCut, Recording, SupervisionSegment +from lhotse.supervision import AlignmentItem + + +@dataclass(frozen=True) +class WERVariant: + name: str + upper_pct: float + exact_zero: bool = False + + +VARIANTS = ( + WERVariant(name="wer_000", upper_pct=0.0, exact_zero=True), + WERVariant(name="wer_000_010", upper_pct=10.0), + WERVariant(name="wer_000_100", upper_pct=100.0), +) + + +def segment_matches_variant(row: dict[str, Any], variant: WERVariant) -> bool: + wer = row.get("wer_pct") + if wer is None or not row.get("alignment"): + return False + value = float(wer) + return value == 0.0 if variant.exact_zero else 0.0 <= value <= variant.upper_pct + + +def _alignment_items(row: dict[str, Any], cut_duration: float) -> list[AlignmentItem]: + items: list[AlignmentItem] = [] + for item in row["alignment"]: + start = max(0.0, float(item["start"])) + end = min(cut_duration, start + float(item["duration"])) + if end <= start: + continue + items.append( + AlignmentItem( + symbol=str(item["symbol"]), + start=round(start, 6), + duration=round(end - start, 6), + ) + ) + return items + + +def _build_cut(row: dict[str, Any], recording: Recording) -> MonoCut | None: + offset = max(0.0, float(row["start"])) + duration = min(float(row["duration"]), recording.duration - offset) + if duration <= 0: + return None + cut_id = f"{row['recording_id']}_{int(row['segment_index']):05d}" + supervision = SupervisionSegment( + id=cut_id, + recording_id=recording.id, + start=0.0, + duration=duration, + channel=0, + text=str(row["text"]), + language="en", + speaker=str(row["speaker_id"]), + alignment={"word": _alignment_items(row, duration)}, + custom={ + "session_id": row["session_id"], + "segment_index": int(row["segment_index"]), + "text_raw": row["text_raw"], + "pred_text": row["pred_text"], + "wer_pct": float(row["wer_pct"]), + "alignment_source": "fastmss_textgrid", + }, + ) + return MonoCut( + id=cut_id, + start=offset, + duration=duration, + channel=0, + recording=recording, + supervisions=[supervision], + custom={"wer_pct": float(row["wer_pct"])}, + ) + + +def build_variant( + audit_path: Path, + output_dir: Path, + variant: WERVariant, + recording_cache: dict[str, Recording], +) -> dict[str, Any]: + cuts: list[MonoCut] = [] + matched = invalid_bounds = 0 + with audit_path.open(encoding="utf-8") as stream: + for line in stream: + row = json.loads(line) + if not segment_matches_variant(row, variant): + continue + matched += 1 + recording_id = str(row["recording_id"]) + if recording_id not in recording_cache: + recording_cache[recording_id] = Recording.from_file( + row["audio_filepath"], + recording_id=recording_id, + ) + cut = _build_cut(row, recording_cache[recording_id]) + if cut is None: + invalid_bounds += 1 + continue + cuts.append(cut) + + variant_dir = output_dir / variant.name + variant_dir.mkdir(parents=True, exist_ok=True) + cut_path = variant_dir / "cuts.jsonl.gz" + CutSet.from_cuts(cuts).to_file(cut_path) + summary = { + "variant": variant.name, + "exact_zero": variant.exact_zero, + "upper_wer_pct": variant.upper_pct, + "matched_segments": matched, + "cuts_written": len(cuts), + "invalid_audio_bounds": invalid_bounds, + "recordings_referenced": len({cut.recording_id for cut in cuts}), + "cut_manifest": str(cut_path), + } + (variant_dir / "summary.json").write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return summary + + +def build_all_variants(audit_path: Path, output_dir: Path) -> list[dict[str, Any]]: + """Build all configured Lhotse WER variants from one shard audit.""" + recording_cache: dict[str, Recording] = {} + summaries = [build_variant(audit_path, output_dir, variant, recording_cache) for variant in VARIANTS] + summary_path = output_dir / "summary.json" + summary_path.write_text(json.dumps(summaries, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return summaries + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--audit-jsonl", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + args = parser.parse_args() + + if not args.audit_jsonl.is_file(): + parser.error(f"audit JSONL does not exist: {args.audit_jsonl}") + summaries = build_all_variants(args.audit_jsonl, args.output_dir) + summary_path = args.output_dir / "summary.json" + print(json.dumps(summaries, indent=2, sort_keys=True)) + print(f"Summary: {summary_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/cluster/merge_outputs.sh b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/cluster/merge_outputs.sh new file mode 100755 index 0000000000..ce055be15a --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/cluster/merge_outputs.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Merge successful cluster shards into dataset-level WER and Lhotse outputs. + +set -euo pipefail + +: "${OUTPUT_DIR:?Missing OUTPUT_DIR}" +: "${PARAKEET_ROOT:?Missing PARAKEET_ROOT}" +: "${ASR_ENV:?Missing ASR_ENV}" + +WER_THRESHOLD_PCT="${WER_THRESHOLD_PCT:-100}" +PYTHON="$ASR_ENV/bin/python" +if [[ ! -x "$PYTHON" ]]; then + echo "ERROR: merge node cannot access ASR_ENV=$ASR_ENV" >&2 + exit 1 +fi + +"$PYTHON" "$PARAKEET_ROOT/analyze_wer_distribution.py" \ + --output-dir "$OUTPUT_DIR" \ + --applied-threshold-pct "$WER_THRESHOLD_PCT" + +"$PYTHON" "$PARAKEET_ROOT/merge_lhotse_variants.py" \ + --output-dir "$OUTPUT_DIR" + +echo "[$(date -Is)] Dataset-level WER and Lhotse outputs completed" diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/cluster/run_multinode.sh b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/cluster/run_multinode.sh new file mode 100755 index 0000000000..2f6f9dede4 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/cluster/run_multinode.sh @@ -0,0 +1,126 @@ +#!/bin/bash +# Submit Parakeet WER shards and a dependent Lhotse/distribution merge job. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TUTORIAL_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +PARAKEET_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" + +DATA_ROOT="${DATA_ROOT:?Set DATA_ROOT to raw session manifests}" +MFA_WORK_DIR="${MFA_WORK_DIR:?Set MFA_WORK_DIR to completed masked-WAV MFA outputs}" +OUTPUT_DIR="${OUTPUT_DIR:?Set OUTPUT_DIR for Parakeet/WER outputs}" +ASR_ENV="${ASR_ENV:?Set ASR_ENV to a shared Parakeet environment}" +MODEL_CACHE_DIR="${MODEL_CACHE_DIR:?Set MODEL_CACHE_DIR to a shared model cache}" +SESSIONS_FILE="${SESSIONS_FILE:-}" +PARAKEET_MODEL="${PARAKEET_MODEL:-nvidia/parakeet-tdt-0.6b-v2}" +ASR_BATCH_SIZE="${ASR_BATCH_SIZE:-16}" +WER_THRESHOLD_PCT="${WER_THRESHOLD_PCT:-100}" + +NUM_NODES="${NUM_NODES:-1}" +MAX_CONCURRENT_NODES="${MAX_CONCURRENT_NODES:-$NUM_NODES}" +CPUS_PER_NODE="${CPUS_PER_NODE:-16}" +GPU_GRES="${GPU_GRES:-gpu:1}" +TIME_LIMIT="${TIME_LIMIT:-08:00:00}" +MEMORY_PER_NODE="${MEMORY_PER_NODE:-64G}" +SLURM_ACCOUNT="${SLURM_ACCOUNT:-}" +SLURM_PARTITION="${SLURM_PARTITION:-}" +JOB_NAME="${JOB_NAME:-david-ai-parakeet-wer}" +PRELOAD_MODEL="${PRELOAD_MODEL:-1}" +CONTAINER_IMAGE="${CONTAINER_IMAGE:-}" +CONTAINER_MOUNTS="${CONTAINER_MOUNTS:-}" + +for path in "$DATA_ROOT" "$MFA_WORK_DIR" "$OUTPUT_DIR" "$ASR_ENV" "$MODEL_CACHE_DIR" "$TUTORIAL_ROOT"; do + if [[ "$path" != /* || "$path" == *,* ]]; then + echo "ERROR: shared paths must be absolute and contain no commas: $path" >&2 + exit 2 + fi +done +if ! [[ "$NUM_NODES" =~ ^[1-9][0-9]*$ && "$MAX_CONCURRENT_NODES" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: NUM_NODES and MAX_CONCURRENT_NODES must be positive integers" >&2 + exit 2 +fi +if [[ -n "$SESSIONS_FILE" && ! -f "$SESSIONS_FILE" ]]; then + echo "ERROR: sessions file does not exist: $SESSIONS_FILE" >&2 + exit 1 +fi +if [[ -n "$SESSIONS_FILE" && ( "$SESSIONS_FILE" != /* || "$SESSIONS_FILE" == *,* ) ]]; then + echo "ERROR: SESSIONS_FILE must be absolute and contain no commas" >&2 + exit 2 +fi +if [[ ! -x "$ASR_ENV/bin/python" ]]; then + echo "ERROR: ASR environment is missing python: $ASR_ENV" >&2 + exit 1 +fi +if ! command -v sbatch >/dev/null 2>&1; then + echo "ERROR: sbatch is not available" >&2 + exit 1 +fi + +mkdir -p "$OUTPUT_DIR/logs/slurm" "$MODEL_CACHE_DIR" +export PYTHONPATH="$(cd "$TUTORIAL_ROOT/../../.." && pwd)" +export NEMO_CACHE_DIR="$MODEL_CACHE_DIR" + +if [[ "$PRELOAD_MODEL" == "1" ]]; then + "$ASR_ENV/bin/python" - "$PARAKEET_MODEL" <<'PY' +import sys +from nemo.collections.asr.models import ASRModel + +ASRModel.from_pretrained(model_name=sys.argv[1], return_model_file=True) +PY +fi + +MOUNTS_B64="" +[[ -n "$CONTAINER_MOUNTS" ]] && + MOUNTS_B64="$(printf "%s" "$CONTAINER_MOUNTS" | base64 | tr -d "\n")" + +EXPORTS="DATA_ROOT=$DATA_ROOT,MFA_WORK_DIR=$MFA_WORK_DIR,OUTPUT_DIR=$OUTPUT_DIR" +EXPORTS+=",ASR_ENV=$ASR_ENV,MODEL_CACHE_DIR=$MODEL_CACHE_DIR" +EXPORTS+=",TUTORIAL_ROOT=$TUTORIAL_ROOT,PARAKEET_ROOT=$PARAKEET_ROOT" +EXPORTS+=",PARAKEET_MODEL=$PARAKEET_MODEL,ASR_BATCH_SIZE=$ASR_BATCH_SIZE" +EXPORTS+=",WER_THRESHOLD_PCT=$WER_THRESHOLD_PCT,SHARD_COUNT=$NUM_NODES" +EXPORTS+=",CONTAINER_IMAGE=$CONTAINER_IMAGE,CONTAINER_MOUNTS_B64=$MOUNTS_B64" +[[ -n "$SESSIONS_FILE" ]] && EXPORTS+=",SESSIONS_FILE=$SESSIONS_FILE" + +ARRAY_ARGS=( + --parsable + --job-name "$JOB_NAME" + --nodes 1 + --ntasks 1 + --cpus-per-task "$CPUS_PER_NODE" + --gres "$GPU_GRES" + --mem "$MEMORY_PER_NODE" + --time "$TIME_LIMIT" + --array "0-$((NUM_NODES - 1))%$MAX_CONCURRENT_NODES" + --output "$OUTPUT_DIR/logs/slurm/${JOB_NAME}_%A_%a.out" + --error "$OUTPUT_DIR/logs/slurm/${JOB_NAME}_%A_%a.err" + --export "$EXPORTS" +) +[[ -n "$SLURM_ACCOUNT" ]] && ARRAY_ARGS+=(--account "$SLURM_ACCOUNT") +[[ -n "$SLURM_PARTITION" ]] && ARRAY_ARGS+=(--partition "$SLURM_PARTITION") + +ARRAY_JOB_ID="$(sbatch "${ARRAY_ARGS[@]}" "$SCRIPT_DIR/run_node.sh")" +ARRAY_JOB_ID="${ARRAY_JOB_ID%%;*}" +echo "Submitted Parakeet array job: $ARRAY_JOB_ID" + +MERGE_EXPORTS="OUTPUT_DIR=$OUTPUT_DIR,PARAKEET_ROOT=$PARAKEET_ROOT,ASR_ENV=$ASR_ENV" +MERGE_EXPORTS+=",WER_THRESHOLD_PCT=$WER_THRESHOLD_PCT" +MERGE_ARGS=( + --parsable + --job-name "${JOB_NAME}-merge" + --nodes 1 + --ntasks 1 + --cpus-per-task 4 + --mem 32G + --time 02:00:00 + --dependency "afterok:$ARRAY_JOB_ID" + --output "$OUTPUT_DIR/logs/slurm/${JOB_NAME}_merge_%j.out" + --error "$OUTPUT_DIR/logs/slurm/${JOB_NAME}_merge_%j.err" + --export "$MERGE_EXPORTS" +) +[[ -n "$SLURM_ACCOUNT" ]] && MERGE_ARGS+=(--account "$SLURM_ACCOUNT") +[[ -n "$SLURM_PARTITION" ]] && MERGE_ARGS+=(--partition "$SLURM_PARTITION") + +MERGE_JOB_ID="$(sbatch "${MERGE_ARGS[@]}" "$SCRIPT_DIR/merge_outputs.sh")" +MERGE_JOB_ID="${MERGE_JOB_ID%%;*}" +echo "Submitted dependent merge job: $MERGE_JOB_ID" diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/cluster/run_node.sh b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/cluster/run_node.sh new file mode 100755 index 0000000000..d12faff3e1 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/cluster/run_node.sh @@ -0,0 +1,81 @@ +#!/bin/bash +# Execute one Parakeet WER shard on one GPU node. + +set -euo pipefail + +: "${DATA_ROOT:?Missing DATA_ROOT}" +: "${MFA_WORK_DIR:?Missing MFA_WORK_DIR}" +: "${OUTPUT_DIR:?Missing OUTPUT_DIR}" +: "${ASR_ENV:?Missing ASR_ENV}" +: "${MODEL_CACHE_DIR:?Missing MODEL_CACHE_DIR}" +: "${TUTORIAL_ROOT:?Missing TUTORIAL_ROOT}" +: "${PARAKEET_ROOT:?Missing PARAKEET_ROOT}" +: "${SHARD_COUNT:?Missing SHARD_COUNT}" +: "${SLURM_ARRAY_TASK_ID:?This must run as a SLURM array task}" + +SESSIONS_FILE="${SESSIONS_FILE:-}" +PARAKEET_MODEL="${PARAKEET_MODEL:-nvidia/parakeet-tdt-0.6b-v2}" +ASR_BATCH_SIZE="${ASR_BATCH_SIZE:-16}" +WER_THRESHOLD_PCT="${WER_THRESHOLD_PCT:-100}" +CONTAINER_IMAGE="${CONTAINER_IMAGE:-}" +CONTAINER_MOUNTS_B64="${CONTAINER_MOUNTS_B64:-}" +CONTAINER_MOUNTS="" +[[ -n "$CONTAINER_MOUNTS_B64" ]] && + CONTAINER_MOUNTS="$(printf "%s" "$CONTAINER_MOUNTS_B64" | base64 --decode)" + +if [[ -n "$CONTAINER_IMAGE" && "${IN_CONTAINER:-0}" != "1" ]]; then + SRUN_ARGS=(--nodes 1 --ntasks 1 --container-image "$CONTAINER_IMAGE") + [[ -n "$CONTAINER_MOUNTS" ]] && + SRUN_ARGS+=(--container-mounts "$CONTAINER_MOUNTS") + exec srun "${SRUN_ARGS[@]}" \ + env IN_CONTAINER=1 bash "$PARAKEET_ROOT/cluster/run_node.sh" +fi + +if [[ ! -x "$ASR_ENV/bin/python" ]]; then + echo "ERROR: compute node cannot access ASR_ENV=$ASR_ENV" >&2 + exit 1 +fi +for path in "$DATA_ROOT" "$MFA_WORK_DIR/audio_16k_masked" "$MFA_WORK_DIR/textgrids"; do + if [[ ! -d "$path" ]]; then + echo "ERROR: missing required directory: $path" >&2 + exit 1 + fi +done + +export PATH="$ASR_ENV/bin:/usr/local/bin:/usr/bin:/bin" +export PYTHON="$ASR_ENV/bin/python" +export PYTHONPATH="$(cd "$TUTORIAL_ROOT/../../.." && pwd)" +export NEMO_CACHE_DIR="$MODEL_CACHE_DIR" +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export NUMEXPR_MAX_THREADS=1 + +SCRATCH_ROOT="${SLURM_TMPDIR:-/tmp}" +SCRATCH_DIR="$SCRATCH_ROOT/david_ai_parakeet_${SLURM_JOB_ID}_${SLURM_ARRAY_TASK_ID}" +rm -rf "$SCRATCH_DIR" +mkdir -p "$SCRATCH_DIR" +cleanup() { + rm -rf "$SCRATCH_DIR" +} +trap cleanup EXIT + +echo "[$(date -Is)] Node=$(hostname) shard=$SLURM_ARRAY_TASK_ID/$SHARD_COUNT" +env \ + DATA_ROOT="$DATA_ROOT" \ + MFA_WORK_DIR="$MFA_WORK_DIR" \ + OUTPUT_DIR="$OUTPUT_DIR" \ + SESSIONS_FILE="$SESSIONS_FILE" \ + PARAKEET_MODEL="$PARAKEET_MODEL" \ + MODEL_CACHE_DIR="$MODEL_CACHE_DIR" \ + ASR_BATCH_SIZE="$ASR_BATCH_SIZE" \ + ASR_WORKERS=1 \ + WER_THRESHOLD_PCT="$WER_THRESHOLD_PCT" \ + BUILD_LHOTSE=1 \ + SHARD_COUNT="$SHARD_COUNT" \ + SHARD_INDEX="$SLURM_ARRAY_TASK_ID" \ + SCRATCH_DIR="$SCRATCH_DIR" \ + PYTHON="$PYTHON" \ + bash "$PARAKEET_ROOT/run_parakeet_wer.sh" + +echo "[$(date -Is)] Parakeet shard $SLURM_ARRAY_TASK_ID completed" diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/manifest.py b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/manifest.py new file mode 100644 index 0000000000..33ce4071d8 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/manifest.py @@ -0,0 +1,318 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Build segment tasks and write WER-filtered per-speaker manifests.""" + +from __future__ import annotations + +import json +from collections import defaultdict +from dataclasses import dataclass +from itertools import pairwise +from typing import TYPE_CHECKING, Any + +import numpy as np +from stages import normalize_wer_text +from textgrid import TextGrid + +from nemo_curator.tasks import AudioTask + +if TYPE_CHECKING: + from pathlib import Path + +SILENCE_TOKENS = {"", "sil", "sp", "spn", ""} +HISTOGRAM_BOUNDS = (0.0, 10.0, 25.0, 50.0, 75.0, 100.0, 150.0, 200.0, float("inf")) + + +def recording_id(speaker_id: str, session_id: str) -> str: + return f"{speaker_id}_{session_id}_postprocessed" + + +def load_fastmss_words(textgrid_path: Path) -> list[tuple[float, float, str]]: + """Read recording-global word intervals from the FastMSS words tier.""" + textgrid = TextGrid.fromFile(str(textgrid_path)) + tier = textgrid.getFirst("words") + return [ + (float(interval.minTime), float(interval.maxTime), interval.mark.strip()) + for interval in tier.intervals + if interval.mark.strip() not in SILENCE_TOKENS + ] + + +def segment_alignments( + words: list[tuple[float, float, str]], + start: float, + end: float, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Clip recording-global FastMSS words to one segment.""" + absolute: list[dict[str, Any]] = [] + relative: list[dict[str, Any]] = [] + for word_start, word_end, word in words: + clipped_start = max(start, word_start) + clipped_end = min(end, word_end) + if clipped_end <= clipped_start: + continue + absolute.append( + { + "word": word, + "start": round(clipped_start, 6), + "end": round(clipped_end, 6), + } + ) + relative.append( + { + "symbol": word, + "start": round(clipped_start - start, 6), + "duration": round(clipped_end - clipped_start, 6), + } + ) + return absolute, relative + + +def load_session_ids(data_root: Path, sessions_file: Path | None) -> list[str]: + if sessions_file is None: + return sorted(path.name for path in data_root.iterdir() if path.is_dir()) + return sorted( + { + line.strip() + for line in sessions_file.read_text(encoding="utf-8").splitlines() + if line.strip() and not line.lstrip().startswith("#") + } + ) + + +@dataclass(frozen=True) +class SegmentTaskConfig: + data_root: Path + masked_audio_dir: Path + textgrid_dir: Path + sessions_file: Path | None + shard_count: int + shard_index: int + + +def _build_session_tasks(config: SegmentTaskConfig, session_id: str) -> list[AudioTask]: + transcript_path = config.data_root / session_id / "machine_generated_transcript.json" + if not transcript_path.is_file(): + msg = f"missing transcript: {transcript_path}" + raise FileNotFoundError(msg) + payload = json.loads(transcript_path.read_text(encoding="utf-8")) + segments = payload.get("transcript") if isinstance(payload, dict) else None + if not isinstance(segments, list): + msg = f"invalid transcript list: {transcript_path}" + raise TypeError(msg) + + by_speaker: dict[str, list[tuple[int, dict[str, Any]]]] = defaultdict(list) + for segment_index, segment in enumerate(segments): + if isinstance(segment, dict) and segment.get("speaker"): + by_speaker[str(segment["speaker"])].append((segment_index, segment)) + + tasks: list[AudioTask] = [] + for speaker_id, speaker_segments in sorted(by_speaker.items()): + rec_id = recording_id(speaker_id, session_id) + audio_path = config.masked_audio_dir / f"{rec_id}.wav" + textgrid_path = config.textgrid_dir / f"{rec_id}_fastmss.TextGrid" + if not audio_path.is_file(): + msg = f"missing masked speaker WAV: {audio_path}" + raise FileNotFoundError(msg) + if not textgrid_path.is_file(): + msg = f"missing FastMSS TextGrid: {textgrid_path}" + raise FileNotFoundError(msg) + words = load_fastmss_words(textgrid_path) + + for segment_index, segment in speaker_segments: + start = float(segment["start"]) + end = float(segment["end"]) + if end <= start: + continue + text_raw = str(segment.get("text") or "").strip() + absolute_words, alignment = segment_alignments(words, start, end) + tasks.append( + AudioTask( + dataset_name="david_ai_masked", + filepath_key="audio_filepath", + data={ + "session_id": session_id, + "speaker_id": speaker_id, + "recording_id": rec_id, + "segment_index": segment_index, + "start": start, + "end": end, + "duration": round(end - start, 6), + "audio_filepath": str(audio_path), + "text_raw": text_raw, + "text": normalize_wer_text(text_raw), + "fastmss_textgrid": str(textgrid_path), + "words": absolute_words, + "alignment": alignment, + }, + ) + ) + return tasks + + +def build_segment_tasks(config: SegmentTaskConfig) -> list[AudioTask]: + """Build one AudioTask per ground-truth segment for a deterministic shard.""" + session_ids = load_session_ids(config.data_root, config.sessions_file) + selected_ids = [ + session_id + for index, session_id in enumerate(session_ids) + if index % config.shard_count == config.shard_index + ] + tasks: list[AudioTask] = [] + for session_id in selected_ids: + tasks.extend(_build_session_tasks(config, session_id)) + return tasks + + +def _percentiles(values: np.ndarray) -> dict[str, float]: + return { + key: round(float(np.percentile(values, percentile)), 6) + for key, percentile in (("p25", 25), ("p50", 50), ("p75", 75), ("p90", 90), ("p95", 95), ("p99", 99)) + } + + +def build_wer_distribution_from_values( + finite: np.ndarray, + *, + total_segments: int, + applied_threshold_pct: float, +) -> dict[str, Any]: + """Compute histogram, percentiles, and a robust proposed WER threshold.""" + if finite.size == 0: + return { + "segments": total_segments, + "segments_with_wer": 0, + "applied_threshold_pct": applied_threshold_pct, + "recommended_threshold_pct": applied_threshold_pct, + "histogram": [], + "percentiles": {}, + } + + percentiles = _percentiles(finite) + iqr = percentiles["p75"] - percentiles["p25"] + tukey_upper = percentiles["p75"] + 1.5 * iqr + recommended = round(min(100.0, max(25.0, percentiles["p95"], tukey_upper)), 6) + histogram: list[dict[str, Any]] = [] + for lower, upper in pairwise(HISTOGRAM_BOUNDS): + count = int(np.sum((finite >= lower) & (finite < upper))) + histogram.append( + { + "lower_pct": lower, + "upper_pct": None if np.isinf(upper) else upper, + "count": count, + } + ) + return { + "segments": total_segments, + "segments_with_wer": int(finite.size), + "segments_without_reference_wer": total_segments - int(finite.size), + "min_pct": round(float(finite.min()), 6), + "max_pct": round(float(finite.max()), 6), + "mean_pct": round(float(finite.mean()), 6), + "percentiles": percentiles, + "tukey_upper_fence_pct": round(tukey_upper, 6), + "recommended_threshold_pct": recommended, + "applied_threshold_pct": applied_threshold_pct, + "histogram": histogram, + } + + +def build_wer_distribution(rows: list[dict[str, Any]], applied_threshold_pct: float) -> dict[str, Any]: + """Compute a WER report from segment dictionaries.""" + finite = np.asarray( + [float(row["wer_pct"]) for row in rows if row.get("wer_pct") is not None], + dtype=np.float64, + ) + return build_wer_distribution_from_values( + finite, + total_segments=len(rows), + applied_threshold_pct=applied_threshold_pct, + ) + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("w", encoding="utf-8") as stream: + for row in rows: + stream.write(json.dumps(row, ensure_ascii=False) + "\n") + temporary.replace(path) + + +def write_pipeline_outputs( + tasks: list[AudioTask], + *, + output_dir: Path, + threshold_pct: float, + require_fastmss_alignment: bool, +) -> dict[str, Any]: + """Write audit rows, filtered per-speaker manifests, and WER analytics.""" + rows: list[dict[str, Any]] = [] + kept_by_recording: dict[str, list[dict[str, Any]]] = defaultdict(list) + for task in tasks: + row = dict(task.data) + row.pop("segment_audio_filepath", None) + row.pop("source_audio_filepath", None) + reasons: list[str] = [] + if row.get("wer_pct") is None: + reasons.append("empty_reference") + elif float(row["wer_pct"]) > threshold_pct: + reasons.append("wer_above_threshold") + if require_fastmss_alignment and not row.get("alignment"): + reasons.append("missing_fastmss_alignment") + row["keep"] = not reasons + row["rejection_reasons"] = reasons + rows.append(row) + if row["keep"]: + kept_by_recording[str(row["recording_id"])].append(row) + + output_dir.mkdir(parents=True, exist_ok=True) + _write_jsonl(output_dir / "segments_with_wer.jsonl", rows) + manifests_dir = output_dir / "manifests" + for rec_id, recording_rows in kept_by_recording.items(): + manifest_rows = [ + { + "audio_filepath": row["audio_filepath"], + "offset": row["start"], + "duration": row["duration"], + "text": row["text"], + "text_raw": row["text_raw"], + "pred_text": row["pred_text"], + "wer_pct": row["wer_pct"], + "session_id": row["session_id"], + "speaker_id": row["speaker_id"], + "recording_id": row["recording_id"], + "segment_index": row["segment_index"], + "alignment": row["alignment"], + "alignment_source": "fastmss_textgrid", + } + for row in recording_rows + ] + _write_jsonl(manifests_dir / f"{rec_id}.jsonl", manifest_rows) + + report = build_wer_distribution(rows, threshold_pct) + report.update( + { + "kept_segments": sum(bool(row["keep"]) for row in rows), + "rejected_segments": sum(not row["keep"] for row in rows), + "recording_manifests": len(kept_by_recording), + "require_fastmss_alignment": require_fastmss_alignment, + } + ) + (output_dir / "wer_distribution.json").write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return report diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/merge_lhotse_variants.py b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/merge_lhotse_variants.py new file mode 100755 index 0000000000..e0406679fc --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/merge_lhotse_variants.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Merge per-shard Lhotse WER variants into dataset-level CutSets.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from build_lhotse_variants import VARIANTS +from lhotse import CutSet + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + args = parser.parse_args() + + merged_root = args.output_dir / "lhotse_merged" + merged_root.mkdir(parents=True, exist_ok=True) + summaries = [] + for variant in VARIANTS: + paths = sorted(args.output_dir.glob(f"shards/shard_*/lhotse/{variant.name}/cuts.jsonl.gz")) + if not paths: + parser.error(f"no shard CutSets found for {variant.name}") + variant_dir = merged_root / variant.name + variant_dir.mkdir(parents=True, exist_ok=True) + output_path = variant_dir / "cuts.jsonl.gz" + CutSet.from_files(paths, shuffle_iters=False).to_file(output_path) + cut_count = sum(1 for _ in CutSet.from_file(output_path)) + summaries.append( + { + "variant": variant.name, + "shards": len(paths), + "cuts": cut_count, + "cut_manifest": str(output_path), + } + ) + + summary_path = merged_root / "summary.json" + summary_path.write_text(json.dumps(summaries, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(summaries, indent=2, sort_keys=True)) + print(f"Summary: {summary_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/requirements.txt b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/requirements.txt new file mode 100644 index 0000000000..f102f632d5 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/requirements.txt @@ -0,0 +1,5 @@ +nemo_toolkit[asr]>=2.7.2 +nvidia-ml-py +num2words>=0.5.12 +numpy +textgrid diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/run_local_2gpu.sh b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/run_local_2gpu.sh new file mode 100755 index 0000000000..5e8c436110 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/run_local_2gpu.sh @@ -0,0 +1,87 @@ +#!/bin/bash +# Run one local Curator/Xenna pipeline with two visible GPUs. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)" + +DATA_ROOT="${DATA_ROOT:?Set DATA_ROOT to raw session manifests}" +MFA_WORK_DIR="${MFA_WORK_DIR:?Set MFA_WORK_DIR to completed masked-WAV MFA outputs}" +OUTPUT_DIR="${OUTPUT_DIR:?Set OUTPUT_DIR for Parakeet/WER outputs}" +SESSIONS_FILE="${SESSIONS_FILE:-}" +PARAKEET_MODEL="${PARAKEET_MODEL:-nvidia/parakeet-tdt-0.6b-v2}" +MODEL_CACHE_DIR="${MODEL_CACHE_DIR:-$OUTPUT_DIR/model_cache}" +ASR_BATCH_SIZE="${ASR_BATCH_SIZE:-16}" +WER_THRESHOLD_PCT="${WER_THRESHOLD_PCT:-100}" +GPU_IDS="${GPU_IDS:-0,1}" +PYTHON="${PYTHON:-python3}" +SCRATCH_DIR="${SCRATCH_DIR:-/tmp/david_ai_parakeet_2gpu_$$}" +LOG_DIR="${LOG_DIR:-$OUTPUT_DIR/logs}" + +IFS="," read -r GPU_0 GPU_1 EXTRA_GPUS <<< "$GPU_IDS" +if [[ -z "$GPU_0" || -z "$GPU_1" || -n "${EXTRA_GPUS:-}" || "$GPU_0" == "$GPU_1" ]]; then + echo "ERROR: GPU_IDS must contain exactly two distinct IDs, for example 0,1" >&2 + exit 2 +fi +if [[ ! "$GPU_0" =~ ^[0-9]+$ || ! "$GPU_1" =~ ^[0-9]+$ ]]; then + echo "ERROR: GPU IDs must be non-negative integers" >&2 + exit 2 +fi +for path in "$DATA_ROOT" "$MFA_WORK_DIR/audio_16k_masked" "$MFA_WORK_DIR/textgrids"; do + if [[ ! -d "$path" ]]; then + echo "ERROR: required directory does not exist: $path" >&2 + exit 1 + fi +done +if [[ -n "$SESSIONS_FILE" && ! -f "$SESSIONS_FILE" ]]; then + echo "ERROR: sessions file does not exist: $SESSIONS_FILE" >&2 + exit 1 +fi + +mkdir -p "$OUTPUT_DIR" "$MODEL_CACHE_DIR" "$SCRATCH_DIR" "$LOG_DIR" +export PYTHONPATH="$REPO_ROOT" +export NEMO_CACHE_DIR="$MODEL_CACHE_DIR" + +echo "Pre-caching $PARAKEET_MODEL once in $MODEL_CACHE_DIR" +CUDA_VISIBLE_DEVICES="" "$PYTHON" - "$PARAKEET_MODEL" <<'PY' +import sys +from nemo.collections.asr.models import ASRModel + +ASRModel.from_pretrained( + model_name=sys.argv[1], + return_model_file=True, +) +PY + +cleanup() { + rm -rf "$SCRATCH_DIR" +} +trap cleanup EXIT INT TERM + +echo "Starting one Xenna pipeline with visible GPUs $GPU_0,$GPU_1" +env \ + CUDA_VISIBLE_DEVICES="$GPU_0,$GPU_1" \ + DATA_ROOT="$DATA_ROOT" \ + MFA_WORK_DIR="$MFA_WORK_DIR" \ + OUTPUT_DIR="$OUTPUT_DIR" \ + SESSIONS_FILE="$SESSIONS_FILE" \ + PARAKEET_MODEL="$PARAKEET_MODEL" \ + MODEL_CACHE_DIR="$MODEL_CACHE_DIR" \ + ASR_BATCH_SIZE="$ASR_BATCH_SIZE" \ + ASR_WORKERS=2 \ + WER_THRESHOLD_PCT="$WER_THRESHOLD_PCT" \ + SHARD_COUNT=1 \ + SHARD_INDEX=0 \ + SCRATCH_DIR="$SCRATCH_DIR" \ + PYTHON="$PYTHON" \ + bash "$SCRIPT_DIR/run_parakeet_wer.sh" \ + > "$LOG_DIR/local_2gpu.log" 2>&1 + +"$PYTHON" "$SCRIPT_DIR/analyze_wer_distribution.py" \ + --output-dir "$OUTPUT_DIR" \ + --applied-threshold-pct "$WER_THRESHOLD_PCT" + +trap - EXIT INT TERM +cleanup +echo "Two-GPU Parakeet WER pipeline completed successfully" diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/run_parakeet_wer.sh b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/run_parakeet_wer.sh new file mode 100755 index 0000000000..d68a575fd4 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/run_parakeet_wer.sh @@ -0,0 +1,50 @@ +#!/bin/bash +# Run the Parakeet segment-WER pipeline for one deterministic shard. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)" + +DATA_ROOT="${DATA_ROOT:?Set DATA_ROOT to raw session manifests}" +MFA_WORK_DIR="${MFA_WORK_DIR:?Set MFA_WORK_DIR to the completed masked-WAV MFA output}" +OUTPUT_DIR="${OUTPUT_DIR:?Set OUTPUT_DIR for Parakeet/WER manifests}" +SESSIONS_FILE="${SESSIONS_FILE:-}" +PARAKEET_MODEL="${PARAKEET_MODEL:-nvidia/parakeet-tdt-0.6b-v2}" +MODEL_CACHE_DIR="${MODEL_CACHE_DIR:-}" +ASR_BATCH_SIZE="${ASR_BATCH_SIZE:-16}" +ASR_WORKERS="${ASR_WORKERS:-1}" +WER_THRESHOLD_PCT="${WER_THRESHOLD_PCT:-100}" +USE_RECOMMENDED_THRESHOLD="${USE_RECOMMENDED_THRESHOLD:-0}" +BUILD_LHOTSE="${BUILD_LHOTSE:-1}" +SHARD_COUNT="${SHARD_COUNT:-1}" +SHARD_INDEX="${SHARD_INDEX:-0}" +SCRATCH_DIR="${SCRATCH_DIR:-${SLURM_TMPDIR:-/tmp}/david_ai_parakeet_${SLURM_JOB_ID:-local}_${SHARD_INDEX}}" +PYTHON="${PYTHON:-python3}" + +export PYTHONPATH="$REPO_ROOT" +[[ -n "$MODEL_CACHE_DIR" ]] && export NEMO_CACHE_DIR="$MODEL_CACHE_DIR" + +COMMAND=( + "$PYTHON" "$SCRIPT_DIR/run_pipeline.py" + --data-root "$DATA_ROOT" + --masked-audio-dir "$MFA_WORK_DIR/audio_16k_masked" + --textgrid-dir "$MFA_WORK_DIR/textgrids" + --output-dir "$OUTPUT_DIR" + --model-name "$PARAKEET_MODEL" + --asr-batch-size "$ASR_BATCH_SIZE" + --asr-workers "$ASR_WORKERS" + --wer-threshold-pct "$WER_THRESHOLD_PCT" + --scratch-dir "$SCRATCH_DIR" + --shard-count "$SHARD_COUNT" + --shard-index "$SHARD_INDEX" +) +[[ -n "$SESSIONS_FILE" ]] && COMMAND+=(--sessions-file "$SESSIONS_FILE") +[[ -n "$MODEL_CACHE_DIR" ]] && COMMAND+=(--model-cache-dir "$MODEL_CACHE_DIR") +[[ "$USE_RECOMMENDED_THRESHOLD" == "1" ]] && COMMAND+=(--use-recommended-threshold) +[[ "$BUILD_LHOTSE" == "0" ]] && COMMAND+=(--no-build-lhotse) + +echo "Parakeet model: $PARAKEET_MODEL" +echo "Shard: $SHARD_INDEX/$SHARD_COUNT" +echo "Applied WER threshold: $WER_THRESHOLD_PCT%" +exec "${COMMAND[@]}" diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/run_pipeline.py b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/run_pipeline.py new file mode 100755 index 0000000000..b2613dcdca --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/run_pipeline.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run Parakeet segment ASR, WER analysis, and FastMSS manifest filtering.""" + +from __future__ import annotations + +import argparse +import os +import shutil +import tempfile +from pathlib import Path + +from build_lhotse_variants import build_all_variants +from loguru import logger +from manifest import ( + SegmentTaskConfig, + build_segment_tasks, + build_wer_distribution, + write_pipeline_outputs, +) +from stages import ParallelInferenceAsrNemoStage, SegmentClipExtractionStage, SegmentWERStage + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-root", type=Path, required=True) + parser.add_argument("--masked-audio-dir", type=Path, required=True) + parser.add_argument("--textgrid-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--sessions-file", type=Path, default=None) + parser.add_argument("--model-name", default="nvidia/parakeet-tdt-0.6b-v2") + parser.add_argument("--model-cache-dir", type=str, default=None) + parser.add_argument("--asr-batch-size", type=int, default=16) + parser.add_argument("--asr-workers", type=int, default=1) + parser.add_argument("--wer-threshold-pct", type=float, default=100.0) + parser.add_argument("--use-recommended-threshold", action="store_true") + parser.add_argument("--build-lhotse", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument( + "--require-fastmss-alignment", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.add_argument("--scratch-dir", type=Path, default=None) + parser.add_argument("--shard-count", type=int, default=1) + parser.add_argument("--shard-index", type=int, default=0) + args = parser.parse_args() + + if args.shard_count < 1 or not 0 <= args.shard_index < args.shard_count: + parser.error("shard-index must be in [0, shard-count)") + if args.asr_workers < 1: + parser.error("asr-workers must be at least 1") + for path in (args.data_root, args.masked_audio_dir, args.textgrid_dir): + if not path.is_dir(): + parser.error(f"required directory does not exist: {path}") + if args.sessions_file is not None and not args.sessions_file.is_file(): + parser.error(f"sessions file does not exist: {args.sessions_file}") + return args + + +def main() -> int: + args = parse_args() + output_dir = args.output_dir / "shards" / f"shard_{args.shard_index:05d}" + if args.model_cache_dir: + os.environ["NEMO_CACHE_DIR"] = str(Path(args.model_cache_dir).resolve()) + scratch_dir = args.scratch_dir or (Path(tempfile.gettempdir()) / f"david_ai_parakeet_{args.shard_index:05d}") + shutil.rmtree(scratch_dir, ignore_errors=True) + scratch_dir.mkdir(parents=True, exist_ok=True) + + initial_tasks = build_segment_tasks( + SegmentTaskConfig( + data_root=args.data_root.resolve(), + masked_audio_dir=args.masked_audio_dir.resolve(), + textgrid_dir=args.textgrid_dir.resolve(), + sessions_file=args.sessions_file.resolve() if args.sessions_file else None, + shard_count=args.shard_count, + shard_index=args.shard_index, + ) + ) + if not initial_tasks: + logger.info("Shard {} contains no segments", args.shard_index) + return 0 + logger.info("Built {} segment tasks for shard {}/{}", len(initial_tasks), args.shard_index, args.shard_count) + + from nemo_curator.backends.xenna import XennaExecutor + from nemo_curator.pipeline import Pipeline + from nemo_curator.stages.resources import Resources + + pipeline = Pipeline( + name="david-ai-parakeet-segment-wer", + description="Segment Parakeet ASR followed by normalized WER", + stages=[ + SegmentClipExtractionStage(scratch_dir=str(scratch_dir)), + ParallelInferenceAsrNemoStage( + model_name=args.model_name, + cache_dir=None, + filepath_key="segment_audio_filepath", + pred_text_key="pred_text", + batch_size=args.asr_batch_size, + resources=Resources(gpus=1.0), + worker_count=args.asr_workers, + ), + SegmentWERStage(), + ], + ) + + try: + results = pipeline.run( + executor=XennaExecutor(config={"execution_mode": "streaming"}), + initial_tasks=initial_tasks, + ) + result_tasks = list(results or []) + threshold = args.wer_threshold_pct + if args.use_recommended_threshold: + preliminary_rows = [dict(task.data) for task in result_tasks] + threshold = float(build_wer_distribution(preliminary_rows, threshold)["recommended_threshold_pct"]) + report = write_pipeline_outputs( + result_tasks, + output_dir=output_dir, + threshold_pct=threshold, + require_fastmss_alignment=args.require_fastmss_alignment, + ) + if args.build_lhotse: + summaries = build_all_variants( + output_dir / "segments_with_wer.jsonl", + output_dir / "lhotse", + ) + logger.info("Lhotse variants: {}", summaries) + logger.info("WER report: {}", report) + logger.info("Outputs: {}", output_dir) + return 0 + finally: + shutil.rmtree(scratch_dir, ignore_errors=True) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/stages.py b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/stages.py new file mode 100644 index 0000000000..89edddc8ba --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/parakeet_wer/stages.py @@ -0,0 +1,207 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Custom Curator stages for segment extraction and per-segment WER.""" + +from __future__ import annotations + +import hashlib +import re +import unicodedata +import wave +from dataclasses import dataclass, field +from pathlib import Path + +from num2words import num2words + +from nemo_curator.stages.audio.inference.asr.asr_nemo import InferenceAsrNemoStage +from nemo_curator.stages.base import ProcessingStage +from nemo_curator.stages.resources import Resources +from nemo_curator.tasks import AudioTask + +_NUMBER_TOKEN = re.compile(r"^\d+$") +_WORD_TOKEN = re.compile(r"[a-z0-9]+(?:'[a-z0-9]+)?") + + +@dataclass +class ParallelInferenceAsrNemoStage(InferenceAsrNemoStage): + """NeMo ASR stage with an explicit number of one-GPU workers.""" + + worker_count: int = 1 + name: str = "ParakeetSegmentASR" + + def num_workers(self) -> int: + return self.worker_count + + +def normalize_wer_text(text: str) -> str: + """Normalize reference and ASR text consistently before WER.""" + folded = "".join( + character + for character in unicodedata.normalize("NFKD", text.casefold()) + if unicodedata.category(character) != "Mn" + ) + words: list[str] = [] + for token in _WORD_TOKEN.findall(folded): + if _NUMBER_TOKEN.fullmatch(token): + spoken = str(num2words(int(token), lang="en")) + words.extend(_WORD_TOKEN.findall(spoken.casefold())) + else: + words.append(token) + return " ".join(words) + + +def word_error_counts(reference: str, hypothesis: str) -> dict[str, int | float | None]: + """Compute Levenshtein WER and S/D/I counts for normalized text.""" + ref_words = normalize_wer_text(reference).split() + hyp_words = normalize_wer_text(hypothesis).split() + if not ref_words: + return { + "reference_words": 0, + "hypothesis_words": len(hyp_words), + "substitutions": 0, + "deletions": 0, + "insertions": len(hyp_words), + "errors": len(hyp_words), + "wer_pct": 0.0 if not hyp_words else None, + } + + previous = [(index, 0, 0, index) for index in range(len(hyp_words) + 1)] + for ref_index, ref_word in enumerate(ref_words, start=1): + current = [(ref_index, 0, ref_index, 0)] + for hyp_index, hyp_word in enumerate(hyp_words, start=1): + if ref_word == hyp_word: + current.append(previous[hyp_index - 1]) + continue + sub = previous[hyp_index - 1] + delete = previous[hyp_index] + insert = current[hyp_index - 1] + candidates = ( + (sub[0] + 1, sub[1] + 1, sub[2], sub[3]), + (delete[0] + 1, delete[1], delete[2] + 1, delete[3]), + (insert[0] + 1, insert[1], insert[2], insert[3] + 1), + ) + current.append(min(candidates, key=lambda item: (item[0], item[3], item[2], item[1]))) + previous = current + + errors, substitutions, deletions, insertions = previous[-1] + return { + "reference_words": len(ref_words), + "hypothesis_words": len(hyp_words), + "substitutions": substitutions, + "deletions": deletions, + "insertions": insertions, + "errors": errors, + "wer_pct": round(100.0 * errors / len(ref_words), 6), + } + + +@dataclass +class SegmentClipExtractionStage(ProcessingStage[AudioTask, AudioTask]): + """Extract exact manifest intervals from masked mono 16 kHz WAVs.""" + + scratch_dir: str = "" + sample_rate: int = 16000 + minimum_clip_duration: float = 0.1 + name: str = "ExtractMaskedSegment" + resources: Resources = field(default_factory=lambda: Resources(cpus=1.0)) + batch_size: int = 1 + + def inputs(self) -> tuple[list[str], list[str]]: + return [], ["audio_filepath", "start", "end"] + + def outputs(self) -> tuple[list[str], list[str]]: + return [], ["segment_audio_filepath", "source_audio_filepath", "clip_start", "clip_end"] + + def process(self, task: AudioTask) -> AudioTask: + source = Path(task.data["audio_filepath"]) + start = float(task.data["start"]) + end = float(task.data["end"]) + if end <= start: + msg = f"invalid segment interval: {start} >= {end}" + raise ValueError(msg) + + digest = hashlib.sha256( + f"{task.data['recording_id']}:{task.data['segment_index']}:{start}:{end}".encode() + ).hexdigest()[:20] + destination = Path(self.scratch_dir) / f"{digest}.wav" + destination.parent.mkdir(parents=True, exist_ok=True) + + with wave.open(str(source), "rb") as reader: + if reader.getframerate() != self.sample_rate or reader.getnchannels() != 1: + msg = ( + f"masked audio must be mono {self.sample_rate} Hz WAV: " + f"{source} is {reader.getnchannels()}ch/{reader.getframerate()}Hz" + ) + raise ValueError(msg) + total_frames = reader.getnframes() + start_frame = min(total_frames, max(0, round(start * self.sample_rate))) + end_frame = min(total_frames, round(end * self.sample_rate)) + minimum_frames = max(1, round(self.minimum_clip_duration * self.sample_rate)) + if end_frame - start_frame < minimum_frames: + center_frame = round((start + end) * self.sample_rate / 2) + start_frame = max(0, center_frame - minimum_frames // 2) + end_frame = min(total_frames, start_frame + minimum_frames) + start_frame = max(0, end_frame - minimum_frames) + reader.setpos(start_frame) + frames = reader.readframes(max(0, end_frame - start_frame)) + sample_width = reader.getsampwidth() + compression_type = reader.getcomptype() + compression_name = reader.getcompname() + + with wave.open(str(destination), "wb") as writer: + writer.setnchannels(1) + writer.setsampwidth(sample_width) + writer.setframerate(self.sample_rate) + writer.setcomptype(compression_type, compression_name) + writer.writeframes(frames) + + task.data["source_audio_filepath"] = str(source) + task.data["segment_audio_filepath"] = str(destination) + task.data["clip_start"] = round(start_frame / self.sample_rate, 6) + task.data["clip_end"] = round(end_frame / self.sample_rate, 6) + return task + + +@dataclass +class SegmentWERStage(ProcessingStage[AudioTask, AudioTask]): + """Compute normalized per-segment WER and remove the temporary clip.""" + + reference_key: str = "text_raw" + hypothesis_key: str = "pred_text" + cleanup_clips: bool = True + name: str = "ComputeSegmentWER" + resources: Resources = field(default_factory=lambda: Resources(cpus=1.0)) + batch_size: int = 1 + + def inputs(self) -> tuple[list[str], list[str]]: + return [], [self.reference_key, self.hypothesis_key, "segment_audio_filepath"] + + def outputs(self) -> tuple[list[str], list[str]]: + return [], ["wer_pct", "wer_details", "audio_filepath"] + + def process(self, task: AudioTask) -> AudioTask: + clip = Path(task.data["segment_audio_filepath"]) + try: + details = word_error_counts( + str(task.data[self.reference_key]), + str(task.data[self.hypothesis_key]), + ) + task.data["wer_details"] = details + task.data["wer_pct"] = details["wer_pct"] + task.data["audio_filepath"] = task.data["source_audio_filepath"] + return task + finally: + if self.cleanup_clips: + clip.unlink(missing_ok=True) diff --git a/tutorials/audio/david_ai_redelivered_mfa/requirements-dev.txt b/tutorials/audio/david_ai_redelivered_mfa/requirements-dev.txt new file mode 100644 index 0000000000..5e89329490 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/requirements-dev.txt @@ -0,0 +1,3 @@ +-r requirements.txt +pytest +ruff diff --git a/tutorials/audio/david_ai_redelivered_mfa/requirements.txt b/tutorials/audio/david_ai_redelivered_mfa/requirements.txt new file mode 100644 index 0000000000..5887f66568 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/requirements.txt @@ -0,0 +1,3 @@ +numpy +num2words>=0.5.12 +textgrid diff --git a/tutorials/audio/david_ai_redelivered_mfa/tests/cluster/test_cluster_launchers.py b/tutorials/audio/david_ai_redelivered_mfa/tests/cluster/test_cluster_launchers.py new file mode 100644 index 0000000000..dd33ca0cf4 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/tests/cluster/test_cluster_launchers.py @@ -0,0 +1,78 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +SUBMIT = ROOT / "cluster" / "run_multinode.sh" +NODE = ROOT / "cluster" / "run_node.sh" + + +def test_cluster_launchers_are_executable_and_do_not_copy_data() -> None: + for script in (SUBMIT, NODE): + text = script.read_text(encoding="utf-8") + assert os.access(script, os.X_OK) + assert "scp " not in text + assert "rsync " not in text + assert "--export ALL" not in text + assert "--export=ALL" not in text + + +def test_multinode_launcher_uses_explicit_slurm_array_sharding() -> None: + submit = SUBMIT.read_text(encoding="utf-8") + node = NODE.read_text(encoding="utf-8") + + assert '--array "0-$((NUM_NODES - 1))%$MAX_CONCURRENT_NODES"' in submit + assert '--export "$EXPORTS"' in submit + assert 'SHARD_COUNT="$NUM_NODES"' not in submit + assert 'EXPORTS+=",SHARD_COUNT=$NUM_NODES"' in submit + assert 'EXPORTS+=",SESSIONS_FILE=$SESSIONS_FILE"' in submit + assert "CONTAINER_MOUNTS_B64" in submit + assert "CONTAINER_MOUNTS=$CONTAINER_MOUNTS" not in submit + assert 'SHARD_INDEX="$SLURM_ARRAY_TASK_ID"' in node + assert 'SHARD_COUNT="$SHARD_COUNT"' in node + assert 'SESSIONS_FILE="$SESSIONS_FILE"' in node + + +def test_both_pipeline_variants_are_supported() -> None: + submit = SUBMIT.read_text(encoding="utf-8") + node = NODE.read_text(encoding="utf-8") + + assert "opus | wav" in submit + assert 'PIPELINE_DIR="$TUTORIAL_ROOT/$VARIANT"' in node + assert 'bash "$TUTORIAL_ROOT/cluster/run_node.sh"' in node + assert 'bash "$SCRIPT_DIR/run_node.sh"' not in node + + +def test_mfa_scratch_and_model_copies_are_isolated_per_shard_and_worker() -> None: + node = NODE.read_text(encoding="utf-8") + assert "${SLURM_JOB_ID}_${SHARD_INDEX}" in node + assert 'NODE_MFA_ROOT_DIR="$RAM_DIR/model_source"' in node + assert 'export MFA_ROOT_DIR="$NODE_MFA_ROOT_DIR"' in node + assert "stage_model" in node + + for variant in ("opus", "wav"): + ram_session = (ROOT / variant / "david_ai_ram_session.py").read_text(encoding="utf-8") + common = (ROOT / variant / "david_ai_common.py").read_text(encoding="utf-8") + aligner = (ROOT / variant / "david_ai_mfa_align.py").read_text(encoding="utf-8") + + assert 'ram_dir / "mfa_workers" / f"worker_{os.getpid()}"' in ram_session + assert "temp_parent=temp_parent / session_id" in ram_session + assert "align_result.mfa_segments == 0" in ram_session + assert 'models_dir = worker_dir / "models"' in common + assert "shutil.copy2(mfa_dict, local_dict)" in common + assert "shutil.copytree(acoustic_src, local_acoustic)" in common + assert 'env["MFA_ROOT_DIR"] = str(mfa_root)' in common + assert "mfa_subprocess_env(temp_root=temp_root, mfa_root=mfa_root)" in aligner diff --git a/tutorials/audio/david_ai_redelivered_mfa/tests/opus/conftest.py b/tutorials/audio/david_ai_redelivered_mfa/tests/opus/conftest.py new file mode 100644 index 0000000000..596ced60ba --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/tests/opus/conftest.py @@ -0,0 +1,18 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "opus")) diff --git a/tutorials/audio/david_ai_redelivered_mfa/tests/opus/test_opus_ram_session_outputs.py b/tutorials/audio/david_ai_redelivered_mfa/tests/opus/test_opus_ram_session_outputs.py new file mode 100644 index 0000000000..0bec80a0cf --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/tests/opus/test_opus_ram_session_outputs.py @@ -0,0 +1,215 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import zipfile +from pathlib import Path + +import david_ai_common as common +import david_ai_ram_session as ram_session +import pytest +from david_ai_common import ( + fastmss_textgrid_path, + mixed_speaker_audio_path, + recording_id, + recording_textgrid_path, + session_mixed_audio_path, + session_rttm_path, + session_textgrid_path, +) +from david_ai_manifest import normalize_text, resolve_speaker_audio_path +from david_ai_mfa_align import SessionAlignResult +from david_ai_ram_lhotse import write_all_textgrids +from stage_ram_session_pipeline import filter_sessions_from_file, sessions_without_done_flags + + +def _manifest_row(tmp_path: Path, session_id: str, speaker_id: str, start: float, end: float) -> dict: + rec_id = recording_id(speaker_id, session_id) + audio_path = tmp_path / f"{rec_id}.wav" + audio_path.write_bytes(b"source") + return { + "session_id": session_id, + "speaker_id": speaker_id, + "recording_id": rec_id, + "segment_index": 0, + "start": start, + "end": end, + "audio_filepath": str(audio_path), + "audio_filepath_16k": str(audio_path), + } + + +def test_speaker_audio_resolution_priority_and_fallback(tmp_path: Path) -> None: + speaker_id = "speaker" + preprocessed = tmp_path / f"{speaker_id}_preprocessed.wav" + ordinary = tmp_path / f"{speaker_id}.wav" + postprocessed = tmp_path / f"{speaker_id}_postprocessed.wav" + postprocess = tmp_path / f"{speaker_id}_postprocess.wav" + + with pytest.raises(FileNotFoundError): + resolve_speaker_audio_path(tmp_path, speaker_id) + preprocessed.touch() + assert resolve_speaker_audio_path(tmp_path, speaker_id) == preprocessed + ordinary.touch() + assert resolve_speaker_audio_path(tmp_path, speaker_id) == ordinary + postprocessed.touch() + assert resolve_speaker_audio_path(tmp_path, speaker_id) == postprocessed + postprocess.touch() + assert resolve_speaker_audio_path(tmp_path, speaker_id) == postprocess + + +def test_transcript_normalization_is_self_contained() -> None: + assert normalize_text("Café costs 2 dollars — okay!") == "cafe costs two dollars okay" + + +def test_g2p_zip_is_preextracted_into_worker_private_models(tmp_path: Path, monkeypatch) -> None: + archive_path = tmp_path / "g2p.zip" + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("model/meta.json", "{}") + archive.writestr("model/model.fst", b"fst-data") + archive.writestr("model/phones.sym", "1 a") + monkeypatch.setattr(common, "resolve_mfa_g2p_model", lambda _: archive_path) + + model_dir = Path(common._worker_g2p_arg(tmp_path / "worker-models", "test-model")) + + assert model_dir == tmp_path / "worker-models" / "g2p" / "model" + assert (model_dir / "model.fst").read_bytes() == b"fst-data" + + +def test_mix_uses_manifest_boundaries_and_persists_speaker_tracks(tmp_path: Path, monkeypatch) -> None: + session_id = "session" + rows = [ + _manifest_row(tmp_path, session_id, "speaker-a", 1.0, 2.0), + _manifest_row(tmp_path, session_id, "speaker-b", 3.0, 4.0), + ] + captured: dict[str, tuple[list[tuple[float, float]], float]] = {} + + def fake_prepare(src: Path, dst: Path, *, speech_intervals, boundary_indent, **kwargs) -> bool: + captured[src.name] = (speech_intervals, boundary_indent) + dst.parent.mkdir(parents=True, exist_ok=True) + dst.write_bytes(f"prepared:{src.name}".encode()) + return True + + def fake_mix(audio_paths: list[Path], output_path: Path, *, opus_bitrate: str) -> bool: + output_path.write_bytes(b"|".join(path.read_bytes() for path in audio_paths)) + return True + + monkeypatch.setattr(ram_session, "prepare_speaker_audio_for_session_mix", fake_prepare) + monkeypatch.setattr(ram_session, "mix_audio_files", fake_mix) + + audio_mixed_dir = tmp_path / "audio_mixed" + ram_session._mix_session_from_manifest( + session_id, + rows, + audio_mixed_dir=audio_mixed_dir, + session_ram=tmp_path / "scratch", + opus_bitrate="32k", + noise_level=0.0002, + stitch_ms=5.0, + boundary_offset=0.5, + rec_durations={row["recording_id"]: 10.0 for row in rows}, + ) + + assert captured[f"{rows[0]['recording_id']}.wav"] == ([(1.0, 2.0)], 0.5) + assert captured[f"{rows[1]['recording_id']}.wav"] == ([(3.0, 4.0)], 0.5) + for row in rows: + assert mixed_speaker_audio_path(audio_mixed_dir, row["speaker_id"], session_id).is_file() + assert session_mixed_audio_path(audio_mixed_dir, session_id).is_file() + + +def test_write_all_textgrids_writes_both_variants_when_fastmss_is_empty(tmp_path: Path) -> None: + session_id = "session" + rec_id = recording_id("speaker", session_id) + result = SessionAlignResult( + ok=True, + fb_words=[(1.0, 2.0, "speech", "speaker")], + audio_duration=3.0, + recordings=[ + { + "session_id": session_id, + "recording_id": rec_id, + "merged_words": [], + "fb_words": [[1.0, 2.0, "speech"]], + "audio_duration": 3.0, + } + ], + ) + + write_all_textgrids(result, tmp_path) + + assert session_textgrid_path(tmp_path, session_id, variant="fastmss").is_file() + assert session_textgrid_path(tmp_path, session_id, variant="ordinary").is_file() + assert fastmss_textgrid_path(tmp_path, rec_id).is_file() + assert recording_textgrid_path(tmp_path, rec_id, variant="ordinary").is_file() + + +def test_done_flag_is_written_only_after_all_outputs_exist(tmp_path: Path) -> None: + session_id = "session" + speaker_id = "speaker" + row = _manifest_row(tmp_path, session_id, speaker_id, 1.0, 2.0) + work_dir = tmp_path / "work" + audio_mixed_dir = work_dir / "audio_mixed" + textgrid_dir = work_dir / "textgrids" + + with pytest.raises(ram_session.PipelineError): + ram_session._finalize_session_success( + session_id, + [row], + work_dir=work_dir, + audio_mixed_dir=audio_mixed_dir, + textgrid_dir=textgrid_dir, + ) + assert not ram_session.session_done_path(work_dir, session_id).exists() + + rec_id = row["recording_id"] + outputs = [ + session_mixed_audio_path(audio_mixed_dir, session_id), + session_rttm_path(audio_mixed_dir, session_id), + session_textgrid_path(textgrid_dir, session_id, variant="ordinary"), + session_textgrid_path(textgrid_dir, session_id, variant="fastmss"), + mixed_speaker_audio_path(audio_mixed_dir, speaker_id, session_id), + recording_textgrid_path(textgrid_dir, rec_id, variant="ordinary"), + fastmss_textgrid_path(textgrid_dir, rec_id), + ] + for output in outputs: + output.parent.mkdir(parents=True, exist_ok=True) + output.write_bytes(b"output") + + ram_session._finalize_session_success( + session_id, + [row], + work_dir=work_dir, + audio_mixed_dir=audio_mixed_dir, + textgrid_dir=textgrid_dir, + ) + assert ram_session.session_done_path(work_dir, session_id).read_text() == "ok\n" + + +def test_parallel_resume_selects_only_sessions_without_done_flags(tmp_path: Path) -> None: + work_dir = tmp_path / "work" + sessions = [tmp_path / "session-a", tmp_path / "session-b"] + ram_session._mark_session_done(work_dir, "session-a") + + pending = sessions_without_done_flags(sessions, work_dir) + + assert [session.name for session in pending] == ["session-b"] + + +def test_session_list_restricts_discovered_sessions(tmp_path: Path) -> None: + sessions = [tmp_path / name for name in ("session-a", "session-b", "session-c")] + sessions_file = tmp_path / "sessions.txt" + sessions_file.write_text("# subset\nsession-c\n\nsession-a\nmissing-session\n", encoding="utf-8") + + selected = filter_sessions_from_file(sessions, sessions_file) + + assert [session.name for session in selected] == ["session-a", "session-c"] diff --git a/tutorials/audio/david_ai_redelivered_mfa/tests/parakeet_wer/conftest.py b/tutorials/audio/david_ai_redelivered_mfa/tests/parakeet_wer/conftest.py new file mode 100644 index 0000000000..d316375750 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/tests/parakeet_wer/conftest.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[5])) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "parakeet_wer")) diff --git a/tutorials/audio/david_ai_redelivered_mfa/tests/parakeet_wer/test_parakeet_wer_pipeline.py b/tutorials/audio/david_ai_redelivered_mfa/tests/parakeet_wer/test_parakeet_wer_pipeline.py new file mode 100644 index 0000000000..c53074fa5f --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/tests/parakeet_wer/test_parakeet_wer_pipeline.py @@ -0,0 +1,274 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import wave +from pathlib import Path + +from build_lhotse_variants import VARIANTS, segment_matches_variant +from manifest import ( + SegmentTaskConfig, + build_segment_tasks, + build_wer_distribution, + segment_alignments, + write_pipeline_outputs, +) +from stages import ( + ParallelInferenceAsrNemoStage, + SegmentClipExtractionStage, + normalize_wer_text, + word_error_counts, +) + +from nemo_curator.tasks import AudioTask + +ROOT = Path(__file__).resolve().parents[2] + + +def test_normalization_and_wer_above_one_hundred_percent() -> None: + assert normalize_wer_text("Café costs 2 dollars!") == "cafe costs two dollars" + + details = word_error_counts("hello", "hello extra words") + + assert details["insertions"] == 2 + assert details["wer_pct"] == 200.0 + + +def test_fastmss_alignment_is_clipped_and_made_segment_relative() -> None: + words = [(0.5, 1.2, "first"), (1.8, 2.4, "second"), (3.0, 3.5, "outside")] + + absolute, relative = segment_alignments(words, 1.0, 2.0) + + assert absolute == [ + {"word": "first", "start": 1.0, "end": 1.2}, + {"word": "second", "start": 1.8, "end": 2.0}, + ] + assert relative == [ + {"symbol": "first", "start": 0.0, "duration": 0.2}, + {"symbol": "second", "start": 0.8, "duration": 0.2}, + ] + + +def test_distribution_reports_bounded_recommended_threshold() -> None: + rows = [{"wer_pct": value} for value in (0.0, 5.0, 10.0, 25.0, 50.0, 150.0)] + + report = build_wer_distribution(rows, applied_threshold_pct=100.0) + + assert report["segments_with_wer"] == 6 + assert 25.0 <= report["recommended_threshold_pct"] <= 100.0 + assert sum(bucket["count"] for bucket in report["histogram"]) == 6 + + +def test_lhotse_wer_variants_are_nested_and_require_alignment() -> None: + aligned = {"alignment": [{"symbol": "word"}]} + + assert segment_matches_variant({**aligned, "wer_pct": 0.0}, VARIANTS[0]) + assert segment_matches_variant({**aligned, "wer_pct": 0.0}, VARIANTS[1]) + assert segment_matches_variant({**aligned, "wer_pct": 0.0}, VARIANTS[2]) + assert not segment_matches_variant({**aligned, "wer_pct": 5.0}, VARIANTS[0]) + assert segment_matches_variant({**aligned, "wer_pct": 5.0}, VARIANTS[1]) + assert segment_matches_variant({**aligned, "wer_pct": 5.0}, VARIANTS[2]) + assert not segment_matches_variant({**aligned, "wer_pct": 100.01}, VARIANTS[2]) + assert not segment_matches_variant({"alignment": [], "wer_pct": 0.0}, VARIANTS[0]) + + +def test_output_writer_filters_high_wer_and_missing_alignment(tmp_path: Path) -> None: + base = { + "audio_filepath": "/synthetic/masked.wav", + "start": 1.0, + "end": 2.0, + "duration": 1.0, + "text": "reference", + "text_raw": "Reference", + "pred_text": "reference", + "session_id": "session", + "speaker_id": "speaker", + "recording_id": "recording", + "fastmss_textgrid": "/synthetic/recording_fastmss.TextGrid", + "words": [{"word": "reference", "start": 1.1, "end": 1.5}], + } + tasks = [ + AudioTask(data={**base, "segment_index": 0, "wer_pct": 0.0, "alignment": [{"symbol": "reference"}]}), + AudioTask(data={**base, "segment_index": 1, "wer_pct": 150.0, "alignment": [{"symbol": "bad"}]}), + AudioTask(data={**base, "segment_index": 2, "wer_pct": 0.0, "alignment": []}), + ] + + report = write_pipeline_outputs( + tasks, + output_dir=tmp_path, + threshold_pct=100.0, + require_fastmss_alignment=True, + ) + + manifest_rows = [ + json.loads(line) for line in (tmp_path / "manifests" / "recording.jsonl").read_text().splitlines() + ] + audit_rows = [json.loads(line) for line in (tmp_path / "segments_with_wer.jsonl").read_text().splitlines()] + assert report["kept_segments"] == 1 + assert [row["segment_index"] for row in manifest_rows] == [0] + assert audit_rows[1]["rejection_reasons"] == ["wer_above_threshold"] + assert audit_rows[2]["rejection_reasons"] == ["missing_fastmss_alignment"] + + +def test_segment_clip_extraction_uses_exact_interval(tmp_path: Path) -> None: + source = tmp_path / "source.wav" + with wave.open(str(source), "wb") as writer: + writer.setnchannels(1) + writer.setsampwidth(2) + writer.setframerate(16000) + writer.writeframes(b"\x01\x00" * 32000) + task = AudioTask( + data={ + "audio_filepath": str(source), + "recording_id": "recording", + "segment_index": 0, + "start": 0.5, + "end": 1.25, + } + ) + + result = SegmentClipExtractionStage(scratch_dir=str(tmp_path / "scratch")).process(task) + + with wave.open(result.data["segment_audio_filepath"], "rb") as reader: + assert reader.getframerate() == 16000 + assert reader.getnchannels() == 1 + assert reader.getnframes() == 12000 + + +def test_ultrashort_segment_clip_is_padded_to_model_minimum(tmp_path: Path) -> None: + source = tmp_path / "source.wav" + with wave.open(str(source), "wb") as writer: + writer.setnchannels(1) + writer.setsampwidth(2) + writer.setframerate(16000) + writer.writeframes(b"\x01\x00" * 32000) + task = AudioTask( + data={ + "audio_filepath": str(source), + "recording_id": "recording", + "segment_index": 0, + "start": 1.0, + "end": 1.001, + } + ) + + result = SegmentClipExtractionStage(scratch_dir=str(tmp_path / "scratch")).process(task) + + with wave.open(result.data["segment_audio_filepath"], "rb") as reader: + assert reader.getnframes() == 1600 + assert result.data["start"] == 1.0 + assert result.data["end"] == 1.001 + assert round(result.data["clip_end"] - result.data["clip_start"], 6) == 0.1 + + +def test_build_segment_tasks_joins_manifest_masked_audio_and_fastmss(tmp_path: Path) -> None: + data_root = tmp_path / "data" + masked_dir = tmp_path / "audio_16k_masked" + textgrid_dir = tmp_path / "textgrids" + session_dir = data_root / "session" + session_dir.mkdir(parents=True) + masked_dir.mkdir() + textgrid_dir.mkdir() + (session_dir / "machine_generated_transcript.json").write_text( + json.dumps( + { + "transcript": [ + {"speaker": "speaker", "start": 0.0, "end": 1.0, "text": "first"}, + {"speaker": "speaker", "start": 1.0, "end": 2.0, "text": "second"}, + ] + } + ), + encoding="utf-8", + ) + recording = "speaker_session_postprocessed" + (masked_dir / f"{recording}.wav").touch() + (textgrid_dir / f"{recording}_fastmss.TextGrid").write_text( + """File type = "ooTextFile" +Object class = "TextGrid" + +xmin = 0 +xmax = 2 +tiers? +size = 1 +item []: + item [1]: + class = "IntervalTier" + name = "words" + xmin = 0 + xmax = 2 + intervals: size = 2 + intervals [1]: + xmin = 0.1 + xmax = 0.5 + text = "first" + intervals [2]: + xmin = 1.1 + xmax = 1.6 + text = "second" +""", + encoding="utf-8", + ) + + tasks = build_segment_tasks( + SegmentTaskConfig( + data_root=data_root, + masked_audio_dir=masked_dir, + textgrid_dir=textgrid_dir, + sessions_file=None, + shard_count=1, + shard_index=0, + ) + ) + + assert len(tasks) == 2 + assert tasks[0].data["recording_id"] == recording + assert tasks[0].data["alignment"] == [{"symbol": "first", "start": 0.1, "duration": 0.4}] + assert tasks[1].data["alignment"] == [{"symbol": "second", "start": 0.1, "duration": 0.5}] + + +def test_local_two_gpu_launcher_uses_one_xenna_cluster() -> None: + launcher = (ROOT / "parakeet_wer" / "run_local_2gpu.sh").read_text(encoding="utf-8") + + assert 'GPU_IDS="${GPU_IDS:-0,1}"' in launcher + assert 'CUDA_VISIBLE_DEVICES="$GPU_0,$GPU_1"' in launcher + assert "SHARD_COUNT=1" in launcher + assert "SHARD_INDEX=0" in launcher + assert "ASR_WORKERS=2" in launcher + assert 'SCRATCH_DIR="$SCRATCH_DIR"' in launcher + assert "ASRModel.from_pretrained" in launcher + assert 'export NEMO_CACHE_DIR="$MODEL_CACHE_DIR"' in launcher + assert "cache_dir=" not in launcher + assert "analyze_wer_distribution.py" in launcher + assert "local_2gpu.log" in launcher + + +def test_cluster_launcher_builds_and_merges_lhotse_after_array() -> None: + cluster_dir = ROOT / "parakeet_wer" / "cluster" + launcher = (cluster_dir / "run_multinode.sh").read_text(encoding="utf-8") + node = (cluster_dir / "run_node.sh").read_text(encoding="utf-8") + merge = (cluster_dir / "merge_outputs.sh").read_text(encoding="utf-8") + + assert '--dependency "afterok:$ARRAY_JOB_ID"' in launcher + assert "BUILD_LHOTSE=1 \\" in node + assert "merge_lhotse_variants.py" in merge + assert "analyze_wer_distribution.py" in merge + assert "scp " not in launcher + node + merge + assert "rsync " not in launcher + node + merge + + +def test_parallel_asr_stage_requests_two_workers() -> None: + stage = ParallelInferenceAsrNemoStage(model_name="synthetic-model", worker_count=2) + + assert stage.num_workers() == 2 + assert stage.resources.gpus == 0 diff --git a/tutorials/audio/david_ai_redelivered_mfa/tests/wav/conftest.py b/tutorials/audio/david_ai_redelivered_mfa/tests/wav/conftest.py new file mode 100644 index 0000000000..9cfb197e9f --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/tests/wav/conftest.py @@ -0,0 +1,18 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "wav")) diff --git a/tutorials/audio/david_ai_redelivered_mfa/tests/wav/test_wav_ram_session_outputs.py b/tutorials/audio/david_ai_redelivered_mfa/tests/wav/test_wav_ram_session_outputs.py new file mode 100644 index 0000000000..c4a6e6b60e --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/tests/wav/test_wav_ram_session_outputs.py @@ -0,0 +1,251 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import zipfile +from pathlib import Path + +import david_ai_common as common +import david_ai_ram_session as ram_session +import pytest +from david_ai_common import ( + fastmss_textgrid_path, + masked_speaker_audio_path, + masked_speaker_rttm_path, + recording_id, + recording_textgrid_path, + session_mixed_audio_path, + session_rttm_path, + session_textgrid_path, +) +from david_ai_manifest import normalize_text, resolve_speaker_audio_path +from david_ai_mfa_align import SessionAlignResult +from david_ai_ram_lhotse import write_all_textgrids +from stage_ram_session_pipeline import filter_sessions_from_file, sessions_without_done_flags + + +def _manifest_row(tmp_path: Path, session_id: str, speaker_id: str, start: float, end: float) -> dict: + rec_id = recording_id(speaker_id, session_id) + audio_path = tmp_path / f"{rec_id}.wav" + audio_path.write_bytes(b"source") + return { + "session_id": session_id, + "speaker_id": speaker_id, + "recording_id": rec_id, + "segment_index": 0, + "start": start, + "end": end, + "audio_filepath": str(audio_path), + "audio_filepath_16k": str(audio_path), + } + + +def test_speaker_audio_resolution_priority_and_fallback(tmp_path: Path) -> None: + speaker_id = "speaker" + preprocessed = tmp_path / f"{speaker_id}_preprocessed.wav" + ordinary = tmp_path / f"{speaker_id}.wav" + postprocessed = tmp_path / f"{speaker_id}_postprocessed.wav" + postprocess = tmp_path / f"{speaker_id}_postprocess.wav" + + with pytest.raises(FileNotFoundError): + resolve_speaker_audio_path(tmp_path, speaker_id) + preprocessed.touch() + assert resolve_speaker_audio_path(tmp_path, speaker_id) == preprocessed + ordinary.touch() + assert resolve_speaker_audio_path(tmp_path, speaker_id) == ordinary + postprocessed.touch() + assert resolve_speaker_audio_path(tmp_path, speaker_id) == postprocessed + postprocess.touch() + assert resolve_speaker_audio_path(tmp_path, speaker_id) == postprocess + + +def test_transcript_normalization_is_self_contained() -> None: + assert normalize_text("Café costs 2 dollars — okay!") == "cafe costs two dollars okay" + + +def test_g2p_zip_is_preextracted_into_worker_private_models(tmp_path: Path, monkeypatch) -> None: + archive_path = tmp_path / "g2p.zip" + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("model/meta.json", "{}") + archive.writestr("model/model.fst", b"fst-data") + archive.writestr("model/phones.sym", "1 a") + monkeypatch.setattr(common, "resolve_mfa_g2p_model", lambda _: archive_path) + + model_dir = Path(common._worker_g2p_arg(tmp_path / "worker-models", "test-model")) + + assert model_dir == tmp_path / "worker-models" / "g2p" / "model" + assert (model_dir / "model.fst").read_bytes() == b"fst-data" + + +def test_mix_persists_masked_speaker_and_mixed_wavs(tmp_path: Path, monkeypatch) -> None: + session_id = "session" + rows = [ + _manifest_row(tmp_path, session_id, "speaker-a", 1.0, 2.0), + _manifest_row(tmp_path, session_id, "speaker-b", 3.0, 4.0), + ] + captured: dict[str, tuple[list[tuple[float, float]], float]] = {} + + def fake_prepare(src: Path, dst: Path, *, speech_intervals, boundary_indent, **kwargs) -> bool: + assert dst.suffix == ".wav" + captured[src.name] = (speech_intervals, boundary_indent) + dst.parent.mkdir(parents=True, exist_ok=True) + dst.write_bytes(f"prepared:{src.name}".encode()) + return True + + def fake_mix(audio_paths: list[Path], output_path: Path) -> bool: + assert output_path.suffix == ".wav" + assert all(path.suffix == ".wav" for path in audio_paths) + output_path.write_bytes(b"|".join(path.read_bytes() for path in audio_paths)) + return True + + monkeypatch.setattr(ram_session, "prepare_speaker_audio_for_session_mix", fake_prepare) + monkeypatch.setattr(ram_session, "mix_audio_files", fake_mix) + + audio_masked_dir = tmp_path / "audio_16k_masked" + audio_mixed_dir = tmp_path / "audio_mixed" + ram_session._mix_session_from_manifest( + session_id, + rows, + audio_masked_dir=audio_masked_dir, + audio_mixed_dir=audio_mixed_dir, + session_ram=tmp_path / "scratch", + noise_level=0.0002, + stitch_ms=5.0, + boundary_offset=0.5, + rec_durations={row["recording_id"]: 10.0 for row in rows}, + ) + + assert captured[f"{rows[0]['recording_id']}.wav"] == ([(1.0, 2.0)], 0.5) + assert captured[f"{rows[1]['recording_id']}.wav"] == ([(3.0, 4.0)], 0.5) + for row in rows: + assert masked_speaker_audio_path(audio_masked_dir, row["speaker_id"], session_id).is_file() + assert session_mixed_audio_path(audio_mixed_dir, session_id).suffix == ".wav" + assert session_mixed_audio_path(audio_mixed_dir, session_id).is_file() + + +def test_masked_speaker_rttms_are_filtered_from_session_rttm(tmp_path: Path) -> None: + session_id = "session" + rows = [ + _manifest_row(tmp_path, session_id, "speaker-a", 1.0, 2.0), + _manifest_row(tmp_path, session_id, "speaker-b", 3.0, 4.0), + ] + lines = [ + f"SPEAKER {session_id} 1 1.000000 1.000000 speaker-a ", + f"SPEAKER {session_id} 1 3.000000 1.000000 speaker-b ", + ] + + ram_session._write_masked_speaker_rttms( + session_id, + rows, + lines, + audio_masked_dir=tmp_path, + ) + + for speaker_id in ("speaker-a", "speaker-b"): + text = masked_speaker_rttm_path(tmp_path, speaker_id, session_id).read_text() + assert recording_id(speaker_id, session_id) in text + assert f" {speaker_id} " in text + other = "speaker-b" if speaker_id == "speaker-a" else "speaker-a" + assert f" {other} " not in text + + +def test_write_all_textgrids_writes_both_variants_when_fastmss_is_empty(tmp_path: Path) -> None: + session_id = "session" + rec_id = recording_id("speaker", session_id) + result = SessionAlignResult( + ok=True, + fb_words=[(1.0, 2.0, "speech", "speaker")], + audio_duration=3.0, + recordings=[ + { + "session_id": session_id, + "recording_id": rec_id, + "merged_words": [], + "fb_words": [[1.0, 2.0, "speech"]], + "audio_duration": 3.0, + } + ], + ) + + write_all_textgrids(result, tmp_path) + + assert session_textgrid_path(tmp_path, session_id, variant="fastmss").is_file() + assert session_textgrid_path(tmp_path, session_id, variant="ordinary").is_file() + assert fastmss_textgrid_path(tmp_path, rec_id).is_file() + assert recording_textgrid_path(tmp_path, rec_id, variant="ordinary").is_file() + + +def test_done_flag_requires_mixed_wav_rttm_and_textgrids(tmp_path: Path) -> None: + session_id = "session" + speaker_id = "speaker" + row = _manifest_row(tmp_path, session_id, speaker_id, 1.0, 2.0) + work_dir = tmp_path / "work" + audio_masked_dir = work_dir / "audio_16k_masked" + audio_mixed_dir = work_dir / "audio_mixed" + textgrid_dir = work_dir / "textgrids" + + with pytest.raises(ram_session.PipelineError): + ram_session._finalize_session_success( + session_id, + [row], + work_dir=work_dir, + audio_masked_dir=audio_masked_dir, + audio_mixed_dir=audio_mixed_dir, + textgrid_dir=textgrid_dir, + ) + assert not ram_session.session_done_path(work_dir, session_id).exists() + + rec_id = row["recording_id"] + outputs = [ + session_mixed_audio_path(audio_mixed_dir, session_id), + session_rttm_path(audio_mixed_dir, session_id), + session_textgrid_path(textgrid_dir, session_id, variant="ordinary"), + session_textgrid_path(textgrid_dir, session_id, variant="fastmss"), + masked_speaker_audio_path(audio_masked_dir, speaker_id, session_id), + masked_speaker_rttm_path(audio_masked_dir, speaker_id, session_id), + recording_textgrid_path(textgrid_dir, rec_id, variant="ordinary"), + fastmss_textgrid_path(textgrid_dir, rec_id), + ] + for output in outputs: + output.parent.mkdir(parents=True, exist_ok=True) + output.write_bytes(b"output") + + ram_session._finalize_session_success( + session_id, + [row], + work_dir=work_dir, + audio_masked_dir=audio_masked_dir, + audio_mixed_dir=audio_mixed_dir, + textgrid_dir=textgrid_dir, + ) + assert ram_session.session_done_path(work_dir, session_id).read_text() == "ok\n" + + +def test_parallel_resume_selects_only_sessions_without_done_flags(tmp_path: Path) -> None: + work_dir = tmp_path / "work" + sessions = [tmp_path / "session-a", tmp_path / "session-b"] + ram_session._mark_session_done(work_dir, "session-a") + + pending = sessions_without_done_flags(sessions, work_dir) + + assert [session.name for session in pending] == ["session-b"] + + +def test_session_list_restricts_discovered_sessions(tmp_path: Path) -> None: + sessions = [tmp_path / name for name in ("session-a", "session-b", "session-c")] + sessions_file = tmp_path / "sessions.txt" + sessions_file.write_text("# subset\nsession-c\n\nsession-a\nmissing-session\n", encoding="utf-8") + + selected = filter_sessions_from_file(sessions, sessions_file) + + assert [session.name for session in selected] == ["session-a", "session-c"] diff --git a/tutorials/audio/david_ai_redelivered_mfa/wav/README.md b/tutorials/audio/david_ai_redelivered_mfa/wav/README.md new file mode 100644 index 0000000000..925d58e285 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/wav/README.md @@ -0,0 +1,69 @@ +# David AI on-the-fly MFA E2E — 16 kHz WAV outputs + +This is an independent copy of the David AI E2E pipeline. It never creates Opus +files or a raw `audio_16k` directory. + +For every raw session it: + +1. Reads `machine_generated_transcript.json` and resolves one supported WAV per + speaker (`_postprocess`, `_postprocessed`, ordinary `.wav`, then `_preprocessed`). +2. Normalizes transcript rows in memory. +3. Runs MFA with the base `english_us_arpa` dictionary and runtime G2P. +4. Writes the session RTTM. +5. Creates 16 kHz speaker WAVs in node-local scratch. +6. Replaces pauses outside original manifest boundaries ±0.5 seconds with white + noise (`0.0002`, 5 ms smoothing). +7. Saves every masked speaker WAV to `audio_16k_masked` and filters the session + RTTM into a matching per-speaker RTTM. +8. Mixes the masked tracks and saves one mono 16 kHz PCM WAV for the session. +9. Writes ordinary (MFA + fallback) and FastMSS TextGrids. +10. Validates outputs and writes `.done/sessions/.done`. + +Validated done flags provide resume behavior. Re-running the same command skips +completed sessions and processes only sessions without done flags. + +## Run + +```bash +DATA_ROOT=/path/to/raw/sessions \ +WORK_DIR=/path/to/output \ +WORKERS=16 \ +MFA_NUM_JOBS=2 \ +bash run_david_ai_mfa_ram_session.sh +``` + +For a multi-node SLURM run: + +```bash +cd .. + +VARIANT=wav \ +DATA_ROOT=/shared/path/to/raw/sessions \ +WORK_DIR=/shared/path/to/output \ +NUM_NODES=8 \ +WORKERS_PER_NODE=16 \ +MFA_NUM_JOBS=2 \ +bash cluster/run_multinode.sh +``` + +See `../cluster/README.md` for environment and scheduler options. + +## Persistent outputs + +```text +/ +├── audio_16k_masked/ +│ ├── __postprocessed.wav +│ └── __postprocessed.rttm +├── audio_mixed/ +│ ├── .wav # mono 16 kHz PCM s16le +│ └── .rttm +├── textgrids/ +│ ├── .TextGrid +│ ├── _fastmss.TextGrid +│ ├── .TextGrid +│ └── _fastmss.TextGrid +├── .done/ +│ └── sessions/.done +└── logs/ +``` diff --git a/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_common.py b/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_common.py new file mode 100644 index 0000000000..ac1b7d8fa8 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_common.py @@ -0,0 +1,1661 @@ +"""Shared helpers for the David AI MFA pipeline.""" + +from __future__ import annotations + +import json +import logging +import os +import re +import shutil +import subprocess +import threading +import traceback +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from pathlib import Path +from typing import TYPE_CHECKING, TypeVar + +if TYPE_CHECKING: + from collections.abc import Callable + +MFA_ROOT_DIR_DEFAULT = "~/MFA_models" + +POSTPROCESSED_RE = re.compile(r"^(.+)_postprocessed\.wav$") +SILENCE_TOKENS = {"", "sil", "sp", "spn", ""} + +logger = logging.getLogger(__name__) + +# Hard wall-clock cap for any single ffmpeg subprocess. Without this, a wedged +# ffmpeg (e.g. an internal futex deadlock seen when many run in a worker pool) +# blocks its caller forever and hangs the whole shard. On timeout, subprocess +# kills the child and raises TimeoutExpired, which callers treat as a failure. +# Configurable via FFMPEG_TIMEOUT_S: the final multi-speaker `amix` of long +# sessions can exceed the default when all node CPUs are saturated (mix reruns +# use a higher value and lower concurrency). +def _ffmpeg_timeout_s() -> int: + raw = os.environ.get("FFMPEG_TIMEOUT_S", "").strip() + if raw: + try: + value = int(float(raw)) + if value > 0: + return value + except ValueError: + pass + return 600 + + +FFMPEG_TIMEOUT_S = _ffmpeg_timeout_s() + + +def ffmpeg_executable() -> str: + """Return ffmpeg binary path; honors FFMPEG_BIN for cluster static builds.""" + env_bin = os.environ.get("FFMPEG_BIN", "").strip() + if env_bin: + return env_bin + return shutil.which("ffmpeg") or "ffmpeg" + + +T = TypeVar("T") +R = TypeVar("R") + + +class PipelineError(Exception): + """Raised when a pipeline stage cannot continue.""" + + +def mfa_models_root() -> Path: + return Path(os.environ.get("MFA_ROOT_DIR", MFA_ROOT_DIR_DEFAULT)).expanduser().resolve() + + +def resolve_mfa_dict(mfa_dict: str) -> Path: + for candidate in ( + mfa_models_root() / "pretrained_models" / "dictionary" / f"{mfa_dict}.dict", + mfa_models_root() / "pretrained_models" / "dictionary" / f"{mfa_dict}.txt", + Path(mfa_dict).expanduser(), + ): + if candidate.is_file(): + return candidate.resolve() + msg = f"MFA dictionary not found for {mfa_dict!r}" + raise FileNotFoundError(msg) + + +def resolve_mfa_acoustic_model(mfa_acoustic: str) -> str: + """Resolve an acoustic model name to its pretrained zip path. + + Passing the full zip path (instead of the bare model name) lets each MFA + invocation use its own ``--temporary_directory`` without breaking model + lookup, which is required for safe parallel execution. + """ + direct = Path(mfa_acoustic).expanduser() + if direct.is_file() or direct.is_dir(): + return str(direct.resolve()) + for candidate in ( + mfa_models_root() / "pretrained_models" / "acoustic" / f"{mfa_acoustic}.zip", + mfa_models_root() / "pretrained_models" / "acoustic" / mfa_acoustic, + ): + if candidate.is_file() or candidate.is_dir(): + return str(candidate.resolve()) + return mfa_acoustic + + +def partition_list(items: list[T], num_parts: int) -> list[list[T]]: + if num_parts <= 1 or not items: + return [items] + parts: list[list[T]] = [[] for _ in range(num_parts)] + for i, item in enumerate(items): + parts[i % num_parts].append(item) + return [part for part in parts if part] + + +def append_mfa_g2p_args(cmd: list[str], *, g2p_path: str | Path | None) -> None: + if g2p_path: + cmd.extend(["--g2p_model_path", str(g2p_path)]) + + +def mfa_subprocess_env( + *, + temp_root: Path, + mfa_root: Path, +) -> dict[str, str]: + """Build env for ``mfa`` subprocesses. + + Inside the pyxis container, ``PYTHONPATH`` includes ``/opt/venv/...`` which + shadows the packed conda ``pynini``/OpenFST with a different wheel build. + Mixed imports (conda ``montreal_forced_aligner`` + container ``pynini``) + cause ``FstIOError: Read failed`` on G2P ``model.fst``. + """ + env = os.environ.copy() + env["TMPDIR"] = str(temp_root.parent) + env["MFA_ROOT_DIR"] = str(mfa_root) + + conda_lib: Path | None = None + mfa_env_dir = os.environ.get("MFA_ENV", "").strip() + if mfa_env_dir: + candidate = Path(mfa_env_dir) / "lib" + if candidate.is_dir(): + conda_lib = candidate + if conda_lib is None: + mfa_bin = shutil.which("mfa") + if mfa_bin: + candidate = Path(mfa_bin).resolve().parent.parent / "lib" + if candidate.is_dir(): + conda_lib = candidate + if conda_lib is not None: + prev = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = f"{conda_lib}:{prev}" if prev else str(conda_lib) + + pp = env.get("PYTHONPATH", "") + if pp: + kept = [ + p + for p in pp.split(os.pathsep) + if p and "/opt/venv" not in p and "/opt/Export-Deploy" not in p + ] + if kept: + env["PYTHONPATH"] = os.pathsep.join(kept) + else: + env.pop("PYTHONPATH", None) + return env + + +def _extract_g2p_archive(g2p_src: Path, extract_root: Path) -> Path: + import zipfile + + extract_root.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(g2p_src) as archive: + for member in archive.infolist(): + relative = Path(member.filename) + if relative.is_absolute() or ".." in relative.parts: + msg = f"unsafe path in G2P archive: {member.filename}" + raise PipelineError(msg) + if member.is_dir(): + continue + destination = extract_root / relative + destination.parent.mkdir(parents=True, exist_ok=True) + with archive.open(member) as source, destination.open("wb") as target: + shutil.copyfileobj(source, target) + fst_paths = list(extract_root.rglob("model.fst")) + if len(fst_paths) != 1 or fst_paths[0].stat().st_size == 0: + msg = f"invalid G2P archive contents: {g2p_src}" + raise PipelineError(msg) + return fst_paths[0].parent + + +def _worker_g2p_arg(models_dir: Path, mfa_g2p: str) -> str | None: + g2p_src = resolve_mfa_g2p_model(mfa_g2p) + if g2p_src.is_dir(): + local_dir = models_dir / "g2p" / g2p_src.name + if not (local_dir / "model.fst").is_file(): + local_dir.parent.mkdir(parents=True, exist_ok=True) + if local_dir.exists(): + shutil.rmtree(local_dir) + shutil.copytree(g2p_src, local_dir) + return str(local_dir) + if g2p_src.is_file(): + if g2p_src.suffix == ".zip": + return str(_extract_g2p_archive(g2p_src, models_dir / "g2p")) + local = models_dir / g2p_src.name + if not local.is_file(): + shutil.copy2(g2p_src, local) + return str(local) + return None + + +def setup_mfa_worker_root( + worker_dir: Path, + *, + mfa_dict: Path, + mfa_acoustic: str, + mfa_g2p: str | None = None, + source_mfa_root: Path | None = None, +) -> tuple[Path, Path, str, str | None]: + """Prepare an isolated MFA root with local copies of lexicon and acoustic model. + + Returns (mfa_root, local_dict_path, acoustic_model_arg, g2p_model_arg) for ``mfa align``. + """ + worker_dir = worker_dir.resolve() + if worker_dir.exists(): + shutil.rmtree(worker_dir, ignore_errors=True) + mfa_root = worker_dir / "mfa_root" + models_dir = worker_dir / "models" + + models_dir.mkdir(parents=True, exist_ok=True) + mfa_root.mkdir(parents=True, exist_ok=True) + + local_dict = models_dir / mfa_dict.name + shutil.copy2(mfa_dict, local_dict) + + acoustic_src = Path(resolve_mfa_acoustic_model(mfa_acoustic)) + if acoustic_src.is_file() and acoustic_src.suffix == ".zip": + local_zip = models_dir / acoustic_src.name + shutil.copy2(acoustic_src, local_zip) + acoustic_arg = str(local_zip) + _extract_acoustic_zip(local_zip, mfa_root, source_mfa_root=source_mfa_root) + elif acoustic_src.is_dir(): + local_acoustic = models_dir / "acoustic" / acoustic_src.name + local_acoustic.parent.mkdir(parents=True, exist_ok=True) + shutil.copytree(acoustic_src, local_acoustic) + acoustic_arg = str(local_acoustic) + else: + acoustic_arg = resolve_mfa_acoustic_model(mfa_acoustic) + + g2p_arg = _worker_g2p_arg(models_dir, mfa_g2p) if mfa_g2p else None + + global_config = mfa_root / "global_config.yaml" + global_config.write_text( + "\n".join( + [ + "auto_server: true", + "blas_num_threads: 1", + "clean: false", + "cleanup_textgrids: true", + "database_limited_mode: false", + "debug: false", + "num_jobs: 3", + "overwrite: false", + "quiet: false", + "seed: 0", + "single_speaker: false", + f"temporary_directory: {mfa_root}", + "use_mp: true", + "use_postgres: false", + "use_threading: true", + "verbose: false", + ] + ) + + "\n", + encoding="utf-8", + ) + + cmd_hist = mfa_root / "command_history.yaml" + if cmd_hist.exists() or cmd_hist.is_symlink(): + cmd_hist.unlink(missing_ok=True) + cmd_hist.symlink_to("/dev/null") + + return mfa_root, local_dict, acoustic_arg, g2p_arg + + +def _extract_acoustic_zip( + zip_path: Path, + mfa_root: Path, + *, + source_mfa_root: Path | None = None, +) -> None: + import zipfile + + extracted_root = mfa_root / "extracted_models" / "acoustic" + if source_mfa_root is not None: + src_acoustic = source_mfa_root / "extracted_models" / "acoustic" + if src_acoustic.is_dir(): + for src_dir in src_acoustic.iterdir(): + if not src_dir.is_dir(): + continue + dst_dir = extracted_root / src_dir.name + if (dst_dir / "final.mdl").is_file(): + continue + shutil.copytree(src_dir, dst_dir) + return + + extracted_root.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(zip_path) as zf: + zf.extractall(extracted_root) + + for _path in extracted_root.rglob("final.mdl"): + return + msg = f"acoustic zip did not contain final.mdl: {zip_path}" + raise PipelineError(msg) + + +def resolve_mfa_g2p_model(mfa_g2p: str) -> Path: + direct = Path(mfa_g2p).expanduser() + if direct.is_file() or direct.is_dir(): + return direct.resolve() + root = mfa_models_root() + for candidate in ( + root / "extracted_models" / "g2p" / f"{mfa_g2p}_g2p", + root / "extracted_models" / "g2p" / mfa_g2p, + root / "pretrained_models" / "g2p" / mfa_g2p, + root / "pretrained_models" / "g2p" / f"{mfa_g2p}.zip", + ): + if candidate.is_file() or candidate.is_dir(): + return candidate.resolve() + msg = ( + f"MFA G2P model not found for {mfa_g2p!r} under " + f"{root / 'pretrained_models' / 'g2p'} or {root / 'extracted_models' / 'g2p'}" + ) + raise FileNotFoundError( + msg + ) + + +def log_exception(context: str, exc: BaseException) -> None: + logger.error("%s: %s", context, exc) + logger.debug(traceback.format_exc()) + + +def load_jsonl(path: Path) -> list[dict]: + rows: list[dict] = [] + try: + with path.open(encoding="utf-8") as f: + for line_no, raw_line in enumerate(f, start=1): + line = raw_line.strip() + if not line: + continue + try: + rows.append(json.loads(line)) + except json.JSONDecodeError as exc: + msg = f"{path}:{line_no}: invalid JSON: {exc}" + raise ValueError(msg) from exc + except OSError as exc: + msg = f"cannot read {path}: {exc}" + raise PipelineError(msg) from exc + return rows + + +def write_jsonl(path: Path, rows: list[dict]) -> None: + try: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + for row in rows: + f.write(json.dumps(row, ensure_ascii=False) + "\n") + except OSError as exc: + msg = f"cannot write {path}: {exc}" + raise PipelineError(msg) from exc + + +def append_jsonl(path: Path, row: dict, *, lock: threading.Lock | None = None) -> None: + def _write() -> None: + try: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as f: + f.write(json.dumps(row, ensure_ascii=False) + "\n") + except OSError as exc: + msg = f"cannot append to {path}: {exc}" + raise PipelineError(msg) from exc + + if lock is not None: + with lock: + _write() + else: + _write() + + +def thread_temp_root(base: Path) -> Path: + """Per-thread scratch directory under *base* for parallel workers.""" + root = base / f"thread_{threading.get_ident()}" + root.mkdir(parents=True, exist_ok=True) + return root + + +def run_thread_pool( + items: list[T], + fn: Callable[[T], R], + *, + workers: int = 1, +) -> list[R]: + if workers <= 1 or len(items) <= 1: + return [fn(item) for item in items] + with ThreadPoolExecutor(max_workers=workers) as pool: + return list(pool.map(fn, items)) + + +def run_process_pool( + items: list[T], + fn: Callable[[T], R], + *, + workers: int = 1, +) -> list[R]: + if workers <= 1 or len(items) <= 1: + return [fn(item) for item in items] + with ProcessPoolExecutor(max_workers=workers) as pool: + return list(pool.map(fn, items)) + + +def ffprobe_duration(path: Path) -> float: + cmd = [ + "ffprobe", + "-v", + "error", + "-show_entries", + "format=duration", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(path), + ] + try: + result = subprocess.run(cmd, capture_output=True, text=True, check=False) + except OSError as exc: + msg = f"ffprobe not available for {path}: {exc}" + raise RuntimeError(msg) from exc + if result.returncode != 0 or not result.stdout.strip(): + msg = f"ffprobe failed for {path}: {result.stderr[-300:]}" + raise RuntimeError(msg) + try: + return float(result.stdout.strip()) + except ValueError as exc: + msg = f"ffprobe returned non-numeric duration for {path}" + raise RuntimeError(msg) from exc + + +def extract_segment_wav( + src: Path, + dst: Path, + start: float, + end: float, + *, + padding: float = 0.0, + max_duration: float | None = None, +) -> float | None: + """Extract audio for MFA, optionally padded before/after the manifest interval. + + Returns the absolute extract start time in *src*, or None on failure. + The clip spans [max(0, start-padding), min(max_duration, end+padding)] when + *max_duration* is set, otherwise [max(0, start-padding), end+padding]. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + extract_start = max(0.0, start - padding) + extract_end = end + padding + if max_duration is not None: + extract_end = min(max_duration, extract_end) + duration = max(extract_end - extract_start, 0.01) + cmd = [ + ffmpeg_executable(), + "-nostdin", + "-y", + "-ss", + f"{extract_start:.6f}", + "-i", + str(src), + "-t", + f"{duration:.6f}", + "-ac", + "1", + "-acodec", + "pcm_s16le", + str(dst), + ] + try: + result = subprocess.run( + cmd, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + timeout=FFMPEG_TIMEOUT_S, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + logger.warning("ffmpeg segment extract failed to start for %s: %s", src, exc) + return None + if result.returncode != 0: + logger.warning( + "ffmpeg segment extract failed (%s [%.3f-%.3f] pad=%.3f): %s", + src, + extract_start, + extract_end, + padding, + result.stderr[-300:], + ) + return None + return extract_start + + +def map_segment_words_to_recording( + words: list[tuple[float, float, str]], + *, + extract_start: float, + extract_end: float, +) -> list[tuple[float, float, str]]: + """Map MFA word times from a padded segment clip back to recording time. + + All MFA word intervals are kept as aligned (only bounded by the clip MFA ran on). + """ + mapped: list[tuple[float, float, str]] = [] + for start, end, word in words: + abs_start = start + extract_start + abs_end = end + extract_start + if abs_end <= extract_start or abs_start >= extract_end: + continue + mapped.append((abs_start, abs_end, word)) + return mapped + + +def parse_textgrid_words(tg_path: Path) -> list[tuple[float, float, str]]: + import textgrid + + try: + tg = textgrid.TextGrid.fromFile(str(tg_path)) + tier = tg.getFirst("words") + except Exception as exc: + msg = f"failed to parse TextGrid {tg_path}: {exc}" + raise ValueError(msg) from exc + + words: list[tuple[float, float, str]] = [] + for iv in tier.intervals: + mark = (iv.mark or "").strip() + if mark and mark not in SILENCE_TOKENS: + words.append((iv.minTime, iv.maxTime, mark)) + return words + + +def safe_parse_textgrid_words(tg_path: Path) -> list[tuple[float, float, str]]: + try: + return parse_textgrid_words(tg_path) + except ImportError as exc: + msg = ( + "textgrid package is required to parse MFA TextGrids " + f"(pip install textgrid): {exc}" + ) + raise PipelineError( + msg + ) from exc + except Exception as exc: + log_exception(f"TextGrid parse failed for {tg_path}", exc) + return [] + + +def write_textgrid( + words: list[tuple[float, float, str]], + output_path: Path, + *, + xmin: float = 0.0, + xmax: float | None = None, +) -> None: + if xmax is None: + xmax = words[-1][1] + 0.01 if words else xmin + 0.01 + + intervals: list[tuple[float, float, str]] = [] + prev_end = xmin + for start, end, word in sorted(words, key=lambda x: x[0]): + if start > prev_end + 0.001: + intervals.append((prev_end, start, "")) + intervals.append((start, end, word)) + prev_end = end + if prev_end < xmax: + intervals.append((prev_end, xmax, "")) + + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", encoding="utf-8") as f: + f.write('File type = "ooTextFile"\n') + f.write('Object class = "TextGrid"\n\n') + f.write(f"xmin = {xmin}\n") + f.write(f"xmax = {xmax}\n") + f.write("tiers? \n") + f.write("size = 1\n") + f.write("item []:\n") + f.write(" item [1]:\n") + f.write(' class = "IntervalTier"\n') + f.write(' name = "words"\n') + f.write(f" xmin = {xmin}\n") + f.write(f" xmax = {xmax}\n") + f.write(f" intervals: size = {len(intervals)}\n") + for i, (s, e, text) in enumerate(intervals, 1): + safe = text.replace('"', '""') + f.write(f" intervals [{i}]:\n") + f.write(f" xmin = {s}\n") + f.write(f" xmax = {e}\n") + f.write(f' text = "{safe}"\n') + + +def merge_speech_intervals( + intervals: list[tuple[float, float]], + merge_gap: float, +) -> list[tuple[float, float]]: + merged: list[tuple[float, float]] = [] + for start, end in sorted(intervals): + if end <= start: + continue + if merged and (start - merged[-1][1]) <= merge_gap: + merged[-1] = (merged[-1][0], max(merged[-1][1], end)) + else: + merged.append((start, end)) + return merged + + +def merge_tagged_speech_intervals( + intervals: list[tuple[float, float, str]], + merge_gap: float, +) -> list[tuple[float, float, str]]: + merged: list[tuple[float, float, str]] = [] + for start, end, label in sorted(intervals): + if end <= start: + continue + if merged and (start - merged[-1][1]) <= merge_gap: + prev_start, prev_end, prev_label = merged[-1] + merged[-1] = (prev_start, max(prev_end, end), prev_label) + else: + merged.append((start, end, label)) + return merged + + +def textgrid_to_rttm_lines( + tg_path: Path, + *, + speaker: str, + merge_gap: float = 0.2, +) -> list[str]: + file_id = tg_path.stem + try: + intervals = [ + (start, end) + for start, end, _ in parse_textgrid_words(tg_path) + ] + except Exception as exc: + log_exception(f"RTTM conversion failed for {tg_path}", exc) + return [] + + merged = merge_speech_intervals(intervals, merge_gap) + + lines = [] + for start, end in merged: + dur = end - start + lines.append( + f"SPEAKER {file_id} 1 {start:.6f} {dur:.6f} {speaker} " + ) + return lines + + +def discover_sessions(data_root: Path) -> list[Path]: + sessions: list[Path] = [] + for path in data_root.iterdir(): + if path.is_symlink(): + # Cluster data_links are already filtered by the link stage. Avoid + # following every symlink to Lustre before stage 0 can start writing. + sessions.append(path) + elif path.is_dir() and (path / "machine_generated_transcript.json").is_file(): + sessions.append(path) + return sorted(sessions) + + +def load_speaker_count_tsv(path: Path) -> dict[str, int]: + """Load `` `` lines from a speaker-count TSV.""" + counts: dict[str, int] = {} + if not path.is_file(): + return counts + for raw_line in path.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + parts = line.split() + if len(parts) < 2: + continue + try: + counts[parts[1]] = int(parts[0]) + except ValueError: + continue + return counts + + +def load_session_id_list(path: Path) -> list[str]: + """Load one session id per line (comments and blanks ignored).""" + if not path.is_file(): + return [] + ids: list[str] = [] + for raw_line in path.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + ids.append(line.split()[0]) + return ids + + +def filter_sessions_by_ids(sessions: list[Path], session_ids: list[str]) -> list[Path]: + wanted = set(session_ids) + if not wanted: + return sessions + by_name = {session.name: session for session in sessions} + missing = sorted(wanted - set(by_name)) + if missing: + logger.warning("sessions-file: %d id(s) not found under data root (first: %s)", len(missing), missing[0]) + return [by_name[sid] for sid in session_ids if sid in by_name] + + +def order_sessions_by_speaker_priority( + sessions: list[Path], + speaker_counts: dict[str, int], + *, + min_priority_speakers: int, +) -> list[Path]: + """Put sessions with at least *min_priority_speakers* first (higher counts earlier).""" + if min_priority_speakers <= 1 or not speaker_counts: + return sessions + + def sort_key(session_dir: Path) -> tuple[int, int, str]: + count = speaker_counts.get(session_dir.name, 0) + priority_bucket = 0 if count >= min_priority_speakers else 1 + return (priority_bucket, -count, session_dir.name) + + return sorted(sessions, key=sort_key) + + +def recording_id(speaker_id: str, session_id: str) -> str: + return f"{speaker_id}_{session_id}_postprocessed" + + +def masked_speaker_audio_path(audio_masked_dir: Path, speaker_id: str, session_id: str) -> Path: + return audio_masked_dir / f"{recording_id(speaker_id, session_id)}.wav" + + +def masked_speaker_rttm_path(audio_masked_dir: Path, speaker_id: str, session_id: str) -> Path: + return audio_masked_dir / f"{recording_id(speaker_id, session_id)}.rttm" + + +def recording_textgrid_path( + textgrid_dir: Path, + recording_id: str, + *, + variant: str = "ordinary", +) -> Path: + suffix = {"ordinary": "", "fastmss": "_fastmss", "fb": "_fb"}.get(variant, "") + return textgrid_dir / f"{recording_id}{suffix}.TextGrid" + + +def recording_textgrid_paths(textgrid_dir: Path, recording_id: str) -> list[Path]: + ordinary = recording_textgrid_path(textgrid_dir, recording_id, variant="ordinary") + fb_path = recording_textgrid_path(textgrid_dir, recording_id, variant="fb") + if ordinary.is_file() and fb_path.is_file(): + return [ordinary, fb_path] + if ordinary.is_file(): + return [ordinary] + if fb_path.is_file(): + return [fb_path] + return [ordinary] + + +def fastmss_textgrid_path(textgrid_dir: Path, recording_id: str) -> Path: + return recording_textgrid_path(textgrid_dir, recording_id, variant="fastmss") + + +def session_textgrid_path( + textgrid_dir: Path, + session_id: str, + *, + variant: str = "ordinary", +) -> Path: + suffix = {"ordinary": "", "fastmss": "_fastmss"}.get(variant, "") + return textgrid_dir / f"{session_id}{suffix}.TextGrid" + + +def interval_overlaps(start: float, end: float, intervals: list[tuple[float, float]]) -> bool: + return any(start < interval_end and end > interval_start for interval_start, interval_end in intervals) + + +def speech_intervals_from_textgrid(tg_path: Path) -> list[tuple[float, float]]: + return [(start, end) for start, end, _ in parse_textgrid_words(tg_path)] + + +def fb_intervals_for_recording(textgrid_dir: Path, recording_id: str) -> list[tuple[float, float]]: + fb_path = recording_textgrid_path(textgrid_dir, recording_id, variant="fb") + if fb_path.is_file(): + return speech_intervals_from_textgrid(fb_path) + ordinary = recording_textgrid_path(textgrid_dir, recording_id, variant="ordinary") + fastmss = recording_textgrid_path(textgrid_dir, recording_id, variant="fastmss") + if ordinary.is_file() and fastmss.is_file(): + ordinary_words = parse_textgrid_words(ordinary) + fastmss_words = parse_textgrid_words(fastmss) + if len(ordinary_words) > len(fastmss_words): + fastmss_intervals = [(s, e) for s, e, _ in fastmss_words] + return [ + (start, end) + for start, end, word in ordinary_words + if word == "speech" and not interval_overlaps(start, end, fastmss_intervals) + ] + return [] + + +def alignment_items_from_textgrid(tg_path: Path) -> list: + from lhotse.supervision import AlignmentItem + + try: + words = parse_textgrid_words(tg_path) + except Exception as exc: + log_exception(f"alignment extraction failed for {tg_path}", exc) + return [] + + items = [] + for start, end, word in words: + items.append( + AlignmentItem( + symbol=word, + start=round(start, 6), + duration=round(end - start, 6), + ) + ) + return items + + +def alignment_items_for_lhotse( + main_tg_path: Path, + *, + fb_intervals: list[tuple[float, float]] | None = None, +) -> list: + if fb_intervals is None: + fb_intervals = fb_intervals_for_recording(main_tg_path.parent, main_tg_path.stem) + + from lhotse.supervision import AlignmentItem + + try: + words = parse_textgrid_words(main_tg_path) + except Exception as exc: + log_exception(f"alignment extraction failed for {main_tg_path}", exc) + return [] + + items = [] + for start, end, word in words: + if fb_intervals and interval_overlaps(start, end, fb_intervals): + continue + items.append( + AlignmentItem( + symbol=word, + start=round(start, 6), + duration=round(end - start, 6), + ) + ) + return items + + +def write_lhotse_concatenated_textgrid( + main_tg_path: Path, + output_path: Path, + *, + fb_intervals: list[tuple[float, float]] | None = None, + xmax: float | None = None, +) -> None: + if fb_intervals is None: + fb_intervals = fb_intervals_for_recording(main_tg_path.parent, main_tg_path.stem) + + words = parse_textgrid_words(main_tg_path) + if fb_intervals: + words = [ + (start, end, word) + for start, end, word in words + if not interval_overlaps(start, end, fb_intervals) + ] + write_textgrid(words, output_path, xmin=0.0, xmax=xmax) + + +def write_rttm(path: Path, lines: list[str]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8") + + +def speech_rttm_line( + file_id: str, + speaker: str, + start: float, + end: float, + *, + label: str = "speech", +) -> str: + """RTTM speech interval; *label* is written to the subtype field.""" + dur = max(end - start, 0.0) + return ( + f"SPEAKER {file_id} 1 {start:.6f} {dur:.6f} {label} {speaker} " + ) + + +def build_speech_rttm_lines( + file_id: str, + speaker: str, + intervals: list[tuple[float, float]], + *, + label: str = "speech", + merge_gap: float = 0.0, +) -> list[str]: + merged: list[tuple[float, float]] = [] + for start, end in sorted(intervals): + if end <= start: + continue + if merged and (start - merged[-1][1]) <= merge_gap: + merged[-1] = (merged[-1][0], max(merged[-1][1], end)) + else: + merged.append((start, end)) + return [ + speech_rttm_line(file_id, speaker, start, end, label=label) + for start, end in merged + ] + + +def normalization_log_entry( + row: dict, + *, + text_raw: str | None = None, + text_norm: str, + num2words_lang: str, + error: str = "", +) -> dict: + text = text_raw if text_raw is not None else (row.get("text_raw") or row.get("text") or "") + return { + "session_id": row.get("session_id"), + "speaker_id": row.get("speaker_id"), + "recording_id": row.get("recording_id"), + "segment_index": row.get("segment_index"), + "start": float(row["start"]), + "end": float(row["end"]), + "duration": float(row.get("duration", float(row["end"]) - float(row["start"]))), + "text": text, + "text_norm": text_norm, + "changed": text != text_norm, + "num2words_lang": num2words_lang, + "error": error, + } + + +def segment_fallback_log_entry( + seg: dict, + recording_id: str, + *, + reason: str, + detail: str = "", +) -> dict: + return { + "recording_id": recording_id, + "session_id": seg.get("session_id"), + "speaker_id": seg.get("speaker_id"), + "segment_index": seg.get("segment_index"), + "start": float(seg["start"]), + "end": float(seg["end"]), + "duration": float(seg.get("duration", float(seg["end"]) - float(seg["start"]))), + "text_norm": seg.get("text_norm", ""), + "reason": reason, + "detail": detail, + "fallback": "manifest_boundaries", + "rttm_label": "speech", + } + + +def merge_session_rttm( + rttm_paths: list[Path], + session_id: str, + output_path: Path, +) -> int: + parsed: list[tuple[float, str]] = [] + for path in rttm_paths: + if not path.is_file(): + logger.warning("%s: missing per-recording RTTM %s", session_id, path) + continue + try: + content = path.read_text(encoding="utf-8") + except OSError as exc: + log_exception(f"cannot read RTTM {path}", exc) + continue + for raw_line in content.splitlines(): + line = raw_line.strip() + if not line or line.startswith(";"): + continue + parts = line.split() + if len(parts) < 8 or parts[0] != "SPEAKER": + continue + try: + start = float(parts[3]) + except ValueError: + logger.warning("%s: invalid RTTM start time in %s: %s", session_id, path.name, line) + continue + parts[1] = session_id + parsed.append((start, " ".join(parts))) + + parsed.sort(key=lambda x: x[0]) + lines = [line for _, line in parsed] + write_rttm(output_path, lines) + return len(lines) + + +def load_norm_manifest_rows( + manifests_dir: Path, + *, + sessions: list[str] | None = None, +) -> tuple[list[dict], int]: + rows: list[dict] = [] + manifest_errors = 0 + wanted = set(sessions) if sessions else None + for path in sorted(manifests_dir.glob("*_norm.jsonl")): + if path.name == "all_norm.jsonl": + continue + try: + file_rows = load_jsonl(path) + except Exception as exc: + manifest_errors += 1 + log_exception(f"cannot load manifest {path}", exc) + continue + if wanted is not None: + file_rows = [r for r in file_rows if r.get("session_id") in wanted] + rows.extend(file_rows) + return rows, manifest_errors + + +def group_segments_by_recording(rows: list[dict]) -> dict[str, list[dict]]: + from collections import defaultdict + + grouped: dict[str, list[dict]] = defaultdict(list) + for row in rows: + grouped[row["recording_id"]].append(row) + for segments in grouped.values(): + segments.sort(key=lambda r: (r["start"], r["segment_index"])) + return grouped + + +def group_segments_by_session(rows: list[dict]) -> dict[str, list[dict]]: + from collections import defaultdict + + grouped: dict[str, list[dict]] = defaultdict(list) + for row in rows: + grouped[row["session_id"]].append(row) + for segments in grouped.values(): + segments.sort(key=lambda r: (r["start"], r["segment_index"])) + return grouped + + +def group_recordings_by_session(rows: list[dict]) -> dict[str, list[dict]]: + from collections import defaultdict + + grouped: dict[str, list[dict]] = defaultdict(list) + seen: dict[str, set[str]] = defaultdict(set) + for row in rows: + session_id = row["session_id"] + rec_id = row["recording_id"] + if rec_id in seen[session_id]: + continue + seen[session_id].add(rec_id) + grouped[session_id].append( + { + "recording_id": rec_id, + "speaker_id": row["speaker_id"], + "audio_path": Path(row["audio_filepath_16k"]), + } + ) + for recordings in grouped.values(): + recordings.sort(key=lambda e: e["recording_id"]) + return grouped + + +def load_fallback_intervals(fallback_log: Path, recording_id: str) -> list[tuple[float, float]]: + if not fallback_log.is_file(): + return [] + intervals: list[tuple[float, float]] = [] + try: + for row in load_jsonl(fallback_log): + if row.get("recording_id") != recording_id: + continue + intervals.append((float(row["start"]), float(row["end"]))) + except Exception as exc: + log_exception(f"cannot read fallback log for {recording_id}", exc) + return intervals + + +def alignment_items_from_words(words: list[tuple[float, float, str]]) -> list: + from lhotse.supervision import AlignmentItem + + return [ + AlignmentItem( + symbol=word, + start=round(start, 6), + duration=round(end - start, 6), + ) + for start, end, word in words + ] + + +def words_to_json(words: list[tuple[float, float, str]]) -> list[list]: + return [[start, end, word] for start, end, word in words] + + +def words_from_json(raw: list) -> list[tuple[float, float, str]]: + return [(float(s), float(e), str(w)) for s, e, w in raw] + + +def tagged_words_to_json(words: list[tuple[float, float, str, str]]) -> list[list]: + return [[start, end, word, speaker_id] for start, end, word, speaker_id in words] + + +def tagged_words_from_json(raw: list) -> list[tuple[float, float, str, str]]: + out: list[tuple[float, float, str, str]] = [] + for item in raw: + if len(item) == 4: + s, e, w, spk = item + out.append((float(s), float(e), str(w), str(spk))) + else: + s, e, w = item + out.append((float(s), float(e), str(w), "")) + return out + + +def alignment_record( + recording_id: str, + segments: list[dict], + *, + merged_words: list[tuple[float, float, str]], + fb_words: list[tuple[float, float, str]], + audio_duration: float, +) -> dict: + return { + "recording_id": recording_id, + "speaker_id": segments[0]["speaker_id"], + "session_id": segments[0]["session_id"], + "audio_filepath_16k": segments[0]["audio_filepath_16k"], + "audio_duration": audio_duration, + "merged_words": words_to_json(merged_words), + "fb_words": words_to_json(fb_words), + } + + +def session_alignment_record( + session_id: str, + *, + merged_words: list[tuple[float, float, str, str]], + fb_words: list[tuple[float, float, str, str]], + audio_duration: float, + recordings: list[dict], +) -> dict: + return { + "session_id": session_id, + "audio_duration": audio_duration, + "merged_words": tagged_words_to_json(merged_words), + "fb_words": tagged_words_to_json(fb_words), + "recordings": recordings, + } + + +def append_alignment_record( + path: Path, + record: dict, + *, + lock: threading.Lock | None = None, +) -> None: + append_jsonl(path, record, lock=lock) + + +def load_alignment_ids(path: Path) -> set[str]: + if not path.is_file(): + return set() + ids: set[str] = set() + for row in load_jsonl(path): + if "session_id" in row: + ids.add(row["session_id"]) + elif "recording_id" in row: + ids.add(row["recording_id"]) + return ids + + +def load_alignments_by_session(path: Path) -> dict[str, dict]: + by_id: dict[str, dict] = {} + if not path.is_file(): + return by_id + for row in load_jsonl(path): + if "session_id" in row: + by_id[row["session_id"]] = row + return by_id + + +def load_alignments_by_recording(path: Path) -> dict[str, dict]: + by_id: dict[str, dict] = {} + if not path.is_file(): + return by_id + for row in load_jsonl(path): + if "recordings" in row: + for rec in row["recordings"]: + by_id[rec["recording_id"]] = rec + elif "recording_id" in row: + by_id[row["recording_id"]] = row + return by_id + + +def build_rttm_lines_from_words( + recording_id: str, + speaker_id: str, + merged_words: list[tuple[float, float, str]], + fb_words: list[tuple[float, float, str]], + *, + merge_gap: float = 0.2, +) -> list[str]: + tagged: list[tuple[float, float, str]] = [] + for start, end, _ in merged_words: + tagged.append((start, end, "")) + for start, end, _ in fb_words: + tagged.append((start, end, "speech")) + merged = merge_tagged_speech_intervals(tagged, merge_gap) + return [ + speech_rttm_line(recording_id, speaker_id, start, end, label=label) + for start, end, label in merged + ] + + +def build_session_rttm_lines_from_words( + session_id: str, + merged_words: list[tuple[float, float, str, str]], + fb_words: list[tuple[float, float, str, str]], + *, + merge_gap: float = 0.2, +) -> list[str]: + from collections import defaultdict + + by_speaker: dict[str, list[tuple[float, float, str]]] = defaultdict(list) + for start, end, _, speaker_id in merged_words: + by_speaker[speaker_id].append((start, end, "")) + for start, end, _, speaker_id in fb_words: + by_speaker[speaker_id].append((start, end, "speech")) + + lines: list[tuple[float, str]] = [] + for speaker_id, tagged in by_speaker.items(): + merged = merge_tagged_speech_intervals(tagged, merge_gap) + for start, end, label in merged: + line = speech_rttm_line(session_id, speaker_id, start, end, label=label) + lines.append((start, line)) + lines.sort(key=lambda x: x[0]) + return [line for _, line in lines] + + +def merge_session_rttm_from_line_lists( + session_id: str, + line_lists: list[list[str]], +) -> list[str]: + parsed: list[tuple[float, str]] = [] + for lines in line_lists: + for raw_line in lines: + line = raw_line.strip() + if not line or line.startswith(";"): + continue + parts = line.split() + if len(parts) < 8 or parts[0] != "SPEAKER": + continue + try: + start = float(parts[3]) + except ValueError: + continue + parts[1] = session_id + parsed.append((start, " ".join(parts))) + parsed.sort(key=lambda x: x[0]) + return [line for _, line in parsed] + + +def build_recording_rttm_lines( + recording_id: str, + speaker_id: str, + tg_path: Path, + *, + fallback_log: Path | None = None, + merge_gap: float = 0.2, +) -> list[str]: + textgrid_dir = tg_path.parent + tagged: list[tuple[float, float, str]] = [] + for path in recording_textgrid_paths(textgrid_dir, recording_id): + try: + for start, end, word in parse_textgrid_words(path): + label = "speech" if word == "speech" else "" + tagged.append((start, end, label)) + except Exception as exc: + log_exception(f"RTTM conversion failed for {path}", exc) + + fb_path = recording_textgrid_path(textgrid_dir, recording_id, variant="fb") + if fallback_log is not None and not fb_path.is_file(): + for start, end in load_fallback_intervals(fallback_log, recording_id): + tagged.append((start, end, "speech")) + + merged = merge_tagged_speech_intervals(tagged, merge_gap) + return [ + speech_rttm_line(recording_id, speaker_id, start, end, label=label) + for start, end, label in merged + ] + + +def pad_speech_intervals( + speech_intervals: list[tuple[float, float]], + pad: float, + duration: float, +) -> list[tuple[float, float]]: + """Grow each speech interval by *pad* seconds on both sides, clamped to [0, duration]. + + Overlaps created by padding are merged. Used to keep a margin of untouched + audio around speech boundaries so pause noise never abuts speech. + """ + if pad <= 0: + return merge_speech_intervals(speech_intervals, 0.0) + padded = [ + (max(0.0, start - pad), min(duration, end + pad)) + for start, end in speech_intervals + ] + return merge_speech_intervals(padded, 0.0) + + +def invert_intervals( + speech_intervals: list[tuple[float, float]], + duration: float, +) -> list[tuple[float, float]]: + """Return gaps between *speech_intervals* over [0, *duration*] (pause regions).""" + pauses: list[tuple[float, float]] = [] + cursor = 0.0 + for start, end in sorted(speech_intervals): + if start > cursor + 1e-6: + pauses.append((cursor, start)) + cursor = max(cursor, end) + if cursor < duration - 1e-6: + pauses.append((cursor, duration)) + return pauses + + +def parse_rttm_speech_intervals( + lines: list[str], + *, + merge_gap: float = 0.2, +) -> list[tuple[float, float]]: + """Speech intervals from RTTM lines ( and speech subtype labels).""" + raw: list[tuple[float, float]] = [] + for raw_line in lines: + line = raw_line.strip() + if not line or line.startswith(";"): + continue + parts = line.split() + if len(parts) < 8 or parts[0] != "SPEAKER": + continue + label = parts[6] + if label not in {"", "speech"}: + continue + try: + start = float(parts[3]) + dur = float(parts[4]) + except ValueError: + continue + if dur > 0: + raw.append((start, start + dur)) + return merge_speech_intervals(raw, merge_gap) + + +def parse_session_rttm_by_speaker( + lines: list[str], + *, + merge_gap: float = 0.2, +) -> dict[str, list[tuple[float, float]]]: + """Speech intervals per speaker from a session-level RTTM (stage 4 output).""" + from collections import defaultdict + + by_speaker: dict[str, list[tuple[float, float]]] = defaultdict(list) + for raw_line in lines: + line = raw_line.strip() + if not line or line.startswith(";"): + continue + parts = line.split() + if len(parts) < 8 or parts[0] != "SPEAKER": + continue + label = parts[6] + if label not in {"", "speech"}: + continue + speaker_id = parts[7] + try: + start = float(parts[3]) + dur = float(parts[4]) + except ValueError: + continue + if dur > 0: + by_speaker[speaker_id].append((start, start + dur)) + + return { + speaker_id: merge_speech_intervals(intervals, merge_gap) + for speaker_id, intervals in by_speaker.items() + } + + +def session_rttm_path(audio_mixed_dir: Path, session_id: str) -> Path: + return audio_mixed_dir / f"{session_id}.rttm" + + +def load_session_rttm_by_speaker( + rttm_path: Path, + *, + merge_gap: float = 0.2, +) -> dict[str, list[tuple[float, float]]]: + if not rttm_path.is_file(): + return {} + lines = rttm_path.read_text(encoding="utf-8").splitlines() + return parse_session_rttm_by_speaker(lines, merge_gap=merge_gap) + + +def speech_intervals_from_recording_alignment( + rec_row: dict, + *, + merge_gap: float = 0.2, +) -> list[tuple[float, float]]: + merged = words_from_json(rec_row.get("merged_words", [])) + fb = words_from_json(rec_row.get("fb_words", [])) + raw = [(start, end) for start, end, _ in merged] + [(start, end) for start, end, _ in fb] + return merge_speech_intervals(raw, merge_gap) + + +def decode_audio_mono_f32(path: Path, *, target_sr: int = 16000) -> tuple: + import numpy as np + + cmd = [ + ffmpeg_executable(), + "-nostdin", + "-i", + str(path), + "-ar", + str(target_sr), + "-ac", + "1", + "-f", + "f32le", + "pipe:1", + ] + try: + result = subprocess.run( + cmd, + stdin=subprocess.DEVNULL, + capture_output=True, + check=False, + timeout=FFMPEG_TIMEOUT_S, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + logger.warning("ffmpeg decode failed to start for %s: %s", path, exc) + raise + if result.returncode != 0: + msg = result.stderr.decode(errors="replace")[-400:] + msg_0 = f"ffmpeg decode failed for {path}: {msg}" + raise RuntimeError(msg_0) + audio = np.frombuffer(result.stdout, dtype=np.float32) + return audio, target_sr + + +def encode_audio_mono_f32_to_wav( + audio, + dst: Path, + *, + sample_rate: int = 16000, +) -> bool: + dst.parent.mkdir(parents=True, exist_ok=True) + cmd = [ + ffmpeg_executable(), + "-y", + "-f", + "f32le", + "-ar", + str(sample_rate), + "-ac", + "1", + "-i", + "pipe:0", + "-c:a", + "pcm_s16le", + str(dst), + ] + try: + result = subprocess.run( + cmd, + input=audio.tobytes(), + capture_output=True, + check=False, + timeout=FFMPEG_TIMEOUT_S, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + logger.warning("ffmpeg WAV encode failed to start for %s: %s", dst, exc) + return False + if result.returncode != 0: + logger.warning( + "ffmpeg WAV encode failed for %s: %s", + dst, + result.stderr.decode(errors="replace")[-400:], + ) + return False + return True + + +def apply_white_noise_in_pause_intervals( + src: Path, + dst: Path, + pause_intervals: list[tuple[float, float]], + *, + target_sr: int = 16000, + noise_level: float = 0.0002, + seed: int | None = None, + preserve_speech: bool = True, + stitch_ms: float = 5.0, +) -> bool: + """Replace *pause_intervals* with white noise; RTTM speech samples stay intact. + + Speech regions are never modified. Only samples inside each pause interval are + written. When *preserve_speech* is True and *stitch_ms* > 0, a linear + crossfade is applied **inside** the pause at both ends: original pause audio + is blended into white noise at the pause start and blended back out before the + pause end (still within the pause interval). + + When *preserve_speech* is False, the pause interior is filled with noise + with no crossfade. + """ + import numpy as np + + try: + audio, sr = decode_audio_mono_f32(src, target_sr=target_sr) + except RuntimeError as exc: + log_exception(f"decode failed for {src}", exc) + return False + + audio = np.array(audio, dtype=np.float32, copy=True) + n = len(audio) + rng = np.random.default_rng(seed) + stitch = max(0, round(stitch_ms * sr / 1000.0)) if preserve_speech else 0 + + for start, end in pause_intervals: + i0 = max(0, int(start * sr)) + i1 = min(n, int(end * sr)) + length = i1 - i0 + if length <= 0: + continue + + orig_pause = audio[i0:i1].copy() + noise = rng.standard_normal(length, dtype=np.float32) * noise_level + + if stitch > 0: + fade = min(stitch, length // 2) + if fade > 0: + ramp_in = np.linspace(0.0, 1.0, fade, dtype=np.float32) + ramp_out = np.linspace(1.0, 0.0, fade, dtype=np.float32) + noise[:fade] = orig_pause[:fade] * (1.0 - ramp_in) + noise[:fade] * ramp_in + noise[-fade:] = orig_pause[-fade:] * ramp_out + noise[-fade:] * (1.0 - ramp_out) + + audio[i0:i1] = noise + + return encode_audio_mono_f32_to_wav(audio, dst, sample_rate=sr) + + +def prepare_speaker_audio_for_session_mix( + audio_path: Path, + dst: Path, + *, + speech_intervals: list[tuple[float, float]], + audio_duration: float | None = None, + noise_level: float = 0.0002, + seed: int | None = None, + preserve_speech: bool = True, + stitch_ms: float = 5.0, + boundary_indent: float = 0.5, +) -> bool: + """Fill non-speech (pause) regions with white noise before session mixing. + + *boundary_indent* keeps that many seconds of original audio on each side of a + speech interval untouched (pause noise starts 0.5s after speech ends and stops + 0.5s before speech begins by default). + """ + if audio_duration is None: + try: + audio_duration = ffprobe_duration(audio_path) + except RuntimeError: + audio_duration = max((end for _, end in speech_intervals), default=0.0) + 0.01 + + padded_speech = pad_speech_intervals(speech_intervals, boundary_indent, audio_duration) + pause_intervals = invert_intervals(padded_speech, audio_duration) + return apply_white_noise_in_pause_intervals( + audio_path, + dst, + pause_intervals, + noise_level=noise_level, + seed=seed, + preserve_speech=preserve_speech, + stitch_ms=stitch_ms, + ) + + +def session_mixed_audio_path(audio_mixed_dir: Path, session_id: str) -> Path: + return audio_mixed_dir / f"{session_id}.wav" + + +def mix_audio_files(audio_paths: list[Path], output_path: Path) -> bool: + existing = [p for p in audio_paths if p.is_file()] + if not existing: + return False + output_path.parent.mkdir(parents=True, exist_ok=True) + if len(existing) == 1: + import shutil + + try: + shutil.copy2(existing[0], output_path) + return True + except OSError as exc: + log_exception(f"cannot copy mixed audio to {output_path}", exc) + return False + + cmd = [ffmpeg_executable(), "-nostdin", "-y"] + for path in existing: + cmd.extend(["-i", str(path)]) + n = len(existing) + cmd.extend( + [ + "-filter_complex", + f"amix=inputs={n}:duration=longest:dropout_transition=0", + "-ac", + "1", + "-ar", + "16000", + "-acodec", + "pcm_s16le", + str(output_path), + ] + ) + try: + result = subprocess.run( + cmd, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + timeout=FFMPEG_TIMEOUT_S, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + logger.warning("ffmpeg mix failed to start for %s: %s", output_path.name, exc) + return False + if result.returncode != 0: + logger.warning("ffmpeg mix failed for %s: %s", output_path.name, result.stderr[-400:]) + return False + return True + + +def run_main(main_fn) -> None: + """Entry-point wrapper: log tracebacks and return non-zero exit codes.""" + try: + raise SystemExit(main_fn()) + except PipelineError as exc: + logger.exception("Pipeline failed") + raise SystemExit(1) from exc + except KeyboardInterrupt: + logger.exception("Interrupted") + raise SystemExit(130) from None + except Exception as exc: + logger.exception("Unhandled error") + raise SystemExit(1) from exc diff --git a/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_manifest.py b/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_manifest.py new file mode 100644 index 0000000000..8dbc5f3b95 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_manifest.py @@ -0,0 +1,255 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Build normalized MFA rows directly from one raw session, without persisted manifests.""" + +from __future__ import annotations + +import json +import logging +import re +import unicodedata +from typing import TYPE_CHECKING + +from david_ai_common import ( + PipelineError, + log_exception, + recording_id, +) + +if TYPE_CHECKING: + from pathlib import Path + +logger = logging.getLogger(__name__) + +_DIGIT_DECADE_RE = re.compile(r"(? str: + """Insert spaces around punctuation, keeping apostrophes and hyphens word-internal.""" + text = re.sub(r"([^\w\s'-])", r" \1 ", text) + return re.sub(r"\s+", " ", text).strip() + + +def strip_digit_grouping_commas(text: str) -> str: + return _DIGIT_GROUP_COMMA_RE.sub("", text) + + +def _split_verbalized_num2words(spoken: str) -> list[str]: + spoken = _NUM2WORDS_PUNCT_RE.sub(" ", spoken) + return [word for word in spoken.split() if word] + + +def verbalize_digit_string(num_str: str, *, num2words_lang: str) -> str: + from num2words import num2words + + spoken = num2words(int(num_str), lang=num2words_lang) + return " ".join(_split_verbalized_num2words(spoken.casefold())) + + +def verbalize_decade(num_str: str, *, num2words_lang: str) -> str: + n = int(num_str) + if num2words_lang == "en": + if 0 < n < 100 and n % 10 == 0: + base = verbalize_digit_string(num_str, num2words_lang=num2words_lang) + return f"{base[:-1]}ies" if base.endswith("y") else f"{base}s" + if 1000 <= n <= 2090 and n % 10 == 0: + head = n // 100 + decade = n % 100 + if decade: + return ( + f"{verbalize_digit_string(str(head), num2words_lang=num2words_lang)} " + f"{verbalize_decade(str(decade), num2words_lang=num2words_lang)}" + ) + return f"{verbalize_digit_string(num_str, num2words_lang=num2words_lang)} s" + + +def preprocess_spoken_numbers(text: str, *, num2words_lang: str) -> str: + lang = (num2words_lang or "").strip() + if not lang: + return text + + def _decade(match: re.Match[str]) -> str: + return f" {verbalize_decade(match.group(1), num2words_lang=lang)} " + + def _feet_inches(match: re.Match[str]) -> str: + feet = verbalize_digit_string(match.group(1), num2words_lang=lang) + inches = verbalize_digit_string(match.group(2), num2words_lang=lang) + return f"{feet} {inches}" + + def _hyphen_prefix(match: re.Match[str]) -> str: + return f"{verbalize_digit_string(match.group(1), num2words_lang=lang)}{match.group(2)}" + + def _general(match: re.Match[str]) -> str: + return f" {verbalize_digit_string(match.group(0), num2words_lang=lang)} " + + text = _DIGIT_DECADE_RE.sub(_decade, text) + text = _DIGIT_FEET_INCHES_RE.sub(_feet_inches, text) + text = _DIGIT_HYPHEN_PREFIX_RE.sub(_hyphen_prefix, text) + text = _DIGIT_GENERAL_RE.sub(_general, text) + return re.sub(r"\s+", " ", text).strip() + + +def normalize_text(text: str, *, num2words_lang: str = "en") -> str: + """Normalize English transcript text using only tutorial-local helpers.""" + lang = (num2words_lang or "").strip() + try: + prepared = separate_gluing_punctuation(strip_digit_grouping_commas(text)) + prepared = preprocess_spoken_numbers(prepared, num2words_lang=lang) if lang else prepared + prepared = "".join( + character + for character in unicodedata.normalize("NFC", prepared).translate(_SMART_QUOTE_TRANSLATION) + if not unicodedata.category(character).startswith("C") + ).casefold() + rebuilt: list[str] = [] + allowed = _ENGLISH_ALPHABET | _PERMITTED_SYMBOLS + for character in prepared: + category = unicodedata.category(character) + if character in allowed: + rebuilt.append(character) + elif character.isspace() or category.startswith(("Z", "P", "S")): + rebuilt.append(" ") + else: + rebuilt.append(character) + + normalized: list[str] = [] + for token in "".join(rebuilt).split(): + if all(character in allowed for character in token): + normalized.append(token) + continue + folded = "".join( + character + for character in unicodedata.normalize("NFKD", token) + if unicodedata.category(character) != "Mn" + ) + normalized.append(folded if folded and all(character in allowed for character in folded) else "spn") + return " ".join(normalized) + except Exception as exc: + msg = f"normalization failed for text snippet: {text[:80]!r}" + raise ValueError(msg) from exc + + +def resolve_speaker_audio_path(session_dir: Path, speaker_id: str) -> Path: + """Resolve one speaker WAV using the supported filename priority.""" + candidates = ( + session_dir / f"{speaker_id}_postprocess.wav", + session_dir / f"{speaker_id}_postprocessed.wav", + session_dir / f"{speaker_id}.wav", + session_dir / f"{speaker_id}_preprocessed.wav", + ) + for candidate in candidates: + if candidate.is_file(): + return candidate.resolve() + expected = ", ".join(path.name for path in candidates) + msg = f"no speaker audio for {speaker_id}; tried: {expected}" + raise FileNotFoundError(msg) + + +def build_session_rows( + session_dir: Path, + *, + num2words_lang: str = "en", +) -> list[dict]: + """Read one raw session and create normalized rows entirely in memory.""" + session_id = session_dir.name + transcript_path = session_dir / "machine_generated_transcript.json" + if not transcript_path.is_file(): + msg = f"missing transcript: {transcript_path}" + raise FileNotFoundError(msg) + try: + with transcript_path.open(encoding="utf-8") as stream: + payload = json.load(stream) + except json.JSONDecodeError as exc: + msg = f"invalid JSON in {transcript_path}: {exc}" + raise ValueError(msg) from exc + except OSError as exc: + msg = f"cannot read {transcript_path}: {exc}" + raise PipelineError(msg) from exc + + segments = payload.get("transcript") if isinstance(payload, dict) else None + if not isinstance(segments, list): + msg = f"expected transcript list in {transcript_path}" + raise TypeError(msg) + + speaker_ids = { + str(segment["speaker"]) + for segment in segments + if isinstance(segment, dict) and segment.get("speaker") + } + norm_rows: list[dict] = [] + for speaker_id in sorted(speaker_ids): + audio_path = resolve_speaker_audio_path(session_dir, speaker_id) + rec_id = recording_id(speaker_id, session_id) + speaker_segments = [ + segment + for segment in segments + if isinstance(segment, dict) and segment.get("speaker") == speaker_id + ] + + for index, segment in enumerate(speaker_segments): + text_raw = (segment.get("text") or "").strip() + try: + start = float(segment["start"]) + end = float(segment["end"]) + except (KeyError, TypeError, ValueError) as exc: + logger.warning("%s/%s segment %d: invalid boundaries: %s", session_id, speaker_id, index, exc) + continue + if end <= start: + continue + + text_norm = "" + try: + text_norm = normalize_text(text_raw, num2words_lang=num2words_lang) if text_raw else "" + except Exception as exc: + log_exception(f"{session_id}/{speaker_id} segment {index} normalization", exc) + + row = { + "session_id": session_id, + "speaker_id": speaker_id, + "recording_id": rec_id, + "segment_index": index, + "start": start, + "end": end, + "duration": round(end - start, 6), + "text": text_norm, + "text_raw": text_raw, + "text_norm": text_norm, + "audio_filepath": str(audio_path.resolve()), + "audio_filepath_16k": str(audio_path.resolve()), + } + norm_rows.append(row) + return norm_rows diff --git a/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_mfa_align.py b/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_mfa_align.py new file mode 100644 index 0000000000..cd5f6245c7 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_mfa_align.py @@ -0,0 +1,568 @@ +"""Ephemeral MFA alignment used only by the on-the-fly RAM E2E pipeline.""" + +from __future__ import annotations + +import logging +import os +import shutil +import subprocess +import tempfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING + +from david_ai_common import ( + PipelineError, + append_jsonl, + append_mfa_g2p_args, + extract_segment_wav, + ffprobe_duration, + group_segments_by_recording, + log_exception, + map_segment_words_to_recording, + mfa_subprocess_env, + resolve_mfa_acoustic_model, + resolve_mfa_g2p_model, + run_thread_pool, + safe_parse_textgrid_words, + segment_fallback_log_entry, + session_textgrid_path, + words_to_json, + write_textgrid, +) + +if TYPE_CHECKING: + import threading + +logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO) +logger = logging.getLogger(__name__) + +# Guard against a wedged `mfa align` (e.g. SQLite/pynini lock contention seen when +# many run in a worker pool) blocking its caller forever and hanging the whole shard, +# mirroring the FFMPEG_TIMEOUT_S protection already used for ffmpeg subprocesses. +MFA_ALIGN_TIMEOUT_S = 3600 + + +def _segment_extract_workers(num_segments: int) -> int: + raw = os.environ.get("SEG_EXTRACT_WORKERS", "").strip() + if raw: + try: + return max(1, min(int(raw), num_segments)) + except ValueError: + pass + return max(1, min(num_segments, 8)) + + +def _export_segment_for_mfa( + seg: dict, + *, + audio_path: Path, + corpus_dir: Path, + segment_padding: float, + audio_duration: float, +) -> tuple[dict, Path, float] | None: + seg_idx = int(seg["segment_index"]) + seg_wav = corpus_dir / f"seg_{seg_idx:05d}.wav" + seg_txt = corpus_dir / f"seg_{seg_idx:05d}.txt" + seg_start = float(seg["start"]) + seg_end = float(seg["end"]) + extract_start = extract_segment_wav( + audio_path, + seg_wav, + seg_start, + seg_end, + padding=segment_padding, + max_duration=audio_duration, + ) + if extract_start is None: + return None + seg_txt.write_text(seg["text_norm"].strip(), encoding="utf-8") + return seg, seg_wav, extract_start + + +@dataclass +class RecordingAlignResult: + ok: bool + mfa_segments: int = 0 + fallback_segments: int = 0 + fallback_entries: list[dict] = field(default_factory=list) + merged_words: list[tuple[float, float, str]] = field(default_factory=list) + fb_words: list[tuple[float, float, str]] = field(default_factory=list) + audio_duration: float = 0.0 + + +@dataclass +class SessionAlignResult: + ok: bool + mfa_segments: int = 0 + fallback_segments: int = 0 + merged_words: list[tuple[float, float, str, str]] = field(default_factory=list) + fb_words: list[tuple[float, float, str, str]] = field(default_factory=list) + audio_duration: float = 0.0 + recordings: list[dict] = field(default_factory=list) + + +def _apply_segment_fallback( + seg: dict, + recording_id: str, + *, + reason: str, + detail: str, + fb_words: list[tuple[float, float, str]], + fallback_entries: list[dict], +) -> None: + start = float(seg["start"]) + end = float(seg["end"]) + fb_words.append((start, end, "speech")) + fallback_entries.append( + segment_fallback_log_entry( + seg, + recording_id, + reason=reason, + detail=detail, + ) + ) + logger.warning( + "%s segment %s: MFA failed (%s); using manifest speech [%.3f, %.3f]", + recording_id, + seg.get("segment_index"), + reason, + start, + end, + ) + + +def _segment_miss_or_fallback( + seg: dict, + recording_id: str, + *, + reason: str, + detail: str, + fb_words: list[tuple[float, float, str]], + fallback_entries: list[dict], + use_fallback: bool, +) -> None: + if use_fallback: + _apply_segment_fallback( + seg, + recording_id, + reason=reason, + detail=detail, + fb_words=fb_words, + fallback_entries=fallback_entries, + ) + + +def align_recording( + recording_id: str, + segments: list[dict], + *, + mfa_dict: Path, + mfa_acoustic: str, + temp_parent: Path, + num_jobs: int, + fallback_log: Path, + segment_padding: float, + fallback_log_lock: threading.Lock | None = None, + worker_mfa_root: Path | None = None, + worker_acoustic: str | None = None, + worker_g2p: str | None = None, + mfa_g2p: str | None = None, + keep_temp: bool = False, + use_fallback: bool = True, +) -> RecordingAlignResult: + temp_parent.mkdir(parents=True, exist_ok=True) + try: + if keep_temp: + temp_root = temp_parent / f"align_{recording_id}" + if temp_root.exists(): + shutil.rmtree(temp_root) + temp_root.mkdir(parents=True, exist_ok=True) + return _align_recording_impl( + recording_id, + segments, + mfa_dict=mfa_dict, + mfa_acoustic=mfa_acoustic, + temp_root=temp_root, + num_jobs=num_jobs, + fallback_log=fallback_log, + segment_padding=segment_padding, + fallback_log_lock=fallback_log_lock, + worker_mfa_root=worker_mfa_root, + worker_acoustic=worker_acoustic, + worker_g2p=worker_g2p, + mfa_g2p=mfa_g2p, + cleanup_temp=False, + use_fallback=use_fallback, + ) + with tempfile.TemporaryDirectory( + prefix=f"mfa_{recording_id}_", + dir=temp_parent, + ) as td: + return _align_recording_impl( + recording_id, + segments, + mfa_dict=mfa_dict, + mfa_acoustic=mfa_acoustic, + temp_root=Path(td), + num_jobs=num_jobs, + fallback_log=fallback_log, + segment_padding=segment_padding, + fallback_log_lock=fallback_log_lock, + worker_mfa_root=worker_mfa_root, + worker_acoustic=worker_acoustic, + worker_g2p=worker_g2p, + mfa_g2p=mfa_g2p, + cleanup_temp=False, + use_fallback=use_fallback, + ) + except Exception as exc: + log_exception(f"MFA alignment failed for {recording_id}", exc) + return RecordingAlignResult(ok=False) + + +def _align_recording_impl( + recording_id: str, + segments: list[dict], + *, + mfa_dict: Path, + mfa_acoustic: str, + temp_root: Path, + num_jobs: int, + fallback_log: Path, + segment_padding: float, + fallback_log_lock: threading.Lock | None = None, + worker_mfa_root: Path | None = None, + worker_acoustic: str | None = None, + worker_g2p: str | None = None, + mfa_g2p: str | None = None, + cleanup_temp: bool = True, + use_fallback: bool = True, +) -> RecordingAlignResult: + audio_path = Path(segments[0]["audio_filepath_16k"]) + speaker_id = segments[0]["speaker_id"] + if not audio_path.is_file(): + logger.warning("%s: missing 16k audio %s", recording_id, audio_path) + return RecordingAlignResult(ok=False) + + usable = [s for s in segments if (s.get("text_norm") or "").strip()] + if not usable: + logger.warning("%s: no segments with normalized text", recording_id) + return RecordingAlignResult(ok=False) + + try: + audio_duration = ffprobe_duration(audio_path) + except RuntimeError: + audio_duration = max(float(s["end"]) for s in usable) + 0.05 + + corpus_name = f"corpus_{recording_id}" + corpus_dir = temp_root / corpus_name / speaker_id + aligned_dir = temp_root / "aligned" + corpus_dir.mkdir(parents=True, exist_ok=True) + aligned_dir.mkdir(parents=True, exist_ok=True) + + merged_words: list[tuple[float, float, str]] = [] + fb_words: list[tuple[float, float, str]] = [] + fallback_entries: list[dict] = [] + mfa_segments = 0 + + seg_meta: list[tuple[dict, Path, float]] = [] + + def _export_one(seg: dict) -> tuple[dict, Path, float] | None: + try: + return _export_segment_for_mfa( + seg, + audio_path=audio_path, + corpus_dir=corpus_dir, + segment_padding=segment_padding, + audio_duration=audio_duration, + ) + except OSError as exc: + log_exception(f"{recording_id} segment {seg.get('segment_index')} export", exc) + return None + + extract_results = run_thread_pool( + usable, + _export_one, + workers=_segment_extract_workers(len(usable)), + ) + for seg, exported in zip(usable, extract_results, strict=True): + if exported is None: + _segment_miss_or_fallback( + seg, + recording_id, + reason="segment_export_failed", + detail="ffmpeg extract failed", + fb_words=fb_words, + fallback_entries=fallback_entries, + use_fallback=use_fallback, + ) + continue + seg_meta.append(exported) + + mfa_failed_globally = False + if seg_meta: + mfa_root = worker_mfa_root or (temp_root / "mfa_root") + if worker_mfa_root is None: + mfa_root.mkdir(parents=True, exist_ok=True) + acoustic_arg = worker_acoustic or resolve_mfa_acoustic_model(mfa_acoustic) + align_cmd = [ + "mfa", + "align", + str(corpus_dir.parent), + str(mfa_dict), + acoustic_arg, + str(aligned_dir), + ] + align_cmd.append("--clean" if worker_mfa_root is None else "--no_clean") + align_cmd.extend( + [ + "--use_mp", + "-j", + str(num_jobs), + "--beam", + "100", + "--retry_beam", + "400", + "--output_format", + "long_textgrid", + "--uses_speaker_adaptation", + "false", + "-t", + str(mfa_root), + ] + ) + g2p_arg = worker_g2p + if g2p_arg is None and mfa_g2p: + try: + g2p_arg = str(resolve_mfa_g2p_model(mfa_g2p)) + except FileNotFoundError: + logger.warning("%s: MFA G2P model not found for %r", recording_id, mfa_g2p) + append_mfa_g2p_args(align_cmd, g2p_path=g2p_arg) + logger.info("%s: running MFA on %d segments", recording_id, len(seg_meta)) + mfa_env = mfa_subprocess_env(temp_root=temp_root, mfa_root=mfa_root) + try: + result = subprocess.run( + align_cmd, + capture_output=True, + text=True, + env=mfa_env, + timeout=MFA_ALIGN_TIMEOUT_S, + check=False, + ) + except subprocess.TimeoutExpired as exc: + logger.exception("%s: mfa align timed out after %ds", recording_id, MFA_ALIGN_TIMEOUT_S) + mfa_failed_globally = True + detail = f"mfa align timed out after {MFA_ALIGN_TIMEOUT_S}s: {exc}" + except OSError as exc: + logger.exception("%s: mfa align failed to start", recording_id) + mfa_failed_globally = True + detail = str(exc) + else: + detail = result.stderr[-1200:] if result.returncode != 0 else "" + if result.returncode != 0: + logger.error( + "%s: mfa align failed (exit %d): %s", + recording_id, + result.returncode, + detail, + ) + mfa_failed_globally = True + + for seg, seg_wav, extract_start in seg_meta: + if mfa_failed_globally: + _segment_miss_or_fallback( + seg, + recording_id, + reason="mfa_align_failed", + detail=detail, + fb_words=fb_words, + fallback_entries=fallback_entries, + use_fallback=use_fallback, + ) + continue + + tg_path = aligned_dir / speaker_id / f"{seg_wav.stem}.TextGrid" + seg_end = float(seg["end"]) + extract_end = min(audio_duration, seg_end + segment_padding) + words = safe_parse_textgrid_words(tg_path) if tg_path.is_file() else [] + mapped_words = map_segment_words_to_recording( + words, + extract_start=extract_start, + extract_end=extract_end, + ) + if not mapped_words: + reason = "missing_textgrid" if not tg_path.is_file() else "empty_alignment" + _segment_miss_or_fallback( + seg, + recording_id, + reason=reason, + detail=tg_path.name, + fb_words=fb_words, + fallback_entries=fallback_entries, + use_fallback=use_fallback, + ) + continue + + merged_words.extend(mapped_words) + mfa_segments += 1 + + if not merged_words and not fb_words: + logger.warning("%s: no segment output produced", recording_id) + return RecordingAlignResult(ok=False) + + merged_words.sort(key=lambda x: x[0]) + fb_words.sort(key=lambda x: x[0]) + + if use_fallback: + for entry in fallback_entries: + try: + append_jsonl(fallback_log, entry, lock=fallback_log_lock) + except PipelineError as exc: + log_exception(f"cannot write MFA fallback log for {recording_id}", exc) + + if cleanup_temp and temp_root.exists(): + shutil.rmtree(temp_root, ignore_errors=True) + if worker_mfa_root is not None: + stale_db = worker_mfa_root / corpus_name + if stale_db.exists(): + shutil.rmtree(stale_db, ignore_errors=True) + + logger.info( + "%s: aligned %d MFA words, %d fallback segments", + recording_id, + len(merged_words), + len(fallback_entries), + ) + return RecordingAlignResult( + ok=True, + mfa_segments=mfa_segments, + fallback_segments=len(fallback_entries), + fallback_entries=fallback_entries, + merged_words=merged_words, + fb_words=fb_words, + audio_duration=audio_duration, + ) + + +def align_session( + session_id: str, + segments: list[dict], + *, + mfa_dict: Path, + mfa_acoustic: str, + textgrid_dir: Path, + temp_parent: Path, + num_jobs: int, + fallback_log: Path, + segment_padding: float, + fallback_log_lock: threading.Lock | None = None, + worker_mfa_root: Path | None = None, + worker_acoustic: str | None = None, + worker_g2p: str | None = None, + mfa_g2p: str | None = None, + keep_temp: bool = False, + use_fallback: bool = True, + write_textgrids: bool = True, +) -> SessionAlignResult: + by_recording = group_segments_by_recording(segments) + session_merged: list[tuple[float, float, str, str]] = [] + session_fb: list[tuple[float, float, str, str]] = [] + recording_rows: list[dict] = [] + session_duration = 0.0 + mfa_segments = 0 + fallback_segments = 0 + + for rec_id, rec_segments in sorted(by_recording.items()): + result = align_recording( + rec_id, + rec_segments, + mfa_dict=mfa_dict, + mfa_acoustic=mfa_acoustic, + temp_parent=temp_parent, + num_jobs=num_jobs, + fallback_log=fallback_log, + segment_padding=segment_padding, + fallback_log_lock=fallback_log_lock, + worker_mfa_root=worker_mfa_root, + worker_acoustic=worker_acoustic, + worker_g2p=worker_g2p, + mfa_g2p=mfa_g2p, + keep_temp=keep_temp, + use_fallback=use_fallback, + ) + if not result.ok: + logger.warning("%s: speaker recording %s failed", session_id, rec_id) + continue + + speaker_id = rec_segments[0]["speaker_id"] + for start, end, word in result.merged_words: + session_merged.append((start, end, word, speaker_id)) + for start, end, word in result.fb_words: + session_fb.append((start, end, word, speaker_id)) + session_duration = max(session_duration, result.audio_duration) + mfa_segments += result.mfa_segments + fallback_segments += result.fallback_segments + recording_rows.append( + { + "recording_id": rec_id, + "speaker_id": speaker_id, + "session_id": session_id, + "audio_filepath_16k": rec_segments[0]["audio_filepath_16k"], + "audio_duration": result.audio_duration, + "merged_words": words_to_json(result.merged_words), + "fb_words": words_to_json(result.fb_words), + } + ) + + if not session_merged and not session_fb: + logger.warning("%s: no session alignment output", session_id) + return SessionAlignResult(ok=False) + + session_merged.sort(key=lambda x: x[0]) + session_fb.sort(key=lambda x: x[0]) + max_seg_end = max((end for _, end, _, _ in session_merged + session_fb), default=0.0) + xmax = max(session_duration, max_seg_end) + 0.01 + + fastmss_words = [(s, e, w) for s, e, w, _ in session_merged] + ordinary_words = sorted( + [(s, e, w) for s, e, w, _ in session_merged] + [(s, e, w) for s, e, w, _ in session_fb], + key=lambda x: x[0], + ) + if write_textgrids: + write_textgrid( + fastmss_words, + session_textgrid_path(textgrid_dir, session_id, variant="fastmss"), + xmin=0.0, + xmax=xmax, + ) + write_textgrid( + ordinary_words, + session_textgrid_path(textgrid_dir, session_id, variant="ordinary"), + xmin=0.0, + xmax=xmax, + ) + logger.info( + "%s: session TextGrids (%d MFA words, %d fallback, %d speakers)", + session_id, + len(session_merged), + len(session_fb), + len(recording_rows), + ) + else: + logger.info( + "%s: MFA alignment (%d words, %d fallback, %d speakers; TextGrids skipped)", + session_id, + len(session_merged), + len(session_fb), + len(recording_rows), + ) + return SessionAlignResult( + ok=True, + mfa_segments=mfa_segments, + fallback_segments=fallback_segments, + merged_words=session_merged, + fb_words=session_fb, + audio_duration=xmax, + recordings=recording_rows, + ) diff --git a/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_ram_lhotse.py b/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_ram_lhotse.py new file mode 100644 index 0000000000..181a83319e --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_ram_lhotse.py @@ -0,0 +1,92 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TextGrid writers used by the on-the-fly RAM E2E pipeline.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from david_ai_common import ( + fastmss_textgrid_path, + recording_textgrid_path, + session_textgrid_path, + words_from_json, + write_textgrid, +) + +if TYPE_CHECKING: + from pathlib import Path + + from david_ai_mfa_align import SessionAlignResult + + +def write_recording_textgrids(align_result: SessionAlignResult, textgrid_dir: Path) -> None: + """Write ordinary and FastMSS TextGrids for every speaker recording.""" + textgrid_dir.mkdir(parents=True, exist_ok=True) + for rec_row in align_result.recordings: + rec_id = rec_row["recording_id"] + merged = words_from_json(rec_row["merged_words"]) + fallback = words_from_json(rec_row["fb_words"]) + duration = float(rec_row.get("audio_duration", 0.0)) + max_word_end = max((end for _, end, _ in merged + fallback), default=0.0) + xmax = max(duration, max_word_end) + 0.01 + + write_textgrid( + merged, + fastmss_textgrid_path(textgrid_dir, rec_id), + xmin=0.0, + xmax=xmax, + ) + write_textgrid( + sorted(merged + fallback, key=lambda word: word[0]), + recording_textgrid_path(textgrid_dir, rec_id), + xmin=0.0, + xmax=xmax, + ) + + +def write_session_textgrids(align_result: SessionAlignResult, textgrid_dir: Path) -> None: + """Write ordinary and FastMSS TextGrids for the mixed session timeline.""" + session_id = align_result.recordings[0]["session_id"] if align_result.recordings else "" + if not session_id: + return + + fastmss_words = [(start, end, word) for start, end, word, _ in align_result.merged_words] + ordinary_words = sorted( + fastmss_words + [(start, end, word) for start, end, word, _ in align_result.fb_words], + key=lambda word: word[0], + ) + max_word_end = max((end for _, end, _ in ordinary_words), default=0.0) + xmax = max(float(align_result.audio_duration), max_word_end) + 0.01 + textgrid_dir.mkdir(parents=True, exist_ok=True) + + write_textgrid( + fastmss_words, + session_textgrid_path(textgrid_dir, session_id, variant="fastmss"), + xmin=0.0, + xmax=xmax, + ) + write_textgrid( + ordinary_words, + session_textgrid_path(textgrid_dir, session_id, variant="ordinary"), + xmin=0.0, + xmax=xmax, + ) + + +def write_all_textgrids(align_result: SessionAlignResult, textgrid_dir: Path) -> None: + """Persist session and per-recording TextGrids in both required variants.""" + write_session_textgrids(align_result, textgrid_dir) + write_recording_textgrids(align_result, textgrid_dir) diff --git a/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_ram_session.py b/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_ram_session.py new file mode 100644 index 0000000000..f27d0b19ce --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/wav/david_ai_ram_session.py @@ -0,0 +1,444 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-session worker for the strict on-the-fly RAM E2E pipeline.""" + +from __future__ import annotations + +import os +import shutil +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from david_ai_common import ( + PipelineError, + build_session_rttm_lines_from_words, + fastmss_textgrid_path, + ffprobe_duration, + group_recordings_by_session, + group_segments_by_recording, + log_exception, + masked_speaker_audio_path, + masked_speaker_rttm_path, + mix_audio_files, + prepare_speaker_audio_for_session_mix, + recording_id, + recording_textgrid_path, + run_thread_pool, + session_mixed_audio_path, + session_rttm_path, + session_textgrid_path, + setup_mfa_worker_root, + write_rttm, +) +from david_ai_manifest import build_session_rows +from david_ai_mfa_align import align_session +from david_ai_ram_lhotse import write_all_textgrids + +if TYPE_CHECKING: + from pathlib import Path + +_PROCESS_MFA: dict | None = None + + +def _lazy_mfa_worker( + *, + ram_dir: Path, + mfa_dict: Path, + mfa_acoustic: str, + mfa_g2p: str, +) -> tuple[Path, Path, str, str | None, Path]: + """Initialize one ephemeral MFA model/database root per process.""" + global _PROCESS_MFA + if _PROCESS_MFA is None: + worker_dir = ram_dir / "mfa_workers" / f"worker_{os.getpid()}" + worker_dir.mkdir(parents=True, exist_ok=True) + mfa_root, local_dict, acoustic_arg, g2p_arg = setup_mfa_worker_root( + worker_dir, + mfa_dict=mfa_dict, + mfa_acoustic=mfa_acoustic, + mfa_g2p=mfa_g2p, + ) + _PROCESS_MFA = { + "mfa_root": mfa_root, + "local_dict": local_dict, + "acoustic_arg": acoustic_arg, + "g2p_arg": g2p_arg, + "temp_parent": worker_dir / "align_temp", + } + cfg = _PROCESS_MFA + return cfg["mfa_root"], cfg["local_dict"], cfg["acoustic_arg"], cfg["g2p_arg"], cfg["temp_parent"] + + +@dataclass +class SessionRamResult: + session_id: str + ok: bool + error: str = "" + + +def session_done_path(work_dir: Path, session_id: str) -> Path: + return work_dir / ".done" / "sessions" / f"{session_id}.done" + + +def is_session_done(work_dir: Path, session_id: str) -> bool: + """Return whether a previous run validated every required session output.""" + return session_done_path(work_dir, session_id).is_file() + + +def _clear_session_done(work_dir: Path, session_id: str) -> None: + session_done_path(work_dir, session_id).unlink(missing_ok=True) + + +def _nonempty(path: Path) -> bool: + return path.is_file() and path.stat().st_size > 0 + + +def _validate_session_outputs( + session_id: str, + norm_rows: list[dict], + *, + audio_masked_dir: Path, + audio_mixed_dir: Path, + textgrid_dir: Path, +) -> None: + """Require every declared deliverable before writing the session success flag.""" + required = [ + session_mixed_audio_path(audio_mixed_dir, session_id), + session_rttm_path(audio_mixed_dir, session_id), + session_textgrid_path(textgrid_dir, session_id, variant="ordinary"), + session_textgrid_path(textgrid_dir, session_id, variant="fastmss"), + ] + speaker_ids = sorted({row["speaker_id"] for row in norm_rows}) + for speaker_id in speaker_ids: + rec_id = recording_id(speaker_id, session_id) + required.extend( + [ + masked_speaker_audio_path(audio_masked_dir, speaker_id, session_id), + masked_speaker_rttm_path(audio_masked_dir, speaker_id, session_id), + recording_textgrid_path(textgrid_dir, rec_id, variant="ordinary"), + fastmss_textgrid_path(textgrid_dir, rec_id), + ] + ) + missing = [str(path) for path in required if not _nonempty(path)] + if missing: + msg = f"missing or empty session outputs: {missing}" + raise PipelineError(msg) + + +def _mark_session_done(work_dir: Path, session_id: str) -> None: + path = session_done_path(work_dir, session_id) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("ok\n", encoding="utf-8") + + +def _finalize_session_success( + session_id: str, + norm_rows: list[dict], + *, + work_dir: Path, + audio_masked_dir: Path, + audio_mixed_dir: Path, + textgrid_dir: Path, +) -> None: + _validate_session_outputs( + session_id, + norm_rows, + audio_masked_dir=audio_masked_dir, + audio_mixed_dir=audio_mixed_dir, + textgrid_dir=textgrid_dir, + ) + _mark_session_done(work_dir, session_id) + + +def _mix_prep_workers(num_speakers: int) -> int: + raw = os.environ.get("MIX_PREP_WORKERS", "").strip() + if raw: + try: + return max(1, min(int(raw), num_speakers)) + except ValueError: + pass + return max(1, num_speakers) + + +def _manifest_speech_intervals_by_recording(norm_rows: list[dict]) -> dict[str, list[tuple[float, float]]]: + """Return exact original manifest boundaries, before the 0.5-second protection offset.""" + return { + rec_id: [(float(row["start"]), float(row["end"])) for row in rows] + for rec_id, rows in group_segments_by_recording(norm_rows).items() + } + + +def _prepare_speaker_tracks_for_mix( + entries: list[dict], + *, + session_id: str, + session_scratch: Path, + audio_masked_dir: Path, + manifest_speech: dict[str, list[tuple[float, float]]], + rec_durations: dict[str, float], + noise_level: float, + stitch_ms: float, + boundary_offset: float, +) -> list[tuple[Path, Path]]: + """Create 16 kHz masked speaker WAVs and their persistent destinations.""" + specs: list[tuple[Path, Path, Path, list[tuple[float, float]], float, int, str]] = [] + for entry in entries: + rec_id = entry["recording_id"] + speaker_id = entry["speaker_id"] + src = entry["audio_path"] + if not src.is_file(): + msg = f"missing source audio {src}" + raise FileNotFoundError(msg) + + speech = manifest_speech.get(rec_id) + if not speech: + msg = f"no original manifest boundaries for {rec_id}" + raise PipelineError(msg) + duration = rec_durations.get(rec_id, 0.0) + if duration <= 0: + try: + duration = ffprobe_duration(src) + except RuntimeError: + duration = max(end for _, end in speech) + 0.01 + + local_dst = session_scratch / f"{rec_id}.wav" + persistent_dst = masked_speaker_audio_path(audio_masked_dir, speaker_id, session_id) + seed = hash((session_id, rec_id)) & 0xFFFFFFFF + specs.append((src, local_dst, persistent_dst, speech, duration, seed, rec_id)) + + def _prepare_one( + spec: tuple[Path, Path, Path, list[tuple[float, float]], float, int, str], + ) -> tuple[Path, Path]: + src, local_dst, persistent_dst, speech, duration, seed, rec_id = spec + if not prepare_speaker_audio_for_session_mix( + src, + local_dst, + speech_intervals=speech, + audio_duration=duration, + noise_level=noise_level, + seed=seed, + preserve_speech=True, + stitch_ms=stitch_ms, + boundary_indent=boundary_offset, + ): + msg = f"pause noise prep failed for {rec_id}" + raise PipelineError(msg) + return local_dst, persistent_dst + + return run_thread_pool(specs, _prepare_one, workers=_mix_prep_workers(len(specs))) + + +def _publish_audio(local_path: Path, output_path: Path) -> None: + """Publish completed local audio without exposing a partial final path.""" + output_path.parent.mkdir(parents=True, exist_ok=True) + temp_path = output_path.with_name(f".{output_path.name}.{os.getpid()}.tmp") + try: + shutil.copyfile(local_path, temp_path) + os.replace(temp_path, output_path) + finally: + if temp_path.is_file(): + temp_path.unlink() + + +def _write_masked_speaker_rttms( + session_id: str, + norm_rows: list[dict], + session_rttm_lines: list[str], + *, + audio_masked_dir: Path, +) -> None: + """Filter the session RTTM into one recording-scoped RTTM per speaker.""" + audio_masked_dir.mkdir(parents=True, exist_ok=True) + for speaker_id in sorted({row["speaker_id"] for row in norm_rows}): + rec_id = recording_id(speaker_id, session_id) + output_lines: list[str] = [] + for line in session_rttm_lines: + parts = line.split() + if len(parts) < 8 or parts[0] != "SPEAKER" or parts[7] != speaker_id: + continue + parts[1] = rec_id + output_lines.append(" ".join(parts)) + if not output_lines: + msg = f"no RTTM intervals for masked speaker {rec_id}" + raise PipelineError(msg) + masked_speaker_rttm_path(audio_masked_dir, speaker_id, session_id).write_text( + "\n".join(output_lines) + "\n", + encoding="utf-8", + ) + + +def _mix_session_from_manifest( + session_id: str, + norm_rows: list[dict], + *, + audio_masked_dir: Path, + audio_mixed_dir: Path, + session_ram: Path, + noise_level: float, + stitch_ms: float, + boundary_offset: float, + rec_durations: dict[str, float], +) -> None: + entries = group_recordings_by_session(norm_rows).get(session_id, []) + if not entries: + msg = "no speaker recordings to mix" + raise PipelineError(msg) + + session_scratch = session_ram / "mix" + if session_scratch.exists(): + shutil.rmtree(session_scratch, ignore_errors=True) + session_scratch.mkdir(parents=True, exist_ok=True) + + prepared_tracks = _prepare_speaker_tracks_for_mix( + entries, + session_id=session_id, + session_scratch=session_scratch, + audio_masked_dir=audio_masked_dir, + manifest_speech=_manifest_speech_intervals_by_recording(norm_rows), + rec_durations=rec_durations, + noise_level=noise_level, + stitch_ms=stitch_ms, + boundary_offset=boundary_offset, + ) + local_mixed = session_scratch / f"{session_id}.wav" + if not mix_audio_files([local for local, _ in prepared_tracks], local_mixed): + msg = "session mix failed" + raise PipelineError(msg) + + for local_path, persistent_path in prepared_tracks: + _publish_audio(local_path, persistent_path) + _publish_audio(local_mixed, session_mixed_audio_path(audio_mixed_dir, session_id)) + + +def process_session_ram( + session_dir: Path, + *, + work_dir: Path, + audio_masked_dir: Path, + audio_mixed_dir: Path, + textgrid_dir: Path, + mfa_dict: Path, + mfa_acoustic: str, + mfa_g2p: str, + ram_dir: Path, + num2words_lang: str = "en", + mfa_num_jobs: int = 2, + segment_padding: float = 0.5, + rttm_merge_gap: float = 0.2, + noise_level: float = 0.0002, + stitch_ms: float = 5.0, + boundary_offset: float = 0.5, +) -> SessionRamResult: + """Run every E2E step from raw transcript/WAV, without reading persisted pipeline state.""" + session_id = session_dir.name + session_ram = ram_dir / "sessions" / session_id + _clear_session_done(work_dir, session_id) + + try: + norm_rows = build_session_rows( + session_dir, + num2words_lang=num2words_lang, + ) + if not norm_rows: + return SessionRamResult(session_id=session_id, ok=False, error="no manifest rows") + + if session_ram.exists(): + shutil.rmtree(session_ram, ignore_errors=True) + session_ram.mkdir(parents=True, exist_ok=True) + audio_masked_dir.mkdir(parents=True, exist_ok=True) + textgrid_dir.mkdir(parents=True, exist_ok=True) + audio_mixed_dir.mkdir(parents=True, exist_ok=True) + + worker_mfa_root, local_dict, acoustic_arg, g2p_arg, temp_parent = _lazy_mfa_worker( + ram_dir=ram_dir, + mfa_dict=mfa_dict, + mfa_acoustic=mfa_acoustic, + mfa_g2p=mfa_g2p, + ) + temp_parent.mkdir(parents=True, exist_ok=True) + align_result = align_session( + session_id, + norm_rows, + mfa_dict=local_dict, + mfa_acoustic=mfa_acoustic, + textgrid_dir=textgrid_dir, + temp_parent=temp_parent / session_id, + num_jobs=mfa_num_jobs, + fallback_log=session_ram / "fallback.jsonl", + segment_padding=segment_padding, + worker_mfa_root=worker_mfa_root, + worker_acoustic=acoustic_arg, + worker_g2p=g2p_arg, + mfa_g2p=mfa_g2p, + keep_temp=False, + use_fallback=True, + write_textgrids=False, + ) + if not align_result.ok: + return SessionRamResult(session_id=session_id, ok=False, error="MFA alignment failed") + if align_result.mfa_segments == 0: + return SessionRamResult( + session_id=session_id, + ok=False, + error="MFA produced zero aligned segments; refusing fallback-only completion", + ) + + rttm_lines = build_session_rttm_lines_from_words( + session_id, + align_result.merged_words, + align_result.fb_words, + merge_gap=rttm_merge_gap, + ) + if not rttm_lines: + return SessionRamResult(session_id=session_id, ok=False, error="empty session RTTM") + write_rttm(session_rttm_path(audio_mixed_dir, session_id), rttm_lines) + _write_masked_speaker_rttms( + session_id, + norm_rows, + rttm_lines, + audio_masked_dir=audio_masked_dir, + ) + + write_all_textgrids(align_result, textgrid_dir) + + rec_durations = { + rec["recording_id"]: float(rec.get("audio_duration", 0.0)) + for rec in align_result.recordings + } + _mix_session_from_manifest( + session_id, + norm_rows, + audio_masked_dir=audio_masked_dir, + audio_mixed_dir=audio_mixed_dir, + session_ram=session_ram, + noise_level=noise_level, + stitch_ms=stitch_ms, + boundary_offset=boundary_offset, + rec_durations=rec_durations, + ) + _finalize_session_success( + session_id, + norm_rows, + work_dir=work_dir, + audio_masked_dir=audio_masked_dir, + audio_mixed_dir=audio_mixed_dir, + textgrid_dir=textgrid_dir, + ) + return SessionRamResult(session_id=session_id, ok=True) + except Exception as exc: + log_exception(f"RAM session pipeline failed for {session_id}", exc) + return SessionRamResult(session_id=session_id, ok=False, error=str(exc)) + finally: + shutil.rmtree(session_ram, ignore_errors=True) diff --git a/tutorials/audio/david_ai_redelivered_mfa/wav/requirements.txt b/tutorials/audio/david_ai_redelivered_mfa/wav/requirements.txt new file mode 100644 index 0000000000..3c8d7e7822 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/wav/requirements.txt @@ -0,0 +1 @@ +-r ../requirements.txt diff --git a/tutorials/audio/david_ai_redelivered_mfa/wav/run_david_ai_mfa_ram_session.sh b/tutorials/audio/david_ai_redelivered_mfa/wav/run_david_ai_mfa_ram_session.sh new file mode 100755 index 0000000000..1416b5baa3 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/wav/run_david_ai_mfa_ram_session.sh @@ -0,0 +1,106 @@ +#!/bin/bash +# Strict on-the-fly David AI mixed-WAV E2E pipeline. +# +# Every unfinished session starts from raw WAVs + machine_generated_transcript.json: +# normalize in memory -> MFA with base dictionary + runtime G2P -> RTTM -> +# manifest-mask speaker WAVs -> mixed 16 kHz PCM WAV -> ordinary and FastMSS +# TextGrids. Masked per-speaker WAVs and their RTTMs are persisted. +# No persisted manifests, shared lexicon, or partial output cache is read. A +# validated session.done flag skips completed sessions and is written only after +# every required output passes validation. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CURATOR_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)" + +DATA_ROOT="${DATA_ROOT:?Set DATA_ROOT to the raw-session directory}" +WORK_DIR="${WORK_DIR:-$SCRIPT_DIR/workdir_e2e_wav}" +AUDIO_MASKED_DIR="${AUDIO_MASKED_DIR:-$WORK_DIR/audio_16k_masked}" +TEXTGRID_DIR="${TEXTGRID_DIR:-$WORK_DIR/textgrids}" +AUDIO_MIXED_DIR="${AUDIO_MIXED_DIR:-$WORK_DIR/audio_mixed}" +LOG_DIR="${LOG_DIR:-$WORK_DIR/logs}" +RAM_DIR="${RAM_DIR:-/tmp/david_ai_ram_session_${SLURM_JOB_ID:-$$}_${SLURM_ARRAY_TASK_ID:-0}}" + +MFA_ROOT_DIR="${MFA_ROOT_DIR:-$HOME/MFA_models}" +MFA_DICT_NAME="${MFA_DICT_NAME:-english_us_arpa}" +MFA_G2P="${MFA_G2P:-english_us_arpa}" +MFA_ACOUSTIC="${MFA_ACOUSTIC:-english_us_arpa}" +MFA_NUM_JOBS="${MFA_NUM_JOBS:-2}" +WORKERS="${WORKERS:-4}" +SEGMENT_PADDING="${SEGMENT_PADDING:-0.5}" +RTTM_MERGE_GAP="${RTTM_MERGE_GAP:-0.2}" +NUM2WORDS_LANG="${NUM2WORDS_LANG:-en}" +SESSIONS_FILE="${SESSIONS_FILE:-}" +SHARD_COUNT="${SHARD_COUNT:-1}" +SHARD_INDEX="${SHARD_INDEX:-0}" + +PYTHON="${PYTHON:-}" +if [[ -z "$PYTHON" && -x "$CURATOR_ROOT/.venv/bin/python" ]]; then + PYTHON="$CURATOR_ROOT/.venv/bin/python" +fi +PYTHON="${PYTHON:-python3}" +MFA_ENV="${MFA_ENV:-$HOME/miniconda3/envs/curator_pain_1}" +if [[ -x "$MFA_ENV/bin/mfa" ]]; then + if [[ -n "${FFMPEG_BIN:-}" && -x "$FFMPEG_BIN" ]]; then + export PATH="${PATH}:$MFA_ENV/bin" + else + export PATH="$MFA_ENV/bin:$PATH" + fi +fi +export MFA_ROOT_DIR + +if [[ -n "${FFMPEG_BIN:-}" && -x "$FFMPEG_BIN" ]]; then + : +elif ! command -v ffmpeg >/dev/null 2>&1; then + echo "ERROR: ffmpeg not on PATH" >&2 + exit 1 +fi +if ! command -v mfa >/dev/null 2>&1; then + echo "ERROR: mfa not on PATH (MFA_ENV=$MFA_ENV)" >&2 + exit 1 +fi +if [[ ! -d "$DATA_ROOT" ]]; then + echo "ERROR: data root does not exist: $DATA_ROOT" >&2 + exit 1 +fi + +mkdir -p "$LOG_DIR" "$AUDIO_MASKED_DIR" "$TEXTGRID_DIR" "$AUDIO_MIXED_DIR" "$RAM_DIR" +RUN_ID="$(date +%Y%m%d_%H%M%S)_${SLURM_JOB_ID:-local}_${SLURM_ARRAY_TASK_ID:-0}" +LOG_FILE="$LOG_DIR/run_e2e_${RUN_ID}.log" +exec > >(tee -a "$LOG_FILE") 2>&1 + +CMD=( + "$PYTHON" "$SCRIPT_DIR/stage_ram_session_pipeline.py" + --data-root "$DATA_ROOT" + --work-dir "$WORK_DIR" + --audio-masked-dir "$AUDIO_MASKED_DIR" + --audio-mixed-dir "$AUDIO_MIXED_DIR" + --textgrid-dir "$TEXTGRID_DIR" + --mfa-dict-name "$MFA_DICT_NAME" + --mfa-acoustic "$MFA_ACOUSTIC" + --mfa-g2p "$MFA_G2P" + --ram-dir "$RAM_DIR" + --num2words-lang "$NUM2WORDS_LANG" + --mfa-num-jobs "$MFA_NUM_JOBS" + --segment-padding "$SEGMENT_PADDING" + --rttm-merge-gap "$RTTM_MERGE_GAP" + --noise-level 0.0002 + --stitch-ms 5 + --boundary-offset 0.5 + --workers "$WORKERS" + --shard-count "$SHARD_COUNT" + --shard-index "$SHARD_INDEX" +) +[[ -n "$SESSIONS_FILE" ]] && CMD+=(--sessions-file "$SESSIONS_FILE") + +echo "[$(date '+%Y-%m-%d %H:%M:%S')] ON-THE-FLY E2E START" +echo "DATA_ROOT=$DATA_ROOT" +echo "WORK_DIR=$WORK_DIR" +echo "WORKERS=$WORKERS MFA_NUM_JOBS=$MFA_NUM_JOBS" +echo "MFA dictionary=$MFA_DICT_NAME runtime_g2p=$MFA_G2P" +echo "Audio output=masked per-speaker + session mixed mono 16 kHz PCM WAV" +echo "Pause mask=original manifest boundaries +/-0.5s, noise=0.0002, smoothing=5ms" +echo "Command: ${CMD[*]}" +"${CMD[@]}" +echo "[$(date '+%Y-%m-%d %H:%M:%S')] ON-THE-FLY E2E DONE" diff --git a/tutorials/audio/david_ai_redelivered_mfa/wav/stage_ram_session_pipeline.py b/tutorials/audio/david_ai_redelivered_mfa/wav/stage_ram_session_pipeline.py new file mode 100755 index 0000000000..315a9d5d95 --- /dev/null +++ b/tutorials/audio/david_ai_redelivered_mfa/wav/stage_ram_session_pipeline.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Strict on-the-fly E2E: raw sessions -> MFA/G2P -> audio, RTTM, and TextGrids.""" + +from __future__ import annotations + +import argparse +import logging +import shutil +import tempfile +from concurrent.futures import ProcessPoolExecutor, as_completed +from dataclasses import dataclass +from pathlib import Path + +from david_ai_common import PipelineError, discover_sessions, resolve_mfa_dict, run_main +from david_ai_ram_session import SessionRamResult, is_session_done, process_session_ram + +logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO) +logger = logging.getLogger(__name__) + + +@dataclass +class RamSessionTask: + session_dir: str + work_dir: str + audio_masked_dir: str + audio_mixed_dir: str + textgrid_dir: str + mfa_dict: str + mfa_acoustic: str + mfa_g2p: str + ram_dir: str + num2words_lang: str + mfa_num_jobs: int + segment_padding: float + rttm_merge_gap: float + noise_level: float + stitch_ms: float + boundary_offset: float + + +def _run_session_task(task: RamSessionTask) -> SessionRamResult: + return process_session_ram( + Path(task.session_dir), + work_dir=Path(task.work_dir), + audio_masked_dir=Path(task.audio_masked_dir), + audio_mixed_dir=Path(task.audio_mixed_dir), + textgrid_dir=Path(task.textgrid_dir), + mfa_dict=Path(task.mfa_dict), + mfa_acoustic=task.mfa_acoustic, + mfa_g2p=task.mfa_g2p, + ram_dir=Path(task.ram_dir), + num2words_lang=task.num2words_lang, + mfa_num_jobs=task.mfa_num_jobs, + segment_padding=task.segment_padding, + rttm_merge_gap=task.rttm_merge_gap, + noise_level=task.noise_level, + stitch_ms=task.stitch_ms, + boundary_offset=task.boundary_offset, + ) + + +def sessions_without_done_flags(sessions: list[Path], work_dir: Path) -> list[Path]: + """Select only sessions that have not completed successfully.""" + return [session for session in sessions if not is_session_done(work_dir, session.name)] + + +def filter_sessions_from_file(sessions: list[Path], sessions_file: Path) -> list[Path]: + """Restrict discovered sessions to IDs listed one per line.""" + requested = { + line.strip() + for line in sessions_file.read_text(encoding="utf-8").splitlines() + if line.strip() and not line.lstrip().startswith("#") + } + available = {session.name for session in sessions} + missing = sorted(requested - available) + if missing: + logger.warning("%d requested session IDs were not found under DATA_ROOT", len(missing)) + return [session for session in sessions if session.name in requested] + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--data-root", type=Path, required=True) + ap.add_argument("--work-dir", type=Path, required=True) + ap.add_argument("--audio-masked-dir", type=Path, default=None) + ap.add_argument("--audio-mixed-dir", type=Path, default=None) + ap.add_argument("--textgrid-dir", type=Path, default=None) + ap.add_argument("--mfa-dict-name", default="english_us_arpa") + ap.add_argument("--mfa-acoustic", default="english_us_arpa") + ap.add_argument("--mfa-g2p", default="english_us_arpa") + ap.add_argument( + "--ram-dir", + type=Path, + default=Path(tempfile.gettempdir()) / "david_ai_ram_session", + ) + ap.add_argument("--num2words-lang", default="en") + ap.add_argument("--mfa-num-jobs", type=int, default=2) + ap.add_argument("--segment-padding", type=float, default=0.5) + ap.add_argument("--rttm-merge-gap", type=float, default=0.2) + ap.add_argument("--noise-level", type=float, default=0.0002) + ap.add_argument("--stitch-ms", type=float, default=5.0) + ap.add_argument("--boundary-offset", type=float, default=0.5) + ap.add_argument("--workers", type=int, default=4) + ap.add_argument("--sessions-file", type=Path, default=None) + ap.add_argument("--shard-count", type=int, default=1) + ap.add_argument("--shard-index", type=int, default=0) + args = ap.parse_args() + + if args.shard_count < 1: + msg = f"--shard-count must be >= 1, got {args.shard_count}" + raise PipelineError(msg) + if not 0 <= args.shard_index < args.shard_count: + msg = f"--shard-index must be in [0, {args.shard_count}), got {args.shard_index}" + raise PipelineError( + msg + ) + + work_dir = args.work_dir.resolve() + data_root = args.data_root.resolve() + audio_masked_dir = (args.audio_masked_dir or work_dir / "audio_16k_masked").resolve() + audio_mixed_dir = (args.audio_mixed_dir or work_dir / "audio_mixed").resolve() + textgrid_dir = (args.textgrid_dir or work_dir / "textgrids").resolve() + ram_dir = args.ram_dir.resolve() + for path in (audio_masked_dir, audio_mixed_dir, textgrid_dir, ram_dir): + path.mkdir(parents=True, exist_ok=True) + + mfa_dict_path = resolve_mfa_dict(args.mfa_dict_name) + logger.info( + "Using base MFA dictionary %s with runtime G2P model %s", + mfa_dict_path, + args.mfa_g2p, + ) + + sessions = discover_sessions(data_root) + if args.sessions_file is not None: + sessions_file = args.sessions_file.resolve() + if not sessions_file.is_file(): + msg = f"sessions file does not exist: {sessions_file}" + raise PipelineError(msg) + sessions = filter_sessions_from_file(sessions, sessions_file) + logger.info("Restricted run to %d sessions from %s", len(sessions), sessions_file) + if not sessions: + msg = f"No sessions under {data_root}" + raise PipelineError(msg) + if args.shard_count > 1: + total = len(sessions) + sessions = [ + session + for index, session in enumerate(sessions) + if index % args.shard_count == args.shard_index + ] + logger.info( + "Shard %d/%d: processing %d of %d raw sessions", + args.shard_index, + args.shard_count, + len(sessions), + total, + ) + + pending_sessions = sessions_without_done_flags(sessions, work_dir) + skipped_sessions = len(sessions) - len(pending_sessions) + workers = max(1, args.workers) + logger.info( + "Resumable E2E START: sessions=%d pending=%d done=%d workers=%d mfa_jobs=%d ram_dir=%s", + len(sessions), + len(pending_sessions), + skipped_sessions, + workers, + args.mfa_num_jobs, + ram_dir, + ) + if not pending_sessions: + logger.info("All %d selected sessions already have validated done flags", len(sessions)) + return 0 + tasks = [ + RamSessionTask( + session_dir=str(session.resolve()), + work_dir=str(work_dir), + audio_masked_dir=str(audio_masked_dir), + audio_mixed_dir=str(audio_mixed_dir), + textgrid_dir=str(textgrid_dir), + mfa_dict=str(mfa_dict_path), + mfa_acoustic=args.mfa_acoustic, + mfa_g2p=args.mfa_g2p, + ram_dir=str(ram_dir), + num2words_lang=args.num2words_lang, + mfa_num_jobs=args.mfa_num_jobs, + segment_padding=args.segment_padding, + rttm_merge_gap=args.rttm_merge_gap, + noise_level=args.noise_level, + stitch_ms=args.stitch_ms, + boundary_offset=args.boundary_offset, + ) + for session in pending_sessions + ] + + ok = fail = completed = 0 + with ProcessPoolExecutor(max_workers=workers) as pool: + futures = [pool.submit(_run_session_task, task) for task in tasks] + for future in as_completed(futures): + result = future.result() + completed += 1 + if result.ok: + ok += 1 + else: + fail += 1 + logger.warning("%s failed: %s", result.session_id, result.error) + if completed % 50 == 0 or completed == len(futures): + logger.info( + "E2E progress: %d/%d (ok=%d fail=%d)", + completed, + len(futures), + ok, + fail, + ) + + shutil.rmtree(ram_dir, ignore_errors=True) + logger.info( + "Resumable E2E DONE: ok=%d fail=%d previously_done=%d workers=%d", + ok, + fail, + skipped_sessions, + workers, + ) + return 0 if fail == 0 and ok == len(pending_sessions) else 1 + + +if __name__ == "__main__": + run_main(main)