diff --git a/.github/workflows/ci_on_ubuntu.yml b/.github/workflows/ci_on_ubuntu.yml index e7729a0b925b..b60e61871d68 100644 --- a/.github/workflows/ci_on_ubuntu.yml +++ b/.github/workflows/ci_on_ubuntu.yml @@ -279,6 +279,46 @@ jobs: source tools/activate_python.sh coverage erase + test_integration_espnet3: + runs-on: ubuntu-latest + needs: process_labels + if: | + github.event.pull_request.draft == false + strategy: + max-parallel: 20 + matrix: + python-version: ["3.10", "3.12"] + pytorch-version: [2.5.1, 2.9.1] + use-conda: [false] + exclude: + - python-version: "3.12" + pytorch-version: 2.5.1 + use-conda: false + steps: + - uses: actions/checkout@master + - uses: actions/cache@v5 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ matrix.pytorch-version }}-${{ hashFiles('**/pyproject.toml') }}-${{ hashFiles('**/Makefile') }} + - name: Make environment + uses: ./.github/actions/prepare-environment + with: + os-version: ubuntu + python-version: ${{ matrix.python-version }} + pytorch-version: ${{ matrix.pytorch-version }} + chainer-version: ${{ matrix.chainer-version }} + use-conda: ${{ matrix.use-conda }} + - name: test python + run: ./ci/test_integration_espnet3.sh + - uses: codecov/codecov-action@v5 + with: + flags: test_integration_espnet3 + - name: coverage erase + continue-on-error: true + run: | + source tools/activate_python.sh + coverage erase + test_import: runs-on: ubuntu-latest needs: process_labels diff --git a/ci/test_integration_espnet3.sh b/ci/test_integration_espnet3.sh new file mode 100755 index 000000000000..bba7fc3e52b7 --- /dev/null +++ b/ci/test_integration_espnet3.sh @@ -0,0 +1,54 @@ +#!/usr/bin/env bash + +set -euo pipefail + +. tools/activate_python.sh +. tools/extra_path.sh + +python="coverage run --append" +cwd=$(pwd) + +gen_dummy_coverage(){ + touch empty.py + ${python} empty.py +} + +python3 -m pip install -e '.[asr]' + +cd ./egs3/mini_an4/asr +gen_dummy_coverage +echo "==== [ESPnet3] ASR ====" +source ./path.sh +run_with_training_config() { + local training_config=$1 + local runner=$2 + local inference_config=$3 + + ln -sfn "${training_config}" conf/training.yaml + ${python} "${runner}" \ + --stages create_dataset train_tokenizer collect_stats train infer measure \ + --training_config conf/training.yaml \ + --inference_config "${inference_config}" \ + --metrics_config conf/metrics.yaml + rm -rf exp data +} + +debug_configs=( + training_asr_rnn_data_aug.yaml + training_asr_rnn.yaml + training_asr_streaming.yaml + training_asr_transformer.yaml + training_asr_transducer.yaml +) + +for training_config in "${debug_configs[@]}"; do + run_with_training_config "${training_config}" run.py conf/inference.yaml +done + +# We need seprate inference config for transducer task +run_with_training_config \ + training_transducer_asr_conformer_rnnt.yaml \ + run.py \ + conf/inference_transducer.yaml + +cd "${cwd}" diff --git a/egs3/TEMPLATE/__init__.py b/egs3/TEMPLATE/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/egs3/TEMPLATE/asr/__init__.py b/egs3/TEMPLATE/asr/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/egs3/TEMPLATE/asr/conf/inference.yaml b/egs3/TEMPLATE/asr/conf/inference.yaml new file mode 100644 index 000000000000..a7dfe10b96ba --- /dev/null +++ b/egs3/TEMPLATE/asr/conf/inference.yaml @@ -0,0 +1,100 @@ +########################################################## +# Inference SETTING # +########################################################## +# Quick usage: +# python run.py --stages infer --inference_config conf/inference.yaml +# +# Outputs: +# - `${inference_dir}//*.scp` +# - each SCP stores `utt_id value` pairs for one inference field + +########################################################## +# PATHS AND INPUTS # +########################################################## +recipe_dir: . +exp_tag: +exp_dir: ${recipe_dir}/exp/${exp_tag} +data_dir: ${recipe_dir}/data +inference_dir: ${exp_dir}/inference + + +# Test split definitions are required for inference. +# `dataset.test[*].name` becomes the output subdirectory name under +# `${inference_dir}`. +# +# Example: +# dataset: +# _target_: espnet3.components.data.data_organizer.DataOrganizer +# test: +# - name: test-clean +# dataset: +# _target_: src.dataset.YourDataset +# manifest_path: ${dataset_dir}/manifest/test-clean.tsv +# - name: test-other +# dataset: +# _target_: src.dataset.YourDataset +# manifest_path: ${dataset_dir}/manifest/test-other.tsv +dataset: + _target_: espnet3.components.data.data_organizer.DataOrganizer + test: + preprocessor: + +########################################################## +# PARALLELIZATION SETTING # +########################################################## +parallel: + env: local + n_workers: 1 + +########################################################## +# MODEL AND OUTPUT # +########################################################## +model: + +# Key name(s) selected from each dataset sample and passed to the model. +# If `dataset[i]` returns `{"speech": ..., "text": ..., "utt_id": ...}` and +# `input_key: speech`, inference uses `dataset[i]["speech"]` as model input. +# If you set `input_key: [speech, text]`, inference builds +# `{"speech": dataset[i]["speech"], "text": dataset[i]["text"]}` and calls +# `model(**...)` with that dict. +input_key: speech + +# Optional output key list. +# These keys must exist in the dict returned by `output_fn`. +# Only the fields listed here are written as SCP files under +# `${inference_dir}//`. +# If omitted, keys are inferred from the first returned sample except for +# `idx_key`. Set this explicitly when you want stable SCP filenames. +# +# output_keys: +# - hyp + +# Sample identifier field name expected in the final inference output dict. +# Inference uses `utt_id` by default, and the final output dict must contain it +# unless you explicitly override `idx_key`. +# This value becomes the first column in generated SCP files and is used to map +# each inference result back to the corresponding dataset sample. +# +#idx_key: utt_id + +# Optional batch size. +# Set to `null` if your model or `output_fn` only supports single-item decoding. +# batch_size: 1 + +# Optional recipe-provided formatter that converts raw model outputs into dicts such as: +# {"utt_id": "...", "hyp": "..."} +# or +# {"utt_id": "...", "hyp": ["nbest1", "nbest2"]} +# List outputs are expanded to `hyp0.scp`, `hyp1.scp`, ... +# If omitted, the model output itself must already be a dict (or a list of dicts +# for batched inference) and must contain `utt_id` unless `idx_key` is +# overridden. +output_fn: src.inference.build_output + +########################################################## +# PROVIDER AND RUNNER SETTING # +########################################################## +provider: + _target_: espnet3.systems.base.inference_provider.InferenceProvider +runner: + _target_: espnet3.systems.base.inference_runner.InferenceRunner diff --git a/egs3/TEMPLATE/asr/conf/metrics.yaml b/egs3/TEMPLATE/asr/conf/metrics.yaml new file mode 100644 index 000000000000..8cd057412b82 --- /dev/null +++ b/egs3/TEMPLATE/asr/conf/metrics.yaml @@ -0,0 +1,36 @@ +########################################################## +# METRIC SCORING SETTING # +########################################################## +# Quick usage: +# python run.py --stages measure --metrics_config conf/metrics.yaml +# +# Assumptions: +# - `conf/inference.yaml` has already produced `/*.scp` files. +# - Each metric reads matching SCP keys such as `ref.scp` and `hyp.scp`. +# - If `dataset.test` is omitted, test sets are discovered from `${inference_dir}` +# subdirectories. + +########################################################## +# PATHS AND INPUTS # +########################################################## +recipe_dir: . +exp_tag: +exp_dir: ${recipe_dir}/exp/${exp_tag} +inference_dir: ${exp_dir}/inference + + +########################################################## +# METRIC DEFINITION # +########################################################## +metrics: + - metric: + _target_: espnet3.systems.asr.metrics.wer.WER + ref_key: ref # reference text key -> `/ref.scp` + hyp_key: hyp # hypothesis text key -> `/hyp.scp` + clean_types: # optional TextCleaner pipeline, e.g. [tacotron] + + - metric: + _target_: espnet3.systems.asr.metrics.cer.CER + ref_key: ref # reference text key -> `/ref.scp` + hyp_key: hyp # hypothesis text key -> `/hyp.scp` + clean_types: # optional TextCleaner pipeline, e.g. [jaconv] diff --git a/egs3/TEMPLATE/asr/conf/training.yaml b/egs3/TEMPLATE/asr/conf/training.yaml new file mode 100644 index 000000000000..58df01d71b12 --- /dev/null +++ b/egs3/TEMPLATE/asr/conf/training.yaml @@ -0,0 +1,200 @@ +########################################################## +# TRAINING SETTING # +########################################################## +# Quick usage: +# python run.py --stages create_dataset --training_config conf/training.yaml +# python run.py --stages train --training_config conf/training.yaml +# +# Main stages and what this config is used for: +# - `create_dataset`: uses `create_dataset` +# - `train_tokenizer`: uses `tokenizer` +# - `collect_stats`: uses `dataset`, `dataloader`, `model`, and `parallel` +# - `train`: uses most fields in this file + +########################################################## +# SYSTEM AND PATH SETTING # +########################################################## +num_device: 1 +num_nodes: 1 + +# Optional ESPnet task class path used by `get_espnet_model(...)`. +# If unset, `model` is instantiated directly via Hydra. +task: + +# Path scaffold relative to the recipe root. +recipe_dir: . +data_dir: ${recipe_dir}/data +exp_tag: +exp_dir: ${recipe_dir}/exp/${exp_tag} +stats_dir: ${recipe_dir}/exp/stats +inference_dir: ${exp_dir}/inference +dataset_dir: /path/to/your/dataset + +########################################################## +# DATASET CREATION SETTING # +########################################################## +# Optional dataset creation hook for the `create_dataset` stage. +# The function is resolved from `create_dataset.func` and called with the +# remaining fields in this block as keyword arguments. +create_dataset: + func: src.creating_dataset.create_dataset + dataset_dir: ${dataset_dir} + +########################################################## +# DATASET DEFINITION # +########################################################## +# Dataset organizer used across collect_stats/train/infer/measure. +# Fill `train`, `valid`, and optionally `test` with named dataset entries. +# +# Example: +# dataset: +# _target_: espnet3.components.data.data_organizer.DataOrganizer +# train: +# - name: train +# dataset: +# _target_: src.dataset.YourDataset +# manifest_path: ${dataset_dir}/manifest/train.tsv +# valid: +# - name: dev +# dataset: +# _target_: src.dataset.YourDataset +# manifest_path: ${dataset_dir}/manifest/dev.tsv +# test: +# - name: test +# dataset: +# _target_: src.dataset.YourDataset +# manifest_path: ${dataset_dir}/manifest/test.tsv +# preprocessor: +# _target_: src.preprocess.build_preprocessor +dataset: + _target_: espnet3.components.data.data_organizer.DataOrganizer + train: + valid: + test: + preprocessor: + +########################################################## +# TOKENIZER SETTING # +########################################################## +# Used by the `train_tokenizer` stage in `espnet3.systems.asr.system.ASRSystem`. +# `text_builder.func` must return either: +# - a path to a text file, or +# - an iterable of training strings +tokenizer: + +########################################################## +# MODEL SETTING # +########################################################## +# Training model definition. +# If `task` is set, this block is passed to `get_espnet_model(task, model_conf)`. +# Thus, you can reuse ESPnet2-style model configs here by setting `task` and +# filling this block with the per-model config that ESPnet2 used. +# If `task` is not set, this block is instantiated directly via Hydra. +model: + +########################################################## +# OPTIMIZER & SCHEDULER SETTING # +########################################################## +optimizer: + _target_: torch.optim.Adam + lr: 0.002 + weight_decay: 0.000001 + +scheduler: + _target_: espnet2.schedulers.warmup_lr.WarmupLR + warmup_steps: 15000 + +########################################################## +# OTHER TRAINING SETTING # +########################################################## +# Checkpoint selection criteria passed to ESPnet3 callbacks. +# Each item is: +# - monitor key +# - number of best checkpoints to keep +# - direction (`min` or `max`) +best_model_criterion: + - - valid/loss + - 3 + - min + +# Optional random seed used by collect_stats/train. +seed: + +# Extra model initialization mode forwarded to ESPnet's initializer. +# Example: xavier_uniform +init: + +########################################################## +# PARALLELIZATION SETTING # +########################################################## +# Parallel backend used by collect_stats/train. +parallel: + env: local + n_workers: 1 + +########################################################## +# DATALOADER SETTING # +########################################################## +dataloader: + collate_fn: + _target_: espnet2.train.collate_fn.CommonCollateFn + int_pad_value: -1 + # Training dataloader / iterator settings. + train: + multiple_iterator: false + num_shards: 1 + iter_factory: + _target_: espnet2.iterators.sequence_iter_factory.SequenceIterFactory + shuffle: true + collate_fn: ${dataloader.collate_fn} + batches: + # Typical batch modes include `unsorted`, `sorted`, and task-specific + # iterator/batcher setups. Adjust this to match your dataset. + type: unsorted + shape_files: + - ${stats_dir}/train/feats_shape + batch_size: 4 + batch_bins: 4000000 + + # Validation dataloader / iterator settings. + valid: + multiple_iterator: false + num_shards: 1 + iter_factory: + _target_: espnet2.iterators.sequence_iter_factory.SequenceIterFactory + shuffle: false + collate_fn: ${dataloader.collate_fn} + batches: + type: ${dataloader.train.iter_factory.batches.type} + shape_files: + - ${stats_dir}/valid/feats_shape + batch_size: ${dataloader.train.iter_factory.batches.batch_size} + batch_bins: ${dataloader.train.iter_factory.batches.batch_bins} + +########################################################## +# TRAINER SETTING # +########################################################## +# These settings map to `lightning.Trainer(...)`. +trainer: + accelerator: auto + devices: ${num_device} + num_nodes: ${num_nodes} + accumulate_grad_batches: 1 + check_val_every_n_epoch: 1 + gradient_clip_val: 1.0 + log_every_n_steps: 100 + max_epochs: 100 + + logger: + - _target_: lightning.pytorch.loggers.TensorBoardLogger + save_dir: ${exp_dir}/tensorboard + name: tb_logger + + strategy: auto + +########################################################## +# FIT SETTING # +########################################################## +# Optional keyword arguments passed directly to `trainer.fit(**fit)`. +# Leave `{}` for the default flow. +fit: {} diff --git a/egs3/TEMPLATE/asr/readme.md b/egs3/TEMPLATE/asr/readme.md new file mode 100644 index 000000000000..95c9732cc405 --- /dev/null +++ b/egs3/TEMPLATE/asr/readme.md @@ -0,0 +1,19 @@ +# ESPnet3 ASR recipe + +## Quick start + +```bash +# 0) Edit configs to set paths (e.g., conf/training.yaml: dataset_dir, create_dataset.func) + +# 1) Convert LibriSpeech to Hugging Face format (run once) +python run.py --stages create_dataset --training_config conf/training.yaml + +# 2) Train with the default Branchformer configuration +python run.py --stages train --training_config conf/training.yaml + +# 3) Decode +python run.py --stages infer --inference_config conf/inference.yaml + +# 4) Score +python run.py --stages measure --metrics_config conf/metrics.yaml +``` diff --git a/egs3/TEMPLATE/asr/run.py b/egs3/TEMPLATE/asr/run.py new file mode 100644 index 000000000000..f707dd6eb593 --- /dev/null +++ b/egs3/TEMPLATE/asr/run.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +"""Generic runner template for System-based experiments.""" + +from __future__ import annotations + +import argparse +import logging +from pathlib import Path +from typing import List, Sequence + +from espnet3.utils.config_utils import load_and_merge_config +from espnet3.utils.logging_utils import configure_logging +from espnet3.utils.stages_utils import ( + parse_cli_and_stage_args, + resolve_stages, + run_stages, +) + +# Default stage list (can be extended/overridden by callers) +DEFAULT_STAGES: List[str] = [ + "create_dataset", + "train_tokenizer", + "collect_stats", + "train", + "infer", + "measure", +] + +logger = logging.getLogger(__name__) + + +def build_parser( + stages: Sequence[str], +) -> argparse.ArgumentParser: + """Build base ArgumentParser and let caller extend it.""" + parser = argparse.ArgumentParser() + + parser.add_argument( + "--stages", + choices=list(stages) + ["all"], + nargs="+", + default=["all"], + help="Which stages to run. Multiple values allowed.", + ) + parser.add_argument( + "--training_config", + default=None, + type=Path, + help=( + "Hydra config for training (passed to load_config_with_defaults). " + "Required for create_dataset/train_tokenizer/collect_stats/train stages." + ), + ) + parser.add_argument( + "--inference_config", + default=None, + type=Path, + help="Hydra config for infer stage.", + ) + parser.add_argument( + "--metrics_config", + default=None, + type=Path, + help="Hydra config for measure stage.", + ) + parser.add_argument( + "--dry_run", + action="store_true", + help="Print what would be executed without actually running stages.", + ) + parser.add_argument( + "--write_requirements", + action="store_true", + help="Write requirements.txt alongside each stage log.", + ) + return parser + + +def main( + args, + system_cls, + stages: Sequence[str] = DEFAULT_STAGES, +) -> None: + stages_to_run = resolve_stages(args.stages, stages) + + # ----------------------------------------- + # Load configs + # ----------------------------------------- + # Keep template_package explicit so the recipe declares which TEMPLATE + # package provides the default configs, instead of relying on path-based + # inference from the user-supplied config location. + training_config = load_and_merge_config( + args.training_config, + template_config_path="conf/training.yaml", + template_package=__package__, + ) + inference_config = load_and_merge_config( + args.inference_config, + template_config_path="conf/inference.yaml", + template_package=__package__, + ) + metrics_config = load_and_merge_config( + args.metrics_config, + template_config_path="conf/metrics.yaml", + template_package=__package__, + ) + + logger = configure_logging() + + # ----------------------------------------- + # Instantiate system + # ----------------------------------------- + system = system_cls( + training_config=training_config, + inference_config=inference_config, + metrics_config=metrics_config, + ) + + # ----------------------------------------- + # Resolve stages and run + # ----------------------------------------- + logger.info("System: %s", system_cls.__name__) + logger.info("Requested stages: %s", args.stages) + logger.info("Resolved stages: %s", stages_to_run) + + # Guardrail: ensure required configs exist for requested stages + pretrain_stages = { + "create_dataset", + "train_tokenizer", + "collect_stats", + "train", + } + required_configs = {} + required_configs.update({stage: training_config for stage in pretrain_stages}) + required_configs.update({"infer": inference_config, "measure": metrics_config}) + missing = [ + s + for s in stages_to_run + if s in required_configs and required_configs[s] is None + ] + if missing: + missing_str = ", ".join(missing) + raise ValueError( + f"Config not provided for stage(s): {missing_str}. " + "Use --training_config/--inference_config/--metrics_config." + ) + run_stages( + system=system, + stages_to_run=stages_to_run, + args=args, + log=logger, + ) + + +if __name__ == "__main__": + parser = build_parser(stages=DEFAULT_STAGES) + args, _ = parse_cli_and_stage_args(parser, stages=DEFAULT_STAGES) + + # Here you should replace `YourSystemClass` with the actual system class + # you want to use for your experiment. + from espnet3.systems.asr.system import ASRSystem + + main( + args=args, + system_cls=ASRSystem, + stages=DEFAULT_STAGES, + ) diff --git a/egs3/TEMPLATE/asr/src/inference.py b/egs3/TEMPLATE/asr/src/inference.py index b8b7399b6733..c6fff782b663 100644 --- a/egs3/TEMPLATE/asr/src/inference.py +++ b/egs3/TEMPLATE/asr/src/inference.py @@ -1,17 +1,17 @@ """Inference output helpers for ASR recipes.""" -def output_fn(*, data, model_output, idx): +def build_output(data, model_output, idx): """Build a dict of outputs for SCP writing.""" - uttid = data.get("uttid", str(idx)) + utt_id = data.get("utt_id", str(idx)) hyp = model_output[0][0] ref = data.get("text", "") - return {"uttid": uttid, "hyp": hyp, "ref": ref} + return {"utt_id": utt_id, "hyp": hyp, "ref": ref} -def output_fn_transducer(*, data, model_output, idx): +def build_output_transducer(data, model_output, idx): """Build a dict of outputs for transducer models.""" - uttid = data.get("uttid", str(idx)) + utt_id = data.get("utt_id", str(idx)) hyp = model_output[0] ref = data.get("text", "") - return {"uttid": uttid, "hyp": hyp, "ref": ref} + return {"utt_id": utt_id, "hyp": hyp, "ref": ref} diff --git a/egs3/__init__.py b/egs3/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/egs3/mini_an4/__init__.py b/egs3/mini_an4/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/egs3/mini_an4/asr/__init__.py b/egs3/mini_an4/asr/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/egs3/mini_an4/asr/conf/inference.yaml b/egs3/mini_an4/asr/conf/inference.yaml new file mode 100644 index 000000000000..96e0f31f927a --- /dev/null +++ b/egs3/mini_an4/asr/conf/inference.yaml @@ -0,0 +1,34 @@ +########################################################## +# Inference SETTING # +########################################################## +# Quick usage: +# python run.py --stages infer --inference_config conf/inference.yaml +# +# Mini AN4 ASR inference config for ESPnet3 integration tests. + +########################################################## +# PATHS AND INPUTS # +########################################################## +exp_tag: ${load_yaml:training.yaml,exp_tag} +inference_dir: ${exp_dir}/inference +dataset_dir: ${data_dir}/mini_an4 + +dataset: + test: + - name: test + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/test.tsv + +########################################################## +# MODEL AND OUTPUT # +########################################################## +model: + _target_: espnet2.bin.asr_inference.Speech2Text + asr_train_config: ${exp_dir}/config.yaml + asr_model_file: ${exp_dir}/last.ckpt + beam_size: 1 + ctc_weight: 0.3 + +input_key: speech +output_fn: src.inference.build_output diff --git a/egs3/mini_an4/asr/conf/inference_transducer.yaml b/egs3/mini_an4/asr/conf/inference_transducer.yaml new file mode 100644 index 000000000000..61a2a9afb486 --- /dev/null +++ b/egs3/mini_an4/asr/conf/inference_transducer.yaml @@ -0,0 +1,32 @@ +########################################################## +# Inference SETTING # +########################################################## +# Quick usage: +# python run.py --stages infer --inference_config conf/inference.yaml +# +# Mini AN4 ASR inference config for ESPnet3 integration tests. +########################################################## +# PATHS AND INPUTS # +########################################################## +exp_dir: ${load_yaml:training.yaml,exp_dir} +inference_dir: ${exp_dir}/inference +dataset_dir: ${load_yaml:training.yaml,dataset_dir} + +dataset: + test: + - name: test + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/test.tsv + +########################################################## +# MODEL AND OUTPUT # +########################################################## +model: + _target_: espnet2.bin.asr_transducer_inference.Speech2Text + asr_training_config: ${exp_dir}/config.yaml + asr_model_file: ${exp_dir}/last.ckpt + beam_size: 1 + +input_key: speech +output_fn: src.inference.build_output_transducer diff --git a/egs3/mini_an4/asr/conf/metrics.yaml b/egs3/mini_an4/asr/conf/metrics.yaml new file mode 100644 index 000000000000..6ef8f210590f --- /dev/null +++ b/egs3/mini_an4/asr/conf/metrics.yaml @@ -0,0 +1,25 @@ +########################################################## +# METRIC SCORING SETTING # +########################################################## +# Quick usage: +# python run.py --stages measure --metrics_config conf/metrics.yaml +# +# Mini AN4 ASR metrics config for ESPnet3 integration tests. + +########################################################## +# PATHS AND INPUTS # +########################################################## +exp_tag: ${load_yaml:training.yaml,exp_tag} +inference_dir: ${exp_dir}/inference + +########################################################## +# METRIC DEFINITION # +########################################################## +metrics: + - metric: + _target_: espnet3.systems.asr.metrics.wer.WER + clean_types: + + - metric: + _target_: espnet3.systems.asr.metrics.cer.CER + clean_types: diff --git a/egs3/mini_an4/asr/conf/training.yaml b/egs3/mini_an4/asr/conf/training.yaml new file mode 120000 index 000000000000..b2b8d88f503f --- /dev/null +++ b/egs3/mini_an4/asr/conf/training.yaml @@ -0,0 +1 @@ +train_transducer_asr_conformer_rnnt_debug.yaml \ No newline at end of file diff --git a/egs3/mini_an4/asr/conf/training_asr_rnn.yaml b/egs3/mini_an4/asr/conf/training_asr_rnn.yaml new file mode 100644 index 000000000000..9f6de38ec692 --- /dev/null +++ b/egs3/mini_an4/asr/conf/training_asr_rnn.yaml @@ -0,0 +1,149 @@ +########################################################## +# TRAINING SETTING # +########################################################## +# This is a debug config for CI. +# Mini AN4 ASR training config for ESPnet3 integration tests. + +########################################################## +# SYSTEM AND PATH SETTING # +########################################################## +task: espnet3.systems.asr.task.ASRTask + +exp_tag: train_asr_rnn_debug +dataset_dir: ${data_dir}/mini_an4 + +########################################################## +# DATASET CREATION SETTING # +########################################################## +# Reuse the downloaded mini_an4 archive to create local manifests. +create_dataset: + func: src.creating_dataset.create_dataset + dataset_dir: ${dataset_dir} + archive_path: ${recipe_dir}/../../egs2/mini_an4/asr1/downloads.tar.gz + +########################################################## +# DATASET DEFINITION # +########################################################## +# Use the tiny train/dev split prepared for CI-sized ASR runs. +dataset: + train: + - name: train_nodev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + valid: + - name: train_dev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_dev.tsv + preprocessor: + _target_: espnet2.train.preprocessor.CommonPreprocessor + fs: 16000 + train: true + # Keep the same augmentation pipeline as the data-aug reference config. + data_aug_effects: + - [0.1, "contrast", {"enhancement_amount": 75.0}] + - [0.1, "highpass", {"cutoff_freq": 5000, "Q": 0.707}] + - [0.1, "equalization", {"center_freq": 1000, "gain": 0, "Q": 0.707}] + - - 0.1 + - - [0.3, "speed_perturb", {"factor": 0.9}] + - [0.3, "speed_perturb", {"factor": 1.1}] + - [0.3, "speed_perturb", {"factor": 1.3}] + data_aug_num: [1, 4] + data_aug_prob: 1.0 + token_type: bpe + token_list: ${tokenizer.save_path}/tokens.txt + bpemodel: ${tokenizer.save_path}/bpe.model + +########################################################## +# TOKENIZER SETTING # +########################################################## +# Build a small BPE model from the training transcript text. +tokenizer: + vocab_size: 30 + character_coverage: 1.0 + model_type: bpe + save_path: ${data_dir}/${tokenizer.model_type}_${tokenizer.vocab_size} + text_builder: + func: src.tokenizer.gather_training_text + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + +########################################################## +# MODEL SETTING # +########################################################## +# Small VGG-RNN encoder/decoder stack for a fast debug run. +model: + vocab_size: ${tokenizer.vocab_size} + token_list: ${tokenizer.save_path}/tokens.txt + + encoder: vgg_rnn + encoder_conf: + num_layers: 1 + hidden_size: 2 + output_size: 2 + + decoder: rnn + decoder_conf: + hidden_size: 2 + + normalize: utterance_mvn + normalize_conf: {} + + model_conf: + ctc_weight: 0.3 + lsm_weight: 0.1 + length_normalized_loss: false + + frontend: default + frontend_conf: + n_fft: 512 + win_length: 400 + hop_length: 160 + +########################################################## +# OPTIMIZER & SCHEDULER SETTING # +########################################################## +optimizer: + _target_: torch.optim.AdamW + lr: 0.001 + weight_decay: 0.0 + +override scheduler: + _target_: torch.optim.lr_scheduler.ReduceLROnPlateau + mode: min + factor: 0.5 + patience: 1 + +########################################################## +# OTHER TRAINING SETTING # +########################################################## +# ReduceLROnPlateau tracks validation loss for LR updates. +val_scheduler_criterion: valid/loss + +best_model_criterion: + - - valid/loss + - 1 + - min + +########################################################## +# DATALOADER SETTING # +########################################################## +# Keep batches tiny so the config stays lightweight in CI. +dataloader: + train: + multiple_iterator: false + num_shards: 1 + iter_factory: + batches: + type: sorted + batch_size: 2 + +########################################################## +# TRAINER SETTING # +########################################################## +trainer: + max_epochs: 1 + limit_train_batches: 1 + limit_val_batches: 1 + + strategy: auto diff --git a/egs3/mini_an4/asr/conf/training_asr_rnn_data_aug.yaml b/egs3/mini_an4/asr/conf/training_asr_rnn_data_aug.yaml new file mode 100644 index 000000000000..f15df91a87b7 --- /dev/null +++ b/egs3/mini_an4/asr/conf/training_asr_rnn_data_aug.yaml @@ -0,0 +1,146 @@ +########################################################## +# TRAINING SETTING # +########################################################## +# Quick usage: +# python run.py --stages create_dataset --training_config conf/training.yaml +# python run.py --stages train --training_config conf/training.yaml +# +# This is a debug config for CI +# Mini AN4 ASR training config for ESPnet3 integration tests. + +########################################################## +# SYSTEM AND PATH SETTING # +########################################################## +task: espnet3.systems.asr.task.ASRTask + +exp_tag: train_asr_rnn_data_aug_debug +dataset_dir: ${data_dir}/mini_an4 + +########################################################## +# DATASET CREATION SETTING # +########################################################## +create_dataset: + func: src.creating_dataset.create_dataset + dataset_dir: ${dataset_dir} + archive_path: ${recipe_dir}/../../egs2/mini_an4/asr1/downloads.tar.gz + +########################################################## +# DATASET DEFINITION # +########################################################## +dataset: + train: + - name: train_nodev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + valid: + - name: train_dev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_dev.tsv + preprocessor: + _target_: espnet2.train.preprocessor.CommonPreprocessor + fs: 16000 + train: true + data_aug_effects: # no need to set the "sample_rate" argument for each effect here + - [0.1, "contrast", {"enhancement_amount": 75.0}] + - [0.1, "highpass", {"cutoff_freq": 5000, "Q": 0.707}] + - [0.1, "equalization", {"center_freq": 1000, "gain": 0, "Q": 0.707}] + - - 0.1 + - - [0.3, "speed_perturb", {"factor": 0.9}] + - [0.3, "speed_perturb", {"factor": 1.1}] + - [0.3, "speed_perturb", {"factor": 1.3}] + data_aug_num: [1, 4] + data_aug_prob: 1.0 + token_type: bpe + token_list: ${tokenizer.save_path}/tokens.txt + bpemodel: ${tokenizer.save_path}/bpe.model + +########################################################## +# TOKENIZER SETTING # +########################################################## +tokenizer: + vocab_size: 30 + character_coverage: 1.0 + model_type: bpe + save_path: ${data_dir}/${tokenizer.model_type}_${tokenizer.vocab_size} + text_builder: + func: src.tokenizer.gather_training_text + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + +########################################################## +# MODEL SETTING # +########################################################## +model: + vocab_size: ${tokenizer.vocab_size} + token_list: ${tokenizer.save_path}/tokens.txt + + encoder: vgg_rnn + encoder_conf: + num_layers: 1 + hidden_size: 2 + output_size: 2 + + decoder: rnn + decoder_conf: + hidden_size: 2 + + normalize: utterance_mvn + normalize_conf: {} + + model_conf: + ctc_weight: 0.3 + lsm_weight: 0.1 + length_normalized_loss: false + + frontend: default + frontend_conf: + n_fft: 512 + win_length: 400 + hop_length: 160 + +########################################################## +# OPTIMIZER & SCHEDULER SETTING # +########################################################## +optimizer: + _target_: torch.optim.AdamW + lr: 0.001 + weight_decay: 0.0 + +override scheduler: + _target_: torch.optim.lr_scheduler.ReduceLROnPlateau + mode: min + factor: 0.5 + patience: 1 + +########################################################## +# OTHER TRAINING SETTING # +########################################################## +val_scheduler_criterion: valid/loss + +best_model_criterion: + - - valid/acc + - 1 + - max + +########################################################## +# DATALOADER SETTING # +########################################################## +dataloader: + train: + iter_factory: + batches: + type: sorted + batch_size: 2 + +########################################################## +# TRAINER SETTING # +########################################################## +# These settings map to `lightning.Trainer(...)`. +trainer: + log_every_n_steps: 1 + max_epochs: 1 + limit_train_batches: 1 + limit_val_batches: 1 + + strategy: auto diff --git a/egs3/mini_an4/asr/conf/training_asr_streaming.yaml b/egs3/mini_an4/asr/conf/training_asr_streaming.yaml new file mode 100644 index 000000000000..1f7fd31deaba --- /dev/null +++ b/egs3/mini_an4/asr/conf/training_asr_streaming.yaml @@ -0,0 +1,141 @@ +########################################################## +# TRAINING SETTING # +########################################################## +# This is a debug config for CI. +# Mini AN4 streaming ASR training config for ESPnet3 integration tests. + +########################################################## +# SYSTEM AND PATH SETTING # +########################################################## +task: espnet3.systems.asr.task.ASRTask + +exp_tag: train_asr_streaming_debug +dataset_dir: ${data_dir}/mini_an4 + +########################################################## +# DATASET CREATION SETTING # +########################################################## +create_dataset: + func: src.creating_dataset.create_dataset + dataset_dir: ${dataset_dir} + archive_path: ${recipe_dir}/../../egs2/mini_an4/asr1/downloads.tar.gz + +########################################################## +# DATASET DEFINITION # +########################################################## +# Standard mini_an4 train/dev manifests plus BPE preprocessing. +dataset: + train: + - name: train_nodev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + valid: + - name: train_dev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_dev.tsv + preprocessor: + _target_: espnet2.train.preprocessor.CommonPreprocessor + train: true + token_type: bpe + token_list: ${tokenizer.save_path}/tokens.txt + bpemodel: ${tokenizer.save_path}/bpe.model + text_cleaner: + +########################################################## +# TOKENIZER SETTING # +########################################################## +tokenizer: + vocab_size: 30 + character_coverage: 1.0 + model_type: bpe + save_path: ${data_dir}/${tokenizer.model_type}_${tokenizer.vocab_size} + text_builder: + func: src.tokenizer.gather_training_text + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + +########################################################## +# MODEL SETTING # +########################################################## +# Contextual block encoder keeps this config focused on streaming behavior. +model: + vocab_size: ${tokenizer.vocab_size} + token_list: ${tokenizer.save_path}/tokens.txt + + encoder: contextual_block_transformer + encoder_conf: + block_size: 40 + hop_size: 16 + look_ahead: 16 + output_size: 2 + attention_heads: 2 + linear_units: 2 + num_blocks: 1 + + decoder: transformer + decoder_conf: + attention_heads: 2 + linear_units: 2 + num_blocks: 1 + + normalize: utterance_mvn + normalize_conf: {} + + model: espnet + model_conf: + ctc_weight: 0.3 + lsm_weight: 0.1 + length_normalized_loss: false + # Streaming config disables feature extraction during stats collection. + extract_feats_in_collect_stats: false + + frontend: default + frontend_conf: + n_fft: 512 + win_length: 400 + hop_length: 160 + +########################################################## +# OPTIMIZER & SCHEDULER SETTING # +########################################################## +optimizer: + lr: 0.005 + +scheduler: + warmup_steps: 100 + +########################################################## +# OTHER TRAINING SETTING # +########################################################## +# Use Xavier init because the model is transformer-based. +init: xavier_uniform + +best_model_criterion: + - - valid/loss + - 1 + - min + +########################################################## +# DATALOADER SETTING # +########################################################## +# Unsorted mini-batches keep the streaming example simple. +dataloader: + train: + multiple_iterator: false + num_shards: 1 + iter_factory: + batches: + type: unsorted + batch_size: 2 + +########################################################## +# TRAINER SETTING # +########################################################## +trainer: + log_every_n_steps: 1 + max_epochs: 1 + limit_train_batches: 1 + limit_val_batches: 1 + + strategy: auto diff --git a/egs3/mini_an4/asr/conf/training_asr_transducer.yaml b/egs3/mini_an4/asr/conf/training_asr_transducer.yaml new file mode 100644 index 000000000000..10e379609c68 --- /dev/null +++ b/egs3/mini_an4/asr/conf/training_asr_transducer.yaml @@ -0,0 +1,143 @@ +########################################################## +# TRAINING SETTING # +########################################################## +# This is a debug config for CI. +# Mini AN4 transducer ASR training config for ESPnet3 integration tests. + +########################################################## +# SYSTEM AND PATH SETTING # +########################################################## +task: espnet3.systems.asr.task.ASRTask + +exp_tag: train_asr_transducer_debug +dataset_dir: ${data_dir}/mini_an4 + +########################################################## +# DATASET CREATION SETTING # +########################################################## +create_dataset: + func: src.creating_dataset.create_dataset + dataset_dir: ${dataset_dir} + archive_path: ${recipe_dir}/../../egs2/mini_an4/asr1/downloads.tar.gz + +########################################################## +# DATASET DEFINITION # +########################################################## +# Share the same tiny manifests across the ASR debug configs. +dataset: + train: + - name: train_nodev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + valid: + - name: train_dev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_dev.tsv + preprocessor: + _target_: espnet2.train.preprocessor.CommonPreprocessor + train: true + token_type: bpe + token_list: ${tokenizer.save_path}/tokens.txt + bpemodel: ${tokenizer.save_path}/bpe.model + text_cleaner: + +########################################################## +# TOKENIZER SETTING # +########################################################## +tokenizer: + vocab_size: 30 + character_coverage: 1.0 + model_type: bpe + save_path: ${data_dir}/${tokenizer.model_type}_${tokenizer.vocab_size} + text_builder: + func: src.tokenizer.gather_training_text + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + +########################################################## +# MODEL SETTING # +########################################################## +# Transformer encoder plus RNNT-style decoder/joint network. +model: + vocab_size: ${tokenizer.vocab_size} + token_list: ${tokenizer.save_path}/tokens.txt + + encoder: transformer + encoder_conf: + output_size: 2 + attention_heads: 2 + linear_units: 2 + num_blocks: 2 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.0 + input_layer: conv1d2 + normalize_before: true + + decoder: transducer + decoder_conf: + rnn_type: lstm + num_layers: 1 # Decoder layers in the prediction network. + hidden_size: 4 # Decoder hidden dimension. + dropout: 0.1 + dropout_embed: 0.2 + + joint_net_conf: + joint_space_size: 4 + + normalize: utterance_mvn + normalize_conf: {} + + model: espnet + model_conf: + ctc_weight: 0.3 + lsm_weight: 0.1 + length_normalized_loss: false + + frontend: default + frontend_conf: + n_fft: 512 + win_length: 400 + hop_length: 160 + +########################################################## +# OPTIMIZER & SCHEDULER SETTING # +########################################################## +optimizer: + lr: 0.005 + +scheduler: + warmup_steps: 100 + +########################################################## +# OTHER TRAINING SETTING # +########################################################## +# Xavier init is the default debug choice for this transformer setup. +init: xavier_uniform + +best_model_criterion: + - - valid/loss + - 1 + - min + +########################################################## +# DATALOADER SETTING # +########################################################## +dataloader: + train: + iter_factory: + batches: + type: unsorted + batch_size: 2 + +########################################################## +# TRAINER SETTING # +########################################################## +trainer: + log_every_n_steps: 1 + max_epochs: 1 + limit_train_batches: 1 + limit_val_batches: 1 + + strategy: auto diff --git a/egs3/mini_an4/asr/conf/training_asr_transformer.yaml b/egs3/mini_an4/asr/conf/training_asr_transformer.yaml new file mode 100644 index 000000000000..2e41e9dbcd84 --- /dev/null +++ b/egs3/mini_an4/asr/conf/training_asr_transformer.yaml @@ -0,0 +1,141 @@ +########################################################## +# TRAINING SETTING # +########################################################## +# This is a debug config for CI. +# Mini AN4 transformer ASR training config for ESPnet3 integration tests. + +########################################################## +# SYSTEM AND PATH SETTING # +########################################################## +task: espnet3.systems.asr.task.ASRTask + +exp_tag: train_asr_transformer_debug +dataset_dir: ${data_dir}/mini_an4 + +########################################################## +# DATASET CREATION SETTING # +########################################################## +create_dataset: + func: src.creating_dataset.create_dataset + dataset_dir: ${dataset_dir} + archive_path: ${recipe_dir}/../../egs2/mini_an4/asr1/downloads.tar.gz + +########################################################## +# DATASET DEFINITION # +########################################################## +# Minimal train/dev split for transformer smoke tests. +dataset: + train: + - name: train_nodev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + valid: + - name: train_dev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_dev.tsv + preprocessor: + _target_: espnet2.train.preprocessor.CommonPreprocessor + train: true + token_type: bpe + token_list: ${tokenizer.save_path}/tokens.txt + bpemodel: ${tokenizer.save_path}/bpe.model + text_cleaner: + +########################################################## +# TOKENIZER SETTING # +########################################################## +tokenizer: + vocab_size: 30 + character_coverage: 1.0 + model_type: bpe + save_path: ${data_dir}/${tokenizer.model_type}_${tokenizer.vocab_size} + text_builder: + func: src.tokenizer.gather_training_text + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + +########################################################## +# MODEL SETTING # +########################################################## +# Tiny transformer encoder/decoder pair to keep CI fast. +model: + vocab_size: ${tokenizer.vocab_size} + token_list: ${tokenizer.save_path}/tokens.txt + + encoder: transformer + encoder_conf: + output_size: 2 + attention_heads: 2 + linear_units: 2 + num_blocks: 2 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.0 + input_layer: conv1d2 + normalize_before: true + + decoder: transformer + decoder_conf: + attention_heads: 2 + linear_units: 2 + num_blocks: 2 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + self_attention_dropout_rate: 0.0 + src_attention_dropout_rate: 0.0 + + normalize: utterance_mvn + normalize_conf: {} + + model: espnet + model_conf: + ctc_weight: 0.3 + lsm_weight: 0.1 + length_normalized_loss: false + + frontend: default + frontend_conf: + n_fft: 512 + win_length: 400 + hop_length: 160 + +########################################################## +# OPTIMIZER & SCHEDULER SETTING # +########################################################## +optimizer: + lr: 0.005 + +scheduler: + warmup_steps: 100 + +########################################################## +# OTHER TRAINING SETTING # +########################################################## +# Xavier init is used to stabilize this transformer baseline. +init: xavier_uniform + +best_model_criterion: + - - valid/loss + - 10 + - min + +########################################################## +# DATALOADER SETTING # +########################################################## +dataloader: + train: + iter_factory: + batches: + type: unsorted + batch_size: 2 + +########################################################## +# TRAINER SETTING # +########################################################## +trainer: + max_epochs: 1 + limit_train_batches: 1 + limit_val_batches: 1 + + strategy: auto diff --git a/egs3/mini_an4/asr/conf/training_lm_rnn.yaml b/egs3/mini_an4/asr/conf/training_lm_rnn.yaml new file mode 100644 index 000000000000..573337f5107f --- /dev/null +++ b/egs3/mini_an4/asr/conf/training_lm_rnn.yaml @@ -0,0 +1,125 @@ +########################################################## +# TRAINING SETTING # +########################################################## +# This is a debug config for CI. +# Mini AN4 LM-RNN training config for ESPnet3 integration tests. + +########################################################## +# SYSTEM AND PATH SETTING # +########################################################## +task: espnet3.systems.asr.task.ASRTask + +exp_tag: train_lm_rnn_debug +dataset_dir: ${data_dir}/mini_an4 + +########################################################## +# DATASET CREATION SETTING # +########################################################## +create_dataset: + func: src.creating_dataset.create_dataset + dataset_dir: ${dataset_dir} + archive_path: ${recipe_dir}/../../egs2/mini_an4/asr1/downloads.tar.gz + +########################################################## +# DATASET DEFINITION # +########################################################## +# Reuse the ASR manifests and tokenize the transcript text for LM training. +dataset: + train: + - name: train_nodev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + valid: + - name: train_dev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_dev.tsv + preprocessor: + _target_: espnet2.train.preprocessor.CommonPreprocessor + train: true + token_type: bpe + token_list: ${tokenizer.save_path}/tokens.txt + bpemodel: ${tokenizer.save_path}/bpe.model + text_cleaner: + +########################################################## +# TOKENIZER SETTING # +########################################################## +tokenizer: + vocab_size: 30 + character_coverage: 1.0 + model_type: bpe + save_path: ${data_dir}/${tokenizer.model_type}_${tokenizer.vocab_size} + text_builder: + func: src.tokenizer.gather_training_text + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + +########################################################## +# MODEL SETTING # +########################################################## +# Tiny recurrent LM for debug-only language-model coverage. +model: + vocab_size: ${tokenizer.vocab_size} + token_list: ${tokenizer.save_path}/tokens.txt + + lm_conf: + nlayers: 1 + unit: 2 + + normalize: utterance_mvn + normalize_conf: {} + + model: espnet + model_conf: + ctc_weight: 0.3 + lsm_weight: 0.1 + length_normalized_loss: false + + frontend: default + frontend_conf: + n_fft: 512 + win_length: 400 + hop_length: 160 + +########################################################## +# OPTIMIZER & SCHEDULER SETTING # +########################################################## +optimizer: + lr: 0.005 + +scheduler: + warmup_steps: 30000 + +########################################################## +# OTHER TRAINING SETTING # +########################################################## +# Xavier init is kept to match the other debug baselines. +init: xavier_uniform + +best_model_criterion: + - - valid/acc + - 10 + - max + +########################################################## +# DATALOADER SETTING # +########################################################## +# Larger batches are fine here because the LM input is text-only. +dataloader: + train: + iter_factory: + batches: + type: sorted + batch_size: 32 + +########################################################## +# TRAINER SETTING # +########################################################## +trainer: + log_every_n_steps: 1 + max_epochs: 100 + limit_train_batches: 1 + limit_val_batches: 1 + + strategy: auto diff --git a/egs3/mini_an4/asr/conf/training_transducer_asr_conformer_rnnt.yaml b/egs3/mini_an4/asr/conf/training_transducer_asr_conformer_rnnt.yaml new file mode 100644 index 000000000000..d35027a259c2 --- /dev/null +++ b/egs3/mini_an4/asr/conf/training_transducer_asr_conformer_rnnt.yaml @@ -0,0 +1,126 @@ +########################################################## +# TRAINING SETTING # +########################################################## +# This is a debug config for CI. +# Mini AN4 conformer-RNNT training config for ESPnet3 integration tests. + +########################################################## +# SYSTEM AND PATH SETTING # +########################################################## +num_device: 1 +num_nodes: 1 +task: espnet3.systems.asr.transducer_task.ASRTransducerTask + +recipe_dir: . +data_dir: ${recipe_dir}/data +exp_tag: train_transducer_asr_conformer_rnnt_debug +exp_dir: ${recipe_dir}/exp/${exp_tag} +stats_dir: ${recipe_dir}/exp/stats +inference_dir: ${exp_dir}/inference +dataset_dir: ${data_dir}/mini_an4 + +########################################################## +# DATASET CREATION SETTING # +########################################################## +create_dataset: + func: src.creating_dataset.create_dataset + dataset_dir: ${dataset_dir} + archive_path: ${recipe_dir}/../../egs2/mini_an4/asr1/downloads.tar.gz + +########################################################## +# DATASET DEFINITION # +########################################################## +# Use the same manifest layout as the other mini_an4 debug configs. +dataset: + train: + - name: train_nodev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + valid: + - name: train_dev + dataset: + _target_: src.dataset.MiniAN4Dataset + manifest_path: ${dataset_dir}/manifest/train_dev.tsv + preprocessor: + _target_: espnet2.train.preprocessor.CommonPreprocessor + train: true + token_type: bpe + token_list: ${tokenizer.save_path}/tokens.txt + bpemodel: ${tokenizer.save_path}/bpe.model + text_cleaner: + +########################################################## +# TOKENIZER SETTING # +########################################################## +tokenizer: + vocab_size: 30 + character_coverage: 1.0 + model_type: bpe + save_path: ${data_dir}/${tokenizer.model_type}_${tokenizer.vocab_size} + text_builder: + func: src.tokenizer.gather_training_text + manifest_path: ${dataset_dir}/manifest/train_nodev.tsv + +########################################################## +# MODEL SETTING # +########################################################## +# Compact conformer encoder and RNNT decoder for task-specific coverage. +model: + vocab_size: ${tokenizer.vocab_size} + token_list: ${tokenizer.save_path}/tokens.txt + + encoder: conformer + encoder_conf: + body_conf: + - block_type: conformer + hidden_size: 2 + linear_size: 4 + heads: 2 + conv_mod_kernel_size: 3 + + decoder_conf: + embed_size: 4 + hidden_size: 4 + + joint_network_conf: + joint_space_size: 4 + +########################################################## +# OPTIMIZER & SCHEDULER SETTING # +########################################################## +optimizer: + lr: 0.005 + +scheduler: + warmup_steps: 30000 + +########################################################## +# OTHER TRAINING SETTING # +########################################################## +best_model_criterion: + - - valid/loss + - 10 + - min + +########################################################## +# DATALOADER SETTING # +########################################################## +# Sorted batches reduce padding for the transducer example. +dataloader: + train: + iter_factory: + batches: + type: sorted + batch_size: 32 + +########################################################## +# TRAINER SETTING # +########################################################## +trainer: + log_every_n_steps: 1 + max_epochs: 1 + limit_train_batches: 1 + limit_val_batches: 1 + + strategy: auto diff --git a/egs3/mini_an4/asr/path.sh b/egs3/mini_an4/asr/path.sh new file mode 100644 index 000000000000..1456ecff938d --- /dev/null +++ b/egs3/mini_an4/asr/path.sh @@ -0,0 +1,6 @@ +#!/bin/bash + +export PYTHONPATH=../../../:../../TEMPLATE/asr:$(pwd):${PYTHONPATH} + +source ../../../tools/activate_python.sh +source ../../../tools/extra_path.sh diff --git a/egs3/mini_an4/asr/run.py b/egs3/mini_an4/asr/run.py new file mode 100644 index 000000000000..46470cc0db75 --- /dev/null +++ b/egs3/mini_an4/asr/run.py @@ -0,0 +1,19 @@ +from egs3.TEMPLATE.asr.run import ( + DEFAULT_STAGES, + build_parser, + main, + parse_cli_and_stage_args, +) +from espnet3.systems.asr.system import ASRSystem + +if __name__ == "__main__": + parser = build_parser( + stages=DEFAULT_STAGES, + ) + args, _ = parse_cli_and_stage_args(parser, stages=DEFAULT_STAGES) + + main( + args=args, + system_cls=ASRSystem, + stages=DEFAULT_STAGES, + ) diff --git a/egs3/mini_an4/asr/src/__init__.py b/egs3/mini_an4/asr/src/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/egs3/mini_an4/asr/src/creating_dataset.py b/egs3/mini_an4/asr/src/creating_dataset.py new file mode 100644 index 000000000000..fc4c4774d50f --- /dev/null +++ b/egs3/mini_an4/asr/src/creating_dataset.py @@ -0,0 +1,154 @@ +"""Prepare the Mini AN4 dataset for ESPnet3 ASR recipes (simplified).""" + +from __future__ import annotations + +import argparse +import logging +import re +import shutil +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + +from espnet3.utils.download_utils import extract_targz, setup_logger + +LOGGER = logging.getLogger(__name__) + +SPH_DIRS = {"train": "an4_clstk", "test": "an4test_clstk"} +TRANSCRIPTS = {"train": "an4_train.transcription", "test": "an4_test.transcription"} + +# Expect lines like " WORDS (a-b-c)" +LINE_RE = re.compile(r"^(?P.+?)\s+\((?P[^)]+)\)\s*$") + + +@dataclass(frozen=True) +class Entry: + utt_id: str + wav_path: Path + text: str + + +def is_gzip(path: Path) -> bool: + return path.is_file() and path.read_bytes()[:2] == b"\x1f\x8b" + + +def ensure_extracted(archive: Path, dataset_dir: Path, logger: logging.Logger) -> Path: + downloads_dir = dataset_dir / "downloads" + an4_root = downloads_dir / "an4" + if an4_root.exists(): + logger.info("Found existing downloads: %s", an4_root) + return an4_root + + if not archive.is_file(): + raise FileNotFoundError(f"Archive not found: {archive}") + if not is_gzip(archive): + raise RuntimeError(f"Archive is not gzip: {archive}") + + dataset_dir.mkdir(parents=True, exist_ok=True) + extract_targz(archive, dataset_dir, logger=logger) + + if not an4_root.exists(): + raise RuntimeError(f"AN4 root not found after extract: {an4_root}") + return an4_root + + +def parse_line(line: str) -> tuple[str, str, str]: + m = LINE_RE.match(line) + if not m: + raise ValueError(f"Malformed transcript line: {line}") + + words = m.group("words") + # Remove and from the transcript (keeps current behavior) + if words.startswith(" "): + words = words[4:] + if words.endswith(" "): + words = words[:-5] + + src = m.group("src") # Example: abc-def-ghi + pre, mid, last = src.split("-") + utt_id = f"{mid}-{pre}-{last}" + return utt_id, src, words + + +def load_entries(an4_root: Path, split: str) -> list[tuple[str, str, str]]: + path = an4_root / "etc" / TRANSCRIPTS[split] + if not path.is_file(): + raise FileNotFoundError(f"Transcript not found: {path}") + + entries: list[tuple[str, str, str]] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if line: + entries.append(parse_line(line)) + + entries.sort(key=lambda x: x[0]) + return entries + + +def convert_sph_to_wav(sph2pipe: str, sph: Path, wav: Path) -> None: + if wav.exists(): + return + wav.parent.mkdir(parents=True, exist_ok=True) + cmd = [sph2pipe, "-f", "wav", "-p", "-c", "1", str(sph)] + with wav.open("wb") as f: + subprocess.run(cmd, stdout=f, check=True) + + +def prepare_split( + an4_root: Path, dataset_dir: Path, split: str, sph2pipe: str +) -> list[Entry]: + wav_dir = dataset_dir / "wav" / split + sph_dir = an4_root / "wav" / SPH_DIRS[split] + + out: list[Entry] = [] + for utt_id, src, text in load_entries(an4_root, split): + mid = src.split("-")[1] + sph = sph_dir / mid / f"{src}.sph" + if not sph.is_file(): + raise FileNotFoundError(f"Missing sph file: {sph}") + + wav = (wav_dir / f"{utt_id}.wav").resolve() + convert_sph_to_wav(sph2pipe, sph, wav) + out.append(Entry(utt_id=utt_id, wav_path=wav, text=text)) + return out + + +def write_manifest(path: Path, entries: Iterable[Entry]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + for e in entries: + f.write(f"{e.utt_id}\t{e.wav_path}\t{e.text}\n") + + +def create_dataset(dataset_dir: Path, archive_path: Path | None = None) -> None: + dataset_dir = Path(dataset_dir) + logger = setup_logger(name="mini_an4.create_dataset") + + archive = Path("../../../egs2/mini_an4/asr1/downloads.tar.gz") + an4_root = ensure_extracted(archive, dataset_dir, logger=logger) + + sph2pipe = shutil.which("sph2pipe") + if not sph2pipe: + raise RuntimeError("sph2pipe not found in PATH. Please install it.") + + train = prepare_split(an4_root, dataset_dir, "train", sph2pipe) + test = prepare_split(an4_root, dataset_dir, "test", sph2pipe) + + if len(train) < 2: + raise RuntimeError("Training data is too small (need >= 2 for dev+nodev).") + + manifest_dir = dataset_dir / "manifest" + write_manifest(manifest_dir / "train_dev.tsv", train[:1]) + write_manifest(manifest_dir / "train_nodev.tsv", train[1:]) + write_manifest(manifest_dir / "test.tsv", test) + + LOGGER.info("Prepared Mini AN4 manifests under %s", manifest_dir) + + +if __name__ == "__main__": + p = argparse.ArgumentParser(description="Prepare Mini AN4 dataset") + p.add_argument("--dataset_dir", type=Path, required=True) + p.add_argument("--archive_path", type=Path, default=None) + a = p.parse_args() + create_dataset(a.dataset_dir, archive_path=a.archive_path) diff --git a/egs3/mini_an4/asr/src/dataset.py b/egs3/mini_an4/asr/src/dataset.py new file mode 100644 index 000000000000..ed022deca65f --- /dev/null +++ b/egs3/mini_an4/asr/src/dataset.py @@ -0,0 +1,57 @@ +"""Dataset loader for Mini AN4 manifests.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import List + +import numpy as np +import soundfile as sf +from torch.utils.data import Dataset as TorchDataset + + +@dataclass(frozen=True) +class ManifestEntry: + utt_id: str + wav_path: Path + text: str + + +def _read_manifest(manifest_path: Path) -> List[ManifestEntry]: + entries: List[ManifestEntry] = [] + with manifest_path.open("r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + utt_id, wav_path, text = line.split("\t", maxsplit=2) + entries.append( + ManifestEntry( + utt_id=utt_id, + wav_path=Path(wav_path), + text=text, + ) + ) + if not entries: + raise RuntimeError(f"Manifest is empty: {manifest_path}") + return entries + + +class MiniAN4Dataset(TorchDataset): + def __init__(self, manifest_path: str | Path) -> None: + manifest_path = Path(manifest_path) + if not manifest_path.is_file(): + raise FileNotFoundError(f"Manifest not found: {manifest_path}") + self._entries = _read_manifest(manifest_path) + + def __len__(self) -> int: + return len(self._entries) + + def __getitem__(self, idx: int): + entry = self._entries[int(idx)] + array, _sr = sf.read(str(entry.wav_path)) + return { + "speech": np.asarray(array, dtype=np.float32), + "text": entry.text, + } diff --git a/egs3/mini_an4/asr/src/inference.py b/egs3/mini_an4/asr/src/inference.py new file mode 100644 index 000000000000..696d32ed8eb3 --- /dev/null +++ b/egs3/mini_an4/asr/src/inference.py @@ -0,0 +1,17 @@ +"""Inference output helpers for Mini AN4 ASR recipes.""" + + +def build_output(data, model_output, idx): + """Build a dict of outputs for SCP writing.""" + utt_id = data.get("utt_id", str(idx)) + hyp = model_output[0][0] + ref = data.get("text", "") + return {"utt_id": utt_id, "hyp": hyp, "ref": ref} + + +def build_output_transducer(data, model_output, idx): + """Build a dict of outputs for transducer models.""" + utt_id = data.get("utt_id", str(idx)) + hyp = model_output[0] + ref = data.get("text", "") + return {"utt_id": utt_id, "hyp": hyp, "ref": ref} diff --git a/egs3/mini_an4/asr/src/tokenizer.py b/egs3/mini_an4/asr/src/tokenizer.py new file mode 100644 index 000000000000..56e677070b6d --- /dev/null +++ b/egs3/mini_an4/asr/src/tokenizer.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +from pathlib import Path +from typing import List + + +def gather_training_text(manifest_path: Path) -> List[str]: + """Collect transcript texts from a TSV manifest file.""" + manifest_path = Path(manifest_path) + if not manifest_path.is_file(): + raise FileNotFoundError(f"Manifest not found: {manifest_path}") + + texts: list[str] = [] + with manifest_path.open("r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + parts = line.split("\t", maxsplit=2) + if len(parts) != 3: + raise ValueError(f"Invalid manifest line: {line}") + texts.append(parts[2]) + + if not texts: + raise RuntimeError(f"No text found in manifest: {manifest_path}") + return texts diff --git a/espnet3/components/data/data_organizer.py b/espnet3/components/data/data_organizer.py index 281c68e709f0..7f1550103aa2 100644 --- a/espnet3/components/data/data_organizer.py +++ b/espnet3/components/data/data_organizer.py @@ -143,7 +143,7 @@ class DataOrganizer: Example (training + validation): >>> organizer = DataOrganizer( - ... train=train_configs, + ... train=training_configs, ... valid=valid_configs, ... preprocessor=MyPreprocessor() ... ) @@ -181,10 +181,12 @@ def build_dataset_list(config_list): dataset = config.dataset if isinstance(dataset, (dict, DictConfig)): dataset = instantiate(dataset) - if hasattr(config, "transform"): + + if hasattr(config, "transform") and config.transform is not None: transform = config.transform else: transform = do_nothing + if isinstance(transform, (dict, DictConfig)): transform = instantiate(transform) @@ -226,11 +228,21 @@ def build_dataset_list(config_list): self.test_sets = {} if test is not None: for config in test: + if isinstance(config, dict): + config = DatasetConfig(**config) + dataset = config.dataset - if hasattr(config, "transform"): + if isinstance(dataset, (dict, DictConfig)): + dataset = instantiate(dataset) + + if hasattr(config, "transform") and config.transform is not None: transform = config.transform else: transform = do_nothing + + if isinstance(transform, (dict, DictConfig)): + transform = instantiate(transform) + self.test_sets[config.name] = DatasetWithTransform( dataset, transform, diff --git a/espnet3/components/modeling/lightning_module.py b/espnet3/components/modeling/lightning_module.py index 15825803b8e6..ac22519da668 100644 --- a/espnet3/components/modeling/lightning_module.py +++ b/espnet3/components/modeling/lightning_module.py @@ -248,6 +248,17 @@ def configure_optimizers(self): AssertionError: If configuration rules are violated. ValueError: If neither optimizer configuration is provided. """ + + def _get_val_scheduler_monitor(): + criterion = getattr(self.config, "val_scheduler_criterion", None) + if criterion is None: + return None + if not isinstance(criterion, str): + raise ValueError( + "val_scheduler_criterion must be a string like 'valid/loss'" + ) + return criterion + if getattr(self.config, "optimizer", None) and getattr( self.config, "scheduler", None ): @@ -352,19 +363,22 @@ def configure_optimizers(self): "`schedulers`" ) - self._log_training_summary( - logger, - self.model, - optimizer=optimizer, - scheduler=scheduler, - ) + def _build_lr_scheduler_config(sch): + config = {"scheduler": sch, "interval": "step"} + monitor = _get_val_scheduler_monitor() + if monitor is not None: + config["interval"] = "epoch" + config["monitor"] = monitor + return config + + if isinstance(scheduler, list): + lr_scheduler = [_build_lr_scheduler_config(sch) for sch in scheduler] + else: + lr_scheduler = _build_lr_scheduler_config(scheduler) return { "optimizer": optimizer, - "lr_scheduler": { - "scheduler": scheduler, - "interval": "step", # assuming lr scheduler is updated per step - }, + "lr_scheduler": lr_scheduler, } @staticmethod @@ -561,6 +575,11 @@ def collect_stats(self): ) for mode in ["train", "valid"]: + if mode == "train": + dataset_config.preprocessor.train = True + else: + dataset_config.preprocessor.train = False + collect_stats( model_config=OmegaConf.to_container(self.config.model, resolve=True), dataset_config=dataset_config, diff --git a/espnet3/components/trainers/trainer.py b/espnet3/components/trainers/trainer.py index 8c48a778a2c7..3c6804000874 100644 --- a/espnet3/components/trainers/trainer.py +++ b/espnet3/components/trainers/trainer.py @@ -1,5 +1,6 @@ """Trainer class for the espnet3 package.""" +import copy import warnings from argparse import Namespace from typing import Any, Dict, Union @@ -75,28 +76,18 @@ def __init__( # Accelerator accelerator = _get_or_initialize(self.config, "accelerator", "auto") - if hasattr(self.config, "accelerator"): - self._del_config_key("accelerator") # strategy strategy = _get_or_initialize(self.config, "strategy", "auto") - if hasattr(self.config, "strategy"): - self._del_config_key("strategy") # logger logger = _get_or_initialize(self.config, "logger") - if logger is not None: - self._del_config_key("logger") # profiler profiler = _get_or_initialize(self.config, "profiler") - if profiler is not None: - self._del_config_key("profiler") # plugins plugins = _get_or_initialize(self.config, "plugins") - if plugins is not None: - self._del_config_key("plugins") # Callbacks callbacks = get_default_callbacks( @@ -110,7 +101,6 @@ def __init__( ), "callbacks should be a list" for callback in self.config.callbacks: callbacks.append(instantiate(callback)) - self._del_config_key("callbacks") # Since espnet's sampler requires to set the following configs: # Reload dataloaders every epoch to reuse ESPnet's dataloader @@ -144,6 +134,17 @@ def __init__( if self.model.is_espnet_sampler: self.config.use_distributed_sampler = False + trainer_config = copy.deepcopy(self.config) + for key in ( + "accelerator", + "strategy", + "logger", + "profiler", + "plugins", + "callbacks", + ): + self._del_config_key_on(trainer_config, key) + # Set up the trainer self.trainer = lightning.Trainer( accelerator=accelerator, @@ -153,7 +154,7 @@ def __init__( logger=logger, profiler=profiler, plugins=plugins, - **self.config, + **trainer_config, ) def _del_config_key(self, key): @@ -162,6 +163,14 @@ def _del_config_key(self, key): elif isinstance(self.config, dict): self.config.pop(key) + @staticmethod + def _del_config_key_on(config, key): + if isinstance(config, DictConfig) or isinstance(config, Namespace): + if hasattr(config, key): + delattr(config, key) + elif isinstance(config, dict): + config.pop(key, None) + def fit(self, *args, **kwargs): """Start the training loop using Lightning's fit method. diff --git a/espnet3/systems/asr/system.py b/espnet3/systems/asr/system.py index b3f76a8df2b2..d3215376ba12 100644 --- a/espnet3/systems/asr/system.py +++ b/espnet3/systems/asr/system.py @@ -44,23 +44,23 @@ class ASRSystem(BaseSystem): Additional stage log paths: | Stage | Path reference | |--- |--- | - | train_tokenizer | train_config.tokenizer.save_path | + | train_tokenizer | training_config.tokenizer.save_path | """ def __init__( self, - train_config=None, - infer_config=None, - measure_config=None, + training_config=None, + inference_config=None, + metrics_config=None, **kwargs, ) -> None: """Initialize the ASR system with ASR-specific stage mappings.""" super().__init__( - train_config=train_config, - infer_config=infer_config, - measure_config=measure_config, + training_config=training_config, + inference_config=inference_config, + metrics_config=metrics_config, stage_log_mapping={ - "train_tokenizer": "train_config.tokenizer.save_path", + "train_tokenizer": "training_config.tokenizer.save_path", }, **kwargs, ) @@ -68,7 +68,7 @@ def __init__( def create_dataset(self, *args, **kwargs): """Create datasets using the configured helper function. - The callable is resolved from ``train_config.create_dataset.func`` and + The callable is resolved from ``training_config.create_dataset.func`` and invoked with the remaining configuration values. Raises: @@ -77,10 +77,10 @@ def create_dataset(self, *args, **kwargs): self._reject_stage_args("create_dataset", args, kwargs) logger.info("ASRSystem.create_dataset(): starting dataset creation process") start = time.perf_counter() - config = getattr(self.train_config, "create_dataset", None) + config = getattr(self.training_config, "create_dataset", None) if config is None or not getattr(config, "func", None): raise RuntimeError( - "train_config.create_dataset.func must be set to run create_dataset" + "training_config.create_dataset.func must be set to run create_dataset" ) fn = load_function(config.func) extra = {k: v for k, v in config.items() if k != "func"} @@ -100,28 +100,24 @@ def train(self, *args, **kwargs): training before delegating to the base training routine. Raises: - RuntimeError: If ``train_config.dataset_dir`` is not set. + RuntimeError: If ``training_config.dataset_dir`` is not set. """ self._reject_stage_args("train", args, kwargs) logger.info("ASRSystem.train(): starting training process") - dataset_dir = getattr(self.train_config, "dataset_dir", None) + dataset_dir = getattr(self.training_config, "dataset_dir", None) if dataset_dir is None: - raise RuntimeError("train_config.dataset_dir must be set for training.") + raise RuntimeError("training_config.dataset_dir must be set for training.") # Train tokenizer if not trained previously - tokenizer_path = ( - Path(self.train_config.tokenizer.save_path) - / f"{self.train_config.tokenizer.model_type}.model" - ) - if not tokenizer_path.exists(): + if not self._has_tokenizer(): self.train_tokenizer() # Proceed with standard training return super().train() def _has_tokenizer(self) -> bool: - tokenizer_config = self.train_config.tokenizer + tokenizer_config = self.training_config.tokenizer output_path = Path(tokenizer_config.save_path) model = output_path / f"{tokenizer_config.model_type}.model" vocab = output_path / f"{tokenizer_config.model_type}.vocab" @@ -130,7 +126,7 @@ def _has_tokenizer(self) -> bool: def train_tokenizer(self, *args, **kwargs): """Train a SentencePiece tokenizer based on configured text. - The text builder configured in ``train_config.tokenizer.text_builder`` + The text builder configured in ``training_config.tokenizer.text_builder`` is used to generate training text, which is then saved and consumed by the SentencePiece trainer. @@ -142,9 +138,8 @@ def train_tokenizer(self, *args, **kwargs): if self._has_tokenizer(): logger.info("Tokenizer already exists. Skipping train_tokenizer().") return - start = time.perf_counter() - tokenizer_config = getattr(self.train_config, "tokenizer", None) + tokenizer_config = getattr(self.training_config, "tokenizer", None) builder_config = ( getattr(tokenizer_config, "text_builder", None) if tokenizer_config @@ -152,7 +147,7 @@ def train_tokenizer(self, *args, **kwargs): ) if builder_config is None or not getattr(builder_config, "func", None): raise RuntimeError( - "train_config.tokenizer.text_builder.func must be set to build " + "training_config.tokenizer.text_builder.func must be set to build " "tokenizer text." ) builder = load_function(builder_config.func) @@ -177,7 +172,7 @@ def train_tokenizer(self, *args, **kwargs): raise RuntimeError( "Tokenizer text_builder returned no text. Check dataset preparation." ) - output_path = Path(self.train_config.tokenizer.save_path) + output_path = Path(self.training_config.tokenizer.save_path) output_path.mkdir(parents=True, exist_ok=True) train_text_path = getattr(tokenizer_config, "train_file", None) if train_text_path: @@ -193,15 +188,15 @@ def train_tokenizer(self, *args, **kwargs): with open(train_text_path, "w", encoding="utf-8") as f: f.write("\n".join(texts)) - logger.info(f"Training tokenizer: {self.train_config.tokenizer.model_type}") - logger.info(f"Tokenizer output: {self.train_config.tokenizer.save_path}") + logger.info(f"Training tokenizer: {self.training_config.tokenizer.model_type}") + logger.info(f"Tokenizer output: {self.training_config.tokenizer.save_path}") # Example placeholder: train_sentencepiece( train_text_path, output_path, - self.train_config.tokenizer.vocab_size, - model_type=self.train_config.tokenizer.model_type, + self.training_config.tokenizer.vocab_size, + model_type=self.training_config.tokenizer.model_type, ) logger.info( "Tokenizer training completed in %.2fs", time.perf_counter() - start diff --git a/espnet3/systems/asr/transducer_task.py b/espnet3/systems/asr/transducer_task.py new file mode 100644 index 000000000000..1feca7e656ad --- /dev/null +++ b/espnet3/systems/asr/transducer_task.py @@ -0,0 +1,428 @@ +"""ASR Transducer Task.""" + +import argparse +import logging +from typing import Callable, Collection, Dict, List, Optional, Tuple + +import numpy as np +import torch +from typeguard import typechecked + +from espnet2.asr.frontend.abs_frontend import AbsFrontend +from espnet2.asr.frontend.default import DefaultFrontend +from espnet2.asr.frontend.windowing import SlidingWindow +from espnet2.asr.specaug.abs_specaug import AbsSpecAug +from espnet2.asr.specaug.specaug import SpecAug +from espnet2.asr_transducer.decoder.abs_decoder import AbsDecoder +from espnet2.asr_transducer.decoder.mega_decoder import MEGADecoder +from espnet2.asr_transducer.decoder.rnn_decoder import RNNDecoder +from espnet2.asr_transducer.decoder.rwkv_decoder import RWKVDecoder +from espnet2.asr_transducer.decoder.stateless_decoder import StatelessDecoder +from espnet2.asr_transducer.encoder.encoder import Encoder +from espnet2.asr_transducer.espnet_transducer_model import ESPnetASRTransducerModel +from espnet2.asr_transducer.joint_network import JointNetwork +from espnet2.layers.abs_normalize import AbsNormalize +from espnet2.layers.global_mvn import GlobalMVN +from espnet2.layers.utterance_mvn import UtteranceMVN +from espnet2.tasks.abs_task import AbsTask +from espnet2.text.phoneme_tokenizer import g2p_choices +from espnet2.train.class_choices import ClassChoices +from espnet2.train.collate_fn import CommonCollateFn +from espnet2.train.preprocessor import CommonPreprocessor +from espnet2.train.trainer import Trainer +from espnet2.utils.get_default_kwargs import get_default_kwargs +from espnet2.utils.nested_dict_action import NestedDictAction +from espnet2.utils.types import float_or_none, int_or_none, str2bool, str_or_none + +frontend_choices = ClassChoices( + name="frontend", + classes=dict( + default=DefaultFrontend, + sliding_window=SlidingWindow, + ), + type_check=AbsFrontend, + default="default", +) +specaug_choices = ClassChoices( + "specaug", + classes=dict( + specaug=SpecAug, + ), + type_check=AbsSpecAug, + default=None, + optional=True, +) +normalize_choices = ClassChoices( + "normalize", + classes=dict( + global_mvn=GlobalMVN, + utterance_mvn=UtteranceMVN, + ), + type_check=AbsNormalize, + default="utterance_mvn", + optional=True, +) +decoder_choices = ClassChoices( + "decoder", + classes=dict( + mega=MEGADecoder, + rnn=RNNDecoder, + rwkv=RWKVDecoder, + stateless=StatelessDecoder, + ), + type_check=AbsDecoder, + default="rnn", +) + + +class ASRTransducerTask(AbsTask): + """ASR Transducer Task definition.""" + + num_optimizers: int = 1 + + class_choices_list = [ + frontend_choices, + specaug_choices, + normalize_choices, + decoder_choices, + ] + + trainer = Trainer + + @classmethod + def add_task_arguments(cls, parser: argparse.ArgumentParser): + """Add Transducer task arguments. + + Args: + cls: ASRTransducerTask object. + parser: Transducer arguments parser. + + """ + group = parser.add_argument_group(description="Task related.") + + required = parser.get_default("required") + required += ["token_list"] + + group.add_argument( + "--token_list", + type=str_or_none, + default=None, + help="Integer-string mapper for tokens.", + ) + group.add_argument( + "--input_size", + type=int_or_none, + default=None, + help="The number of dimensions for input features.", + ) + group.add_argument( + "--init", + type=str_or_none, + default=None, + help="Type of model initialization to use.", + ) + group.add_argument( + "--model_conf", + action=NestedDictAction, + default=get_default_kwargs(ESPnetASRTransducerModel), + help="The keyword arguments for the model class.", + ) + group.add_argument( + "--encoder_conf", + action=NestedDictAction, + default={}, + help="The keyword arguments for the encoder class.", + ) + group.add_argument( + "--joint_network_conf", + action=NestedDictAction, + default={}, + help="The keyword arguments for the joint network class.", + ) + + group = parser.add_argument_group(description="Preprocess related.") + + group.add_argument( + "--use_preprocessor", + type=str2bool, + default=True, + help="Whether to apply preprocessing to input data.", + ) + group.add_argument( + "--token_type", + type=str, + default="bpe", + choices=["bpe", "char", "word", "phn"], + help="The type of tokens to use during tokenization.", + ) + group.add_argument( + "--bpemodel", + type=str_or_none, + default=None, + help="The path of the sentencepiece model.", + ) + group.add_argument( + "--non_linguistic_symbols", + type=str_or_none, + help="The 'non_linguistic_symbols' file path.", + ) + group.add_argument( + "--cleaner", + type=str_or_none, + choices=[None, "tacotron", "jaconv", "vietnamese"], + default=None, + help="Text cleaner to use.", + ) + group.add_argument( + "--g2p", + type=str_or_none, + choices=g2p_choices, + default=None, + help="g2p method to use if --token_type=phn.", + ) + group.add_argument( + "--speech_volume_normalize", + type=float_or_none, + default=None, + help="Normalization value for maximum amplitude scaling.", + ) + group.add_argument( + "--rir_scp", + type=str_or_none, + default=None, + help="The RIR SCP file path.", + ) + group.add_argument( + "--rir_apply_prob", + type=float, + default=1.0, + help="The probability of the applied RIR convolution.", + ) + group.add_argument( + "--noise_scp", + type=str_or_none, + default=None, + help="The path of noise SCP file.", + ) + group.add_argument( + "--noise_apply_prob", + type=float, + default=1.0, + help="The probability of the applied noise addition.", + ) + group.add_argument( + "--noise_db_range", + type=str, + default="13_15", + help="The range of the noise decibel level.", + ) + + for class_choices in cls.class_choices_list: + # Append -- and --_conf. + # e.g. --decoder and --decoder_conf + class_choices.add_arguments(group) + + @classmethod + @typechecked + def build_collate_fn(cls, args: argparse.Namespace, train: bool) -> Callable[ + [Collection[Tuple[str, Dict[str, np.ndarray]]]], + Tuple[List[str], Dict[str, torch.Tensor]], + ]: + """Build collate function. + + Args: + cls: ASRTransducerTask object. + args: Task arguments. + train: Training mode. + + Return: + : Callable collate function. + + """ + return CommonCollateFn(float_pad_value=0.0, int_pad_value=-1) + + @classmethod + @typechecked + def build_preprocess_fn( + cls, args: argparse.Namespace, train: bool + ) -> Optional[Callable[[str, Dict[str, np.array]], Dict[str, np.ndarray]]]: + """Build pre-processing function. + + Args: + cls: ASRTransducerTask object. + args: Task arguments. + train: Training mode. + + Return: + : Callable pre-processing function. + + """ + if args.use_preprocessor: + retval = CommonPreprocessor( + train=train, + token_type=args.token_type, + token_list=args.token_list, + bpemodel=args.bpemodel, + non_linguistic_symbols=args.non_linguistic_symbols, + text_cleaner=args.cleaner, + g2p_type=args.g2p, + rir_scp=args.rir_scp if hasattr(args, "rir_scp") else None, + rir_apply_prob=( + args.rir_apply_prob if hasattr(args, "rir_apply_prob") else 1.0 + ), + noise_scp=args.noise_scp if hasattr(args, "noise_scp") else None, + noise_apply_prob=( + args.noise_apply_prob if hasattr(args, "noise_apply_prob") else 1.0 + ), + noise_db_range=( + args.noise_db_range if hasattr(args, "noise_db_range") else "13_15" + ), + speech_volume_normalize=( + args.speech_volume_normalize if hasattr(args, "rir_scp") else None + ), + ) + else: + retval = None + + return retval + + @classmethod + def required_data_names( + cls, train: bool = True, inference: bool = False + ) -> Tuple[str, ...]: + """Determine required data based on the task mode. + + Args: + cls: ASRTransducerTask object. + train: Training mode. + inference: Inference mode. + + Return: + retval: Required task data. + + """ + if not inference: + retval = ("speech", "text") + else: + retval = ("speech",) + + return retval + + @classmethod + def optional_data_names( + cls, train: bool = True, inference: bool = False + ) -> Tuple[str, ...]: + """Specify optional data based on the task mode. + + Args: + cls: ASRTransducerTask object. + train: Training mode. + inference: Inference mode. + + Return: + retval: Optional task data. + + """ + retval = () + + return retval + + @classmethod + @typechecked + def build_model(cls, args: argparse.Namespace) -> ESPnetASRTransducerModel: + """Require data based on the task mode. + + Args: + cls: ASRTransducerTask object. + args: Task arguments. + + Return: + model: ASR Transducer model. + + """ + if isinstance(args.token_list, str): + with open(args.token_list, encoding="utf-8") as f: + token_list = [line.rstrip() for line in f] + + # Overwriting token_list to keep it as "portable". + args.token_list = list(token_list) + elif isinstance(args.token_list, (tuple, list)): + token_list = list(args.token_list) + else: + raise RuntimeError("token_list must be str or list") + vocab_size = len(token_list) + + if hasattr(args, "scheduler_conf"): + args.model_conf["warmup_steps"] = args.scheduler_conf.get( + "warmup_steps", 25000 + ) + + logging.info(f"Vocabulary size: {vocab_size}") + + # 1. frontend + if args.input_size is None: + # Extract features in the model + frontend_class = frontend_choices.get_class(args.frontend) + frontend = frontend_class(**args.frontend_conf) + input_size = frontend.output_size() + else: + # Give features from data-loader + frontend = None + input_size = args.input_size + + # 2. Data augmentation for spectrogram + if args.specaug is not None: + specaug_class = specaug_choices.get_class(args.specaug) + specaug = specaug_class(**args.specaug_conf) + else: + specaug = None + + # 3. Normalization layer + if args.normalize is not None: + normalize_class = normalize_choices.get_class(args.normalize) + normalize = normalize_class(**args.normalize_conf) + else: + normalize = None + + # 4. Encoder + encoder = Encoder(input_size, **args.encoder_conf) + encoder_output_size = encoder.output_size + + # 5. Decoder + decoder_class = decoder_choices.get_class(args.decoder) + + decoder = decoder_class( + vocab_size, + **args.decoder_conf, + ) + decoder_output_size = decoder.output_size + + # 6. Joint Network + joint_network = JointNetwork( + vocab_size, + encoder_output_size, + decoder_output_size, + **args.joint_network_conf, + ) + + # 7. Build model + model = ESPnetASRTransducerModel( + vocab_size=vocab_size, + token_list=token_list, + frontend=frontend, + specaug=specaug, + normalize=normalize, + encoder=encoder, + decoder=decoder, + joint_network=joint_network, + **args.model_conf, + ) + + # 8. Initialize model + if args.init is not None: + raise NotImplementedError( + "Currently not supported.", + "Initialization part will be reworked in a short future.", + ) + + return model diff --git a/espnet3/systems/base/inference.py b/espnet3/systems/base/inference.py index 997ce5821154..6f98fb2a23df 100644 --- a/espnet3/systems/base/inference.py +++ b/espnet3/systems/base/inference.py @@ -90,17 +90,15 @@ def infer(config: DictConfig): config.test_set = test_name output_fn_path = getattr(config, "output_fn", None) - if not output_fn_path: - raise RuntimeError("infer_config.output_fn must be set.") - - _load_output_fn(output_fn_path) + if output_fn_path: + _load_output_fn(output_fn_path) input_key = getattr(config, "input_key", None) if input_key is None: - raise RuntimeError("infer_config.input_key must be set.") + raise RuntimeError("inference_config.input_key must be set.") if isinstance(input_key, (list, tuple)) and not input_key: - raise RuntimeError("infer_config.input_key must not be empty.") + raise RuntimeError("inference_config.input_key must not be empty.") output_keys = getattr(config, "output_keys", None) if output_keys is not None: @@ -109,14 +107,14 @@ def infer(config: DictConfig): elif not isinstance(output_keys, (list, tuple)): output_keys = list(output_keys) if not output_keys: - raise RuntimeError("infer_config.output_keys must not be empty.") + raise RuntimeError("inference_config.output_keys must not be empty.") - idx_key = getattr(config, "idx_key", "uttid") + idx_key = getattr(config, "idx_key", "utt_id") batch_size = getattr(config, "batch_size", None) provider_config = getattr(config, "provider", None) if provider_config is None: - raise RuntimeError("infer_config.provider must be set.") + raise RuntimeError("inference_config.provider must be set.") raw_params = getattr(provider_config, "params", {}) or {} if OmegaConf.is_config(raw_params): provider_params = OmegaConf.to_container(raw_params, resolve=True) @@ -124,11 +122,12 @@ def infer(config: DictConfig): provider_params = dict(raw_params) provider_params["input_key"] = input_key - provider_params["output_fn_path"] = output_fn_path + if output_fn_path: + provider_params["output_fn_path"] = output_fn_path provider = instantiate( provider_config, - infer_config=config, + inference_config=config, params=provider_params, _recursive_=False, ) @@ -136,13 +135,14 @@ def infer(config: DictConfig): hyp_keys = output_keys if output_keys is not None else [] runner_config = getattr(config, "runner", None) if runner_config is None: - raise RuntimeError("infer_config.runner must be set.") + raise RuntimeError("inference_config.runner must be set.") runner_kwargs = { "provider": provider, "async_mode": False, "idx_key": idx_key, "hyp_key": hyp_keys, + "ref_key": [], "batch_size": batch_size, } runner = instantiate(runner_config, **runner_kwargs) @@ -150,25 +150,32 @@ def infer(config: DictConfig): raise TypeError( f"{type(runner).__name__} must provide inference runner attributes" ) - dataset_length = len(provider.build_dataset(config)) + dataset = provider.build_dataset(config) + dataset_length = len(dataset) + if dataset_length == 0: + raise RuntimeError( + f"Test dataset '{test_name}' is empty. " + "Please check the dataset manifest and preprocessing outputs." + ) logger.info("===> Processing %d samples..", dataset_length) out = runner(list(range(dataset_length))) if out is None: raise RuntimeError("Async inference is not supported in this entrypoint.") # Runner can return nested lists. normalize to flat list. results = _flatten_results(out) + if not results: + raise RuntimeError("No inference results available.") + first = results[0] + resolved_idx_key = runner.resolve_idx_key(first) if output_keys is None: - if not results: - raise RuntimeError("No inference results available.") - first = results[0] - output_keys = [key for key in first.keys() if key != runner.idx_key] + output_keys = [key for key in first.keys() if key != resolved_idx_key] if not output_keys: raise RuntimeError("No output keys found in inference results.") - # Convert output dicts into per-key SCP lines (uttid + value). + # Convert output dicts into per-key SCP lines (utt_id + value). scp_lines = _collect_scp_lines( results, - idx_key=runner.idx_key, + idx_key=resolved_idx_key, hyp_keys=output_keys, ref_keys=[], ) diff --git a/espnet3/systems/base/inference_provider.py b/espnet3/systems/base/inference_provider.py index 82a6e84e2dfe..5a45776de29e 100644 --- a/espnet3/systems/base/inference_provider.py +++ b/espnet3/systems/base/inference_provider.py @@ -7,7 +7,7 @@ import torch from hydra.utils import instantiate -from omegaconf import DictConfig +from omegaconf import DictConfig, OmegaConf from espnet3.parallel.env_provider import EnvironmentProvider from espnet3.utils.logging_utils import log_instance_dict @@ -30,14 +30,14 @@ class InferenceProvider(EnvironmentProvider, ABC): pickle-safe for Dask. Args: - infer_config (DictConfig): Hydra configuration used to build dataset/model. + inference_config (DictConfig): Hydra configuration used to build dataset/model. params (Dict[str, Any] | None): Optional additional key-value pairs that will be merged into the returned environment (e.g., device, tokenizer, beam size). Notes: - Subclasses must implement ``build_dataset`` and ``build_model``. - - ``self.infer_config.update(self.params)`` allows lightweight overrides + - ``self.inference_config.update(self.params)`` allows lightweight overrides (e.g., runtime overrides) but avoid mutating deep structures unless intended. """ @@ -45,10 +45,10 @@ class InferenceProvider(EnvironmentProvider, ABC): # TODO(Masao) Add detailed description on Runner/Provider in the document. def __init__( - self, infer_config: DictConfig = None, params: Dict[str, Any] | None = None + self, inference_config: DictConfig = None, params: Dict[str, Any] | None = None ): """Initialize InferenceProvider object.""" - super().__init__(infer_config) + super().__init__(inference_config) self.params = params or {} self.config.update(self.params) @@ -157,7 +157,13 @@ def build_dataset(config: DictConfig): - Rely on fields already present in ``config`` instead of reading global state whenever possible. """ - organizer = instantiate(config.dataset) + if isinstance(config, dict): + config = OmegaConf.create(config) + organizer = instantiate(config.dataset) + + if isinstance(config, DictConfig): + organizer = instantiate(config.dataset) + test_set = config.test_set logger.info("Building dataset for test set: %s", test_set) return organizer.test[test_set] @@ -193,6 +199,9 @@ def build_model(config: DictConfig): - Do not perform training/optimization here, this is for inference setup only. """ + if isinstance(config, dict): + config = OmegaConf.create(config) + device = "cuda" if torch.cuda.is_available() else "cpu" if device == "cuda": device_id = os.getenv("CUDA_VISIBLE_DEVICES", "0").split(",")[0].strip() diff --git a/espnet3/systems/base/inference_runner.py b/espnet3/systems/base/inference_runner.py index df7dea4394f4..74818543a34d 100644 --- a/espnet3/systems/base/inference_runner.py +++ b/espnet3/systems/base/inference_runner.py @@ -19,11 +19,14 @@ class InferenceRunner(BaseRunner): function. The key names are configurable via ``idx_key`` and ``hyp_key``/``ref_key``. ``hyp_key`` and ``ref_key`` may be a single string or a list of strings to support multiple hypothesis/reference - fields. + fields. ``idx_key`` is the key used to map each inference result to + its source dataset index when writing SCP files. Output format requirements: - - The result is a dict with exactly the configured keys. - - ``idx_key`` value is a scalar (list/tuple is not allowed). + - The result is a dict with the configured keys plus any extra fields. + - A sample identifier key must exist under ``idx_key`` so SCP outputs + can map each result back to the corresponding dataset sample. + - The sample identifier must be a single value, not a list or tuple. - ``hyp_key`` and ``ref_key`` values may be scalars or lists/tuples. If lists are returned, each entry is written to its own SCP file (e.g., ``hyp0.scp``, ``hyp1.scp``). @@ -32,12 +35,22 @@ class InferenceRunner(BaseRunner): def __init__( self, provider: EnvironmentProvider, - idx_key: str = "idx", + idx_key: str = "utt_id", hyp_key: str | Sequence[str] = "hyp", ref_key: str | Sequence[str] = "ref", **kwargs, ) -> None: - """Initialize the inference runner with output key settings.""" + """Initialize the inference runner with output key settings. + + Args: + provider: Environment provider that supplies dataset/model/env. + idx_key: Output dict key used as the sample identifier written in + the first column of each SCP line. This ties each inference + result back to its dataset sample. Defaults to ``"utt_id"``. + hyp_key: Hypothesis key or keys expected in the output dict. + ref_key: Reference key or keys expected in the output dict. + **kwargs: Forwarded to ``BaseRunner``. + """ super().__init__(provider, **kwargs) self.idx_key = idx_key self.hyp_key = ( @@ -47,6 +60,16 @@ def __init__( list(ref_key) if isinstance(ref_key, (list, tuple, ListConfig)) else ref_key ) + def resolve_idx_key(self, output: Dict[str, Any]) -> str: + """Validate that the configured sample-identifier key exists in output.""" + if self.idx_key not in output: + raise ValueError( + "Inference output must include the configured sample identifier " + "key used to map SCP results back to dataset samples. " + f"idx_key={self.idx_key!r}" + ) + return self.idx_key + def _validate_output(self, output: Dict[str, Any]) -> None: if not isinstance(output, dict): raise TypeError( @@ -63,7 +86,8 @@ def _validate_output(self, output: Dict[str, Any]) -> None: if isinstance(self.ref_key, (list, tuple)) else [self.ref_key] ) - expected = {self.idx_key, *hyp_keys, *ref_keys} + idx_key = self.resolve_idx_key(output) + expected = {idx_key, *hyp_keys, *ref_keys} actual = set(output.keys()) missing = expected - actual if missing: @@ -72,10 +96,10 @@ def _validate_output(self, output: Dict[str, Any]) -> None: f"missing={sorted(missing)}" ) - idx_value = output[self.idx_key] + idx_value = output[idx_key] if isinstance(idx_value, (list, tuple)): raise TypeError( - f"'{self.idx_key}' must be a scalar, got {type(idx_value).__name__}" + f"'{idx_key}' must be a single value, not {type(idx_value).__name__}" ) @staticmethod @@ -86,14 +110,14 @@ def forward(idx, dataset=None, model=None, **kwargs): idx: Integer index or an iterable of integer indices into the dataset. dataset: Dataset providing inference entries. model: Inference model callable on the configured input. - **kwargs: Expects ``input_key`` and ``output_fn_path``. + **kwargs: Expects ``input_key`` and optionally ``output_fn_path``. Returns: Dict containing ``idx`` and output fields for a single item, or a list of dicts for batched inputs (as returned by ``output_fn``). Raises: - RuntimeError: If required I/O settings are missing. + RuntimeError: If required input settings are missing. KeyError: If required input keys are missing from the dataset item(s). RuntimeError: If batched inference fails; includes guidance to disable batching when unsupported. @@ -119,9 +143,7 @@ def forward(idx, dataset=None, model=None, **kwargs): raise RuntimeError("input_key must be provided for inference.") input_key = kwargs["input_key"] output_fn_path = kwargs.get("output_fn_path") - if not output_fn_path: - raise RuntimeError("output_fn_path must be provided for inference.") - output_fn = _load_output_fn(output_fn_path) + output_fn = _load_output_fn(output_fn_path) if output_fn_path else None keys = ( list(input_key) @@ -138,6 +160,8 @@ def forward(idx, dataset=None, model=None, **kwargs): raise KeyError(f"Input key '{key}' not found in dataset item.") inputs_dict[key] = data[key] model_output = model(**inputs_dict) + if output_fn is None: + return model_output return output_fn(data=data, model_output=model_output, idx=idx) indices = list(idx) @@ -151,6 +175,8 @@ def forward(idx, dataset=None, model=None, **kwargs): try: model_output = model(**inputs_dict) + if output_fn is None: + return model_output return output_fn(data=data_batch, model_output=model_output, idx=indices) except Exception as exc: # noqa: BLE001 raise RuntimeError( diff --git a/espnet3/systems/base/metric.py b/espnet3/systems/base/metric.py index 8465a7c057e5..b91b4528f556 100644 --- a/espnet3/systems/base/metric.py +++ b/espnet3/systems/base/metric.py @@ -14,20 +14,76 @@ logger = logging.getLogger(__name__) -def measure(measure_config: DictConfig): +def _resolve_test_sets(metrics_config: DictConfig) -> list[str]: + """Resolve metric target test sets from config or inference outputs. + + This helper decides which test-set names should be scored by the + measurement stage. + + Resolution order: + 1. If ``metrics_config.dataset.test`` is defined, use its ``name`` + fields as-is. + 2. Otherwise, scan ``metrics_config.inference_dir`` and treat every + subdirectory as a test set. + + For example, if ``inference_dir`` contains: + + .. code-block:: text + + exp/my_run/inference/ + test-clean/ + test-other/ + + then this function returns ``["test-clean", "test-other"]``. + + Args: + metrics_config: Measurement config containing either + ``dataset.test`` entries or an ``inference_dir`` path. + + Returns: + Sorted list of test-set names to score. + + Raises: + ValueError: If neither ``dataset.test`` nor any subdirectories under + ``inference_dir`` are available. + """ + dataset = getattr(metrics_config, "dataset", None) + test_config = getattr(dataset, "test", None) if dataset is not None else None + if test_config: + return [t.name for t in test_config] + + inference_dir = Path(metrics_config.inference_dir) + test_sets = sorted( + entry.name + for entry in inference_dir.iterdir() + if entry.is_dir() and not entry.name.startswith(".") + ) + if not test_sets: + raise ValueError( + "No test sets found. Specify `metrics_config.dataset.test` or place " + f"test-set subdirectories under inference_dir: {inference_dir}" + ) + logger.info( + "Resolved test sets from inference_dir: %s", + test_sets, + ) + return test_sets + + +def measure(metrics_config: DictConfig): """Compute metrics for each test set and write a metrics JSON file. Args: - measure_config: Omegaconf configuration with inference and metric settings. + metrics_config: Omegaconf configuration with inference and metric settings. Returns: Nested dict keyed by metric class path and test set name. """ - test_sets = [t.name for t in measure_config.dataset.test] + test_sets = _resolve_test_sets(metrics_config) results = {} - assert hasattr(measure_config, "metrics"), "Please specify `metrics`!" + assert hasattr(metrics_config, "metrics"), "Please specify `metrics`!" - for idx, metric_config in enumerate(measure_config.metrics): + for idx, metric_config in enumerate(metrics_config.metrics): metric = instantiate(metric_config.metric) if not isinstance(metric, AbsMetric): raise TypeError(f"{type(metric)} is not a valid AbsMetric instance") @@ -52,15 +108,15 @@ def measure(measure_config: DictConfig): ) inputs = [ref_key, hyp_key] data = load_scp_fields( - inference_dir=Path(measure_config.inference_dir), + inference_dir=Path(metrics_config.inference_dir), test_name=test_name, inputs=inputs, file_suffix=".scp", ) - metric_result = metric(data, test_name, measure_config.inference_dir) + metric_result = metric(data, test_name, metrics_config.inference_dir) results[get_class_path(metric)].update({test_name: metric_result}) - out_path = Path(measure_config.inference_dir) / "metrics.json" + out_path = Path(metrics_config.inference_dir) / "metrics.json" with open(out_path, "w", encoding="utf-8") as f: json.dump(results, f, indent=2, ensure_ascii=False) diff --git a/espnet3/systems/base/system.py b/espnet3/systems/base/system.py index 18f4ffc1a046..5ca0e73bee1c 100644 --- a/espnet3/systems/base/system.py +++ b/espnet3/systems/base/system.py @@ -30,39 +30,39 @@ class BaseSystem: All behavior is config-driven. Args: - train_config (DictConfig | None): Training configuration. - infer_config (DictConfig | None): Inference configuration. - measure_config (DictConfig | None): Measurement configuration. + training_config (DictConfig | None): Training configuration. + inference_config (DictConfig | None): Inference configuration. + metrics_config (DictConfig | None): Measurement configuration. stage_log_mapping (dict | None): Optional overrides for stage log path resolution. Keys are stage names; values are dotted attribute - paths (e.g., ``"train_config.exp_dir"``) or lists/tuples of such + paths (e.g., ``"training_config.exp_dir"``) or lists/tuples of such paths (first non-empty value wins). Stage log mapping (base defaults): | Stage | Path reference | |--- |--- | - | create_dataset | train_config.recipe_dir | - | collect_stats | train_config.stats_dir | - | train | train_config.exp_dir | - | infer | infer_config.inference_dir | - | measure | measure_config.inference_dir | - | pack_model | train_config.exp_dir | - | upload_model | train_config.exp_dir | + | create_dataset | training_config.recipe_dir | + | collect_stats | training_config.stats_dir | + | train | training_config.exp_dir | + | infer | inference_config.inference_dir | + | measure | metrics_config.inference_dir | + | pack_model | training_config.exp_dir | + | upload_model | training_config.exp_dir | Any stage missing from the mapping (or resolving to ``None``) falls back - to the default log directory: ``train_config.exp_dir`` when available, + to the default log directory: ``training_config.exp_dir`` when available, otherwise ``/logs``. Examples: Override a subset of stage log paths: ```python system = BaseSystem( - train_config=train_cfg, - infer_config=infer_cfg, - measure_config=measure_cfg, + training_config=train_cfg, + inference_config=infer_cfg, + metrics_config=measure_cfg, stage_log_mapping={ - "infer": "train_config.exp_dir", - "measure": "train_config.exp_dir", + "infer": "training_config.exp_dir", + "measure": "training_config.exp_dir", }, ) ``` @@ -70,18 +70,18 @@ class BaseSystem: def __init__( self, - train_config: DictConfig | None = None, - infer_config: DictConfig | None = None, - measure_config: DictConfig | None = None, + training_config: DictConfig | None = None, + inference_config: DictConfig | None = None, + metrics_config: DictConfig | None = None, stage_log_mapping: dict | None = None, ) -> None: """Initialize the system with optional stage configs.""" - self.train_config = train_config - self.infer_config = infer_config - self.measure_config = measure_config + self.training_config = training_config + self.inference_config = inference_config + self.metrics_config = metrics_config - if train_config is not None: - self.exp_dir = Path(train_config.exp_dir) + if training_config is not None: + self.exp_dir = Path(training_config.exp_dir) self.exp_dir.mkdir(parents=True, exist_ok=True) else: self.exp_dir = None @@ -92,13 +92,13 @@ def __init__( default_dir = Path.cwd() / "logs" base_mapping = { - "create_dataset": "train_config.recipe_dir", - "collect_stats": "train_config.stats_dir", - "train": "train_config.exp_dir", - "infer": "infer_config.inference_dir", - "measure": "measure_config.inference_dir", - "pack_model": "train_config.exp_dir", - "upload_model": "train_config.exp_dir", + "create_dataset": "training_config.recipe_dir", + "collect_stats": "training_config.stats_dir", + "train": "training_config.exp_dir", + "infer": "inference_config.inference_dir", + "measure": "metrics_config.inference_dir", + "pack_model": "training_config.exp_dir", + "upload_model": "training_config.exp_dir", } mapping = dict(base_mapping) if stage_log_mapping: @@ -112,12 +112,12 @@ def __init__( self.stage_log_dirs[stage] = Path(resolved) logger.info( - "Initialized %s with train_config=%s infer_config=%s " - "measure_config=%s exp_dir=%s", + "Initialized %s with training_config=%s inference_config=%s " + "metrics_config=%s exp_dir=%s", self.__class__.__name__, - train_config is not None, - infer_config is not None, - measure_config is not None, + training_config is not None, + inference_config is not None, + metrics_config is not None, self.exp_dir, ) @@ -126,7 +126,7 @@ def _resolve_stage_log_ref( ) -> str | None: """Resolve stage log mapping references to concrete values. - Supports dotted attribute paths like ``train_config.exp_dir`` and + Supports dotted attribute paths like ``training_config.exp_dir`` and fallbacks via list/tuple entries (first non-empty value wins). """ target = self @@ -170,43 +170,43 @@ def collect_stats(self, *args, **kwargs): self._reject_stage_args("collect_stats", args, kwargs) logger.info( "Collecting stats | exp_dir=%s stats_dir=%s", - getattr(self.train_config, "exp_dir", None), - getattr(self.train_config, "stats_dir", None), + getattr(self.training_config, "exp_dir", None), + getattr(self.training_config, "stats_dir", None), ) - return collect_stats(self.train_config) + return collect_stats(self.training_config) def train(self, *args, **kwargs): """Train the system model.""" self._reject_stage_args("train", args, kwargs) model_target = None - if self.train_config is not None and hasattr(self.train_config, "model"): - model_config = self.train_config.model + if self.training_config is not None and hasattr(self.training_config, "model"): + model_config = self.training_config.model if isinstance(model_config, DictConfig): model_target = model_config.get("_target_") logger.info( "Training start | exp_dir=%s model=%s", - getattr(self.train_config, "exp_dir", None), + getattr(self.training_config, "exp_dir", None), model_target or "", ) - return train(self.train_config) + return train(self.training_config) def infer(self, *args, **kwargs): """Run inference on the configured datasets.""" self._reject_stage_args("infer", args, kwargs) logger.info( "Inference start | inference_dir=%s", - getattr(self.infer_config, "inference_dir", None), + getattr(self.inference_config, "inference_dir", None), ) - return infer(self.infer_config) + return infer(self.inference_config) def measure(self, *args, **kwargs): """Compute evaluation metrics from hypothesis/reference outputs.""" self._reject_stage_args("measure", args, kwargs) logger.info( - "Metrics start | measure_config=%s", - self.measure_config is not None, + "Metrics start | metrics_config=%s", + self.metrics_config is not None, ) - result = measure(self.measure_config) + result = measure(self.metrics_config) logger.info("results: %s", result) return result diff --git a/espnet3/utils/config_utils.py b/espnet3/utils/config_utils.py index 993149164253..3f789c9526fa 100644 --- a/espnet3/utils/config_utils.py +++ b/espnet3/utils/config_utils.py @@ -1,10 +1,16 @@ """Configuration helpers and OmegaConf resolvers for ESPnet3.""" import logging +import re +from importlib import resources from pathlib import Path from omegaconf import DictConfig, ListConfig, OmegaConf +_RELATIVE_RESOLVER_PATTERN = re.compile( + r"\$\{(?Pload_line|load_yaml):(?P[^,}]+)" +) + def load_line(path): """Load lines from a text file and return as a list of strings. @@ -32,15 +38,52 @@ def load_line(path): raise +def load_yaml(path, key=None): + """Load a YAML file and optionally return a nested key. + + This resolver is intended for pulling a single value from another config, + e.g., `${load_yaml:conf/train.yaml,exp_tag}`. + + Args: + path (str or Path): Path to the YAML file. + key (str, optional): Dot-delimited key to select. + + Returns: + Any: The selected value or the full config if key is None. + """ + cfg_path = Path(path) + try: + cfg = OmegaConf.load(cfg_path) + except FileNotFoundError: + logging.error(f"File not found: {path}") + raise + except PermissionError: + logging.error(f"Permission denied when accessing file: {path}") + raise + + if key is None or str(key).strip() == "": + return cfg + + value = OmegaConf.select(cfg, str(key), default=None) + if value is None: + raise KeyError(f"Key not found in YAML: {key}") + return value + + OMEGACONF_ESPNET3_RESOLVER = { "load_line": load_line, + "load_yaml": load_yaml, } for name, resolver in OMEGACONF_ESPNET3_RESOLVER.items(): OmegaConf.register_new_resolver(name, resolver) logging.info(f"Registered ESPnet-3 OmegaConf Resolver: {name}") -def _process_dict_config_entry(entry: DictConfig, base_path: Path) -> list: +def _process_dict_config_entry( + entry: DictConfig, + base_path: Path, + resolve: bool = True, +) -> list: results = [] for key, val in entry.items(): # Skip null entries (can be used for disabling config entries) @@ -51,15 +94,17 @@ def _process_dict_config_entry(entry: DictConfig, base_path: Path) -> list: # or use as-is if path includes '/' composed = f"{key}/{val}" if "/" not in val else val config_path = _build_config_path(base_path, composed) - results.append({key: load_config_with_defaults(str(config_path))}) + results.append( + {key: load_config_with_defaults(str(config_path), resolve=resolve)} + ) return results -def load_config_with_defaults(path: str) -> OmegaConf: +def load_config_with_defaults(path: str, resolve: bool = True) -> OmegaConf: """Load an YAML config with recursive `_self_` merging via Hydra-style `defaults`. This function recursively loads and merges dependent YAML files specified - in the `defaults` key of the given config. It mimics Hydra’s composition mechanism + in the `defaults` key of the given config. It mimics Hydra's composition mechanism without using Hydra's runtime, which makes it suitable for standalone YAML handling (e.g., for distributed training or script-based training setups). @@ -82,12 +127,14 @@ def load_config_with_defaults(path: str) -> OmegaConf: Args: path (str): Path to the main YAML config file. + resolve (bool): Whether to resolve interpolations after merging. Returns: OmegaConf.DictConfig: Fully resolved and merged configuration object. """ base_path = Path(path).parent main_config = OmegaConf.load(path) + _normalize_relative_resolver_paths(main_config, base_path) config_self = main_config.copy() if "defaults" not in main_config: @@ -103,10 +150,14 @@ def load_config_with_defaults(path: str) -> OmegaConf: self_merged = True else: config_path = _build_config_path(base_path, entry) - merged_configs.append(load_config_with_defaults(str(config_path))) + merged_configs.append( + load_config_with_defaults(str(config_path), resolve=resolve) + ) elif isinstance(entry, DictConfig): - merged_configs.extend(_process_dict_config_entry(entry, base_path)) + merged_configs.extend( + _process_dict_config_entry(entry, base_path, resolve=resolve) + ) elif entry == "_self_": merged_configs.append(config_self) @@ -116,14 +167,121 @@ def load_config_with_defaults(path: str) -> OmegaConf: merged_configs.append(config_self) final_config = OmegaConf.merge(*merged_configs) - OmegaConf.resolve(final_config) + + if resolve: + OmegaConf.resolve(final_config) if "defaults" in final_config: del final_config["defaults"] + _ensure_target_convert_all(final_config) return final_config +def load_template_defaults( + template_config_path: str, + template_package: str, +): + """Load packaged TEMPLATE defaults by config kind.""" + resource = resources.files(template_package).joinpath( + *Path(template_config_path).parts + ) + with resources.as_file(resource) as path: + return load_config_with_defaults(str(path), resolve=False) + + +def load_and_merge_config( + config_path: Path | None, + template_config_path: str, + template_package: str | None = None, +): + """Load user config and merge it with TEMPLATE defaults.""" + if config_path is None: + return None + if template_package is None: + template_package = _infer_template_package_from_config_path(config_path) + if template_package is None: + raise ValueError( + "template_package is required when it cannot be inferred from config_path" + ) + default_cfg = load_template_defaults(template_config_path, template_package) + user_cfg = load_config_with_defaults(str(config_path), resolve=False) + merged_cfg = OmegaConf.merge(default_cfg, user_cfg) + OmegaConf.resolve(merged_cfg) + _ensure_target_convert_all(merged_cfg) + return merged_cfg + + +def _ensure_target_convert_all(cfg) -> None: + if isinstance(cfg, DictConfig): + if "_target_" in cfg: + cfg["_convert_"] = "all" + for value in cfg.values(): + _ensure_target_convert_all(value) + elif isinstance(cfg, ListConfig): + for value in cfg: + _ensure_target_convert_all(value) + + +def _normalize_relative_resolver_paths(cfg, base_path: Path) -> None: + if isinstance(cfg, DictConfig): + for key in list(cfg.keys()): + value = cfg._get_node(key) + if isinstance(value, (DictConfig, ListConfig)): + _normalize_relative_resolver_paths(value, base_path) + else: + raw_value = value._value() + if isinstance(raw_value, str): + cfg[key] = _rewrite_relative_resolver_paths(raw_value, base_path) + elif isinstance(cfg, ListConfig): + for idx in range(len(cfg)): + value = cfg._get_node(idx) + if isinstance(value, (DictConfig, ListConfig)): + _normalize_relative_resolver_paths(value, base_path) + else: + raw_value = value._value() + if isinstance(raw_value, str): + cfg[idx] = _rewrite_relative_resolver_paths(raw_value, base_path) + + +def _rewrite_relative_resolver_paths(value: str, base_path: Path) -> str: + def replace(match: re.Match) -> str: + resolver = match.group("resolver") + raw_path = match.group("path").strip() + + quote = "" + normalized_path = raw_path + if ( + len(raw_path) >= 2 + and raw_path[0] == raw_path[-1] + and raw_path[0] + in ( + "'", + '"', + ) + ): + quote = raw_path[0] + normalized_path = raw_path[1:-1].strip() + + # Leave dynamic or already-absolute paths untouched. + if "${" in normalized_path or Path(normalized_path).is_absolute(): + return match.group(0) + + base_relative_path = base_path / normalized_path + workspace_relative_path = Path(normalized_path) + if normalized_path.startswith(("./", "../")) or base_relative_path.exists(): + resolved_path = base_relative_path.absolute().as_posix() + elif workspace_relative_path.exists(): + resolved_path = workspace_relative_path.absolute().as_posix() + else: + resolved_path = base_relative_path.absolute().as_posix() + if quote: + resolved_path = f"{quote}{resolved_path}{quote}" + return f"${{{resolver}:{resolved_path}" + + return _RELATIVE_RESOLVER_PATTERN.sub(replace, value) + + def _build_config_path(base_path: Path, entry: str) -> Path: entry_path = Path(entry) if entry_path.suffix and entry_path.suffix != ".yaml": @@ -134,3 +292,22 @@ def _build_config_path(base_path: Path, entry: str) -> Path: if not entry_path.suffix: entry += ".yaml" return base_path / entry + + +def _infer_template_package_from_config_path(config_path: Path) -> str | None: + parts = config_path.resolve().parts + try: + egs3_index = parts.index("egs3") + except ValueError: + return None + + conf_index = None + for index in range(egs3_index + 1, len(parts)): + if parts[index] == "conf": + conf_index = index + break + if conf_index is None or conf_index - egs3_index < 3: + return None + + task_name = parts[conf_index - 1] + return f"egs3.TEMPLATE.{task_name}" diff --git a/espnet3/utils/logging_utils.py b/espnet3/utils/logging_utils.py index 797e6a6389db..56d134dfe612 100644 --- a/espnet3/utils/logging_utils.py +++ b/espnet3/utils/logging_utils.py @@ -2,6 +2,7 @@ from __future__ import annotations +import argparse import contextvars import logging import os @@ -14,9 +15,10 @@ from datetime import datetime from pathlib import Path from shutil import which -from typing import Mapping +from typing import Any, Mapping import torch +from omegaconf import OmegaConf LOG_FORMAT = ( "[%(hostname)s] %(asctime)s (%(filename)s:%(lineno)d) " @@ -52,6 +54,96 @@ def log_stage(name: str): _LOG_STAGE.reset(token) +def log_stage_metadata( + logger: logging.Logger, + system: Any, + args: argparse.Namespace | None, +) -> None: + """Write per-stage metadata into the active stage log. + + This helper records the invocation context that is useful when auditing a + stage log after the fact: CLI command line, selected config file paths, + environment metadata, and the resolved in-memory configs attached to the + system object. + + Args: + logger (logging.Logger): Logger that should receive the metadata + entries. In practice this is the stage logger configured immediately + before a stage starts. + system (Any): Instantiated system object. When present, the attributes + ``training_config``, ``inference_config``, and ``metrics_config`` are read + and dumped as resolved YAML for reproducibility. + args (argparse.Namespace | None): Parsed CLI namespace. The function + reads ``training_config``, ``inference_config``, ``metrics_config``, and + ``write_requirements`` from this namespace when available. + + Returns: + None: This function is logging-only and does not return a value. + + Examples: + >>> from argparse import Namespace + >>> class DummySystem: + ... training_config = {"exp_dir": "./exp/train"} + ... inference_config = None + ... metrics_config = None + >>> log_stage_metadata( + ... logging.getLogger("espnet3"), + ... system=DummySystem(), + ... args=Namespace( + ... training_config="conf/training.yaml", + ... inference_config=None, + ... metrics_config=None, + ... write_requirements=False, + ... ), + ... ) + + """ + training_config = getattr(system, "training_config", None) + inference_config = getattr(system, "inference_config", None) + metrics_config = getattr(system, "metrics_config", None) + + log_run_metadata( + logger, + argv=sys.argv, + configs={ + "Training": ( + Path(args.training_config) + if args is not None and getattr(args, "training_config", None) + else None + ), + "Inference": ( + Path(args.inference_config) + if args is not None and getattr(args, "inference_config", None) + else None + ), + "Metrics": ( + Path(args.metrics_config) + if args is not None and getattr(args, "metrics_config", None) + else None + ), + }, + write_requirements=bool( + getattr(args, "write_requirements", False) if args is not None else False + ), + ) + log_env_metadata(logger) + if training_config is not None: + logger.info( + "Training config content:\n%s", + OmegaConf.to_yaml(training_config, resolve=True), + ) + if inference_config is not None: + logger.info( + "Inference config content:\n%s", + OmegaConf.to_yaml(inference_config, resolve=True), + ) + if metrics_config is not None: + logger.info( + "Metrics config content:\n%s", + OmegaConf.to_yaml(metrics_config, resolve=True), + ) + + def set_log_format( log_format: str | None = None, date_format: str | None = None, diff --git a/espnet3/utils/scp_utils.py b/espnet3/utils/scp_utils.py index 501854b7fd55..84d1642aea2c 100644 --- a/espnet3/utils/scp_utils.py +++ b/espnet3/utils/scp_utils.py @@ -8,10 +8,15 @@ def read_scp(scp_file): """Read an SCP file into a key/value dictionary.""" with open(scp_file, "r") as f: lines = [line.strip() for line in f.readlines()] - return { - line.split(maxsplit=1)[0].strip(): line.split(maxsplit=1)[1].strip() - for line in lines - } + parsed = {} + for line in lines: + if not line: + continue + parts = line.split(maxsplit=1) + key = parts[0].strip() + value = parts[1].strip() if len(parts) > 1 else "" + parsed[key] = value + return parsed def get_class_path(obj) -> str: diff --git a/espnet3/utils/stages_utils.py b/espnet3/utils/stages_utils.py index d0558471c7e9..acf4716e2aae 100644 --- a/espnet3/utils/stages_utils.py +++ b/espnet3/utils/stages_utils.py @@ -2,13 +2,18 @@ from __future__ import annotations +import argparse import logging import os import time from pathlib import Path -from typing import Any, Callable, Iterable, List, Sequence +from typing import Any, Iterable, List, Sequence, Tuple -from espnet3.utils.logging_utils import log_stage, set_stage_log_handler +from espnet3.utils.logging_utils import ( + log_stage, + log_stage_metadata, + set_stage_log_handler, +) logger = logging.getLogger(__name__) @@ -40,11 +45,11 @@ def _get_process_rank() -> int: def _get_stage_log_mode(system: Any) -> str: - """Return normalized stage_log_mode from train_config, defaulting to rank0.""" + """Return normalized stage_log_mode from training_config, defaulting to rank0.""" mode = "rank0" - train_config = getattr(system, "train_config", None) - if train_config is not None: - mode = getattr(train_config, "stage_log_mode", mode) + training_config = getattr(system, "training_config", None) + if training_config is not None: + mode = getattr(training_config, "stage_log_mode", mode) return str(mode).lower() @@ -64,23 +69,64 @@ def resolve_stages( return [s for s in stages if s in requested_set] +def parse_cli_and_stage_args( + parser: argparse.ArgumentParser, + stages: Sequence[str], +) -> Tuple[argparse.Namespace, List[str]]: + """Parse CLI arguments and expand the requested stage selection. + + This helper is a thin wrapper around ``ArgumentParser.parse_args()`` plus + :func:`resolve_stages`. It keeps the runner entrypoints concise and ensures + the `"all"` shorthand is expanded consistently across recipes. + + Args: + parser (argparse.ArgumentParser): Parser configured by the recipe + entrypoint. It is expected to define a ``--stages`` argument whose + value is compatible with ``resolve_stages(...)``. + stages (Sequence[str]): Ordered list of stage names supported by the + current runner, for example ``["create_dataset", "train", "infer"]``. + + Returns: + Tuple[argparse.Namespace, List[str]]: A pair containing: + - the parsed CLI namespace returned by ``parser.parse_args()`` + - the resolved list of stages to execute, preserving the canonical + order from ``stages`` + + Examples: + >>> parser = argparse.ArgumentParser() + >>> parser.add_argument("--stages", nargs="+", default=["all"]) + >>> # If argv is: ["--stages", "infer", "train"] + >>> args, stages_to_run = parse_cli_and_stage_args( + ... parser, + ... ["create_dataset", "train", "infer"], + ... ) + >>> stages_to_run + ['train', 'infer'] + + Notes: + The returned stage order is not the same as CLI input order when they + differ. The canonical order defined by ``stages`` always wins so stage + execution remains deterministic. + """ + args = parser.parse_args() + stages_to_run = resolve_stages(args.stages, stages) + return args, stages_to_run + + def run_stages( system: Any, stages_to_run: Iterable[str], - dry_run: bool = False, + args: argparse.Namespace | None = None, log: logging.Logger | None = None, - on_stage_start: Callable[[str, logging.Logger], None] | None = None, ) -> None: """Invoke stage methods on ``system`` in order with logging and timing. Args: system: Object providing stage methods named in ``stages_to_run``. stages_to_run: Iterable of stage method names to execute. - dry_run: If True, log intended stages without executing them. + args: Parsed CLI arguments. ``dry_run`` and ``write_requirements`` are + read from here when present. log: Optional logger instance; defaults to module logger. - on_stage_start: Optional hook invoked after stage logging is configured. - This can be used to emit per-stage metadata (configs, environment, - requirements snapshots, etc.) into the newly attached log file. Raises: AttributeError: If a named stage method is missing on ``system``. @@ -88,6 +134,7 @@ def run_stages( Exception: Re-raises any exception from a stage method. """ log = log or logger + dry_run = bool(getattr(args, "dry_run", False)) for stage in stages_to_run: fn = getattr(system, stage, None) if fn is None: @@ -129,8 +176,7 @@ def run_stages( Path(log_dir) if log_dir else None, filename=filename, ) - if on_stage_start is not None: - on_stage_start(stage, log) + log_stage_metadata(log, system=system, args=args) start = time.perf_counter() log.info("=== [START] stage: %s ===", stage) diff --git a/espnet3/utils/task_utils.py b/espnet3/utils/task_utils.py index d2cef2e1351d..e7be586c40ad 100644 --- a/espnet3/utils/task_utils.py +++ b/espnet3/utils/task_utils.py @@ -5,12 +5,12 @@ from pathlib import Path from typing import Dict, Union -import yaml from hydra.utils import get_class from omegaconf import DictConfig, OmegaConf from typeguard import typechecked from espnet2.train.abs_espnet_model import AbsESPnetModel +from espnet2.utils.yaml_no_alias_safe_dump import yaml_no_alias_safe_dump def get_task_class(task_path: str): @@ -35,7 +35,7 @@ def get_espnet_model(task: str, config: Union[Dict, DictConfig]) -> AbsESPnetMod finally: sys.argv = original_argv - if isinstance(config, OmegaConf): + if OmegaConf.is_config(config): default_config.update(OmegaConf.to_container(config, resolve=True)) else: default_config.update(config) @@ -56,7 +56,11 @@ def save_espnet_config( default_config = ez_task.get_default_config() sys.argv = original_argv - resolved_config = OmegaConf.to_container(config, resolve=True) + resolved_config = ( + OmegaConf.to_container(config, resolve=True) + if OmegaConf.is_config(config) + else config + ) # set model config at the root level model_config = resolved_config.pop("model") @@ -65,8 +69,9 @@ def save_espnet_config( default_config.update(model_config) # set the preprocessor config at the root level - if config.dataset is not None and "preprocessor" in config.dataset: - preprocess_config = resolved_config["dataset"].pop("preprocessor") + dataset_config = resolved_config.get("dataset") + if dataset_config is not None and "preprocessor" in dataset_config: + preprocess_config = dataset_config.pop("preprocessor") if "_target_" in preprocess_config: preprocess_config.pop("_target_") default_config.update(preprocess_config) @@ -84,8 +89,11 @@ def save_espnet_config( default_config[k] = list(v) # Save the config to the output directory - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - output_path = output_dir / "config.yaml" + output_path = Path(output_dir) + if output_path.suffix: + output_path.parent.mkdir(parents=True, exist_ok=True) + else: + output_path.mkdir(parents=True, exist_ok=True) + output_path = output_path / "config.yaml" with open(output_path, "w", encoding="utf-8") as f: - yaml.dump(default_config, f) + f.write(yaml_no_alias_safe_dump(default_config, indent=4, sort_keys=False)) diff --git a/logs/stage_a.log b/logs/stage_a.log new file mode 100644 index 000000000000..2bcc517d7e53 --- /dev/null +++ b/logs/stage_a.log @@ -0,0 +1,7 @@ +[MSI] 2026-03-12 16:47:21 EDT (stages_utils.py:193) ERROR: [stage_a] Stage 'stage_a' failed after 0.00s +Traceback (most recent call last): + File "/mnt/c/Users/might/Documents/git/espnets/espnet3/espnet3/utils/stages_utils.py", line 184, in run_stages + fn() + File "/mnt/c/Users/might/Documents/git/espnets/espnet3/test/espnet3/utils/test_stages.py", line 68, in stage_a + raise ValueError("boom") +ValueError: boom diff --git a/logs/stage_a1.log b/logs/stage_a1.log new file mode 100644 index 000000000000..e271a9983156 --- /dev/null +++ b/logs/stage_a1.log @@ -0,0 +1,5 @@ +[MSI] 2026-03-12 16:47:09 EDT (stages_utils.py:186) ERROR: [stage_a] Stage 'stage_a' failed (bad arguments) +Traceback (most recent call last): + File "/mnt/c/Users/might/Documents/git/espnets/espnet3/espnet3/utils/stages_utils.py", line 184, in run_stages + fn() +TypeError: test_run_stages_typeerror_wrapped..BadSystem.stage_a() missing 1 required positional argument: 'arg' diff --git a/pyproject.toml b/pyproject.toml index c261c7f3a844..38bac5b992e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,7 +69,10 @@ version = {attr = "espnet2.__version__"} readme = {file = ["README.md"], content-type = "text/markdown"} [tool.setuptools.packages.find] -include = ["espnet*"] +include = ["espnet*", "egs3*"] + +[tool.setuptools.package-data] +egs3 = ["egs3/TEMPLATE/**/*.yaml"] @@ -78,6 +81,7 @@ asr = [ "ctc-segmentation @ git+https://github.com/espnet/ctc-segmentation.git@9b9ea1d", "editdistance", "opt_einsum", + "jiwer", ] tts = [ diff --git a/test/egs3/TEMPLATE/asr/test_run.py b/test/egs3/TEMPLATE/asr/test_run.py new file mode 100644 index 000000000000..1c666ca5ca7a --- /dev/null +++ b/test/egs3/TEMPLATE/asr/test_run.py @@ -0,0 +1,80 @@ +from pathlib import Path + +import pytest + +from egs3.TEMPLATE.asr.run import _load_and_merge_config +from espnet3.utils.config_utils import ( + _resolve_template_config_filename, + load_template_defaults, +) + + +def test_resolve_template_config_filename() -> None: + assert _resolve_template_config_filename("training_config") == "training.yaml" + assert _resolve_template_config_filename("inference_config") == "inference.yaml" + assert _resolve_template_config_filename("metrics_config") == "metrics.yaml" + + +def test_resolve_template_config_filename_invalid() -> None: + with pytest.raises(ValueError, match="Unknown config argument name"): + _resolve_template_config_filename("bad_config") + + +def test_load_template_defaults_train_contains_expected_targets() -> None: + cfg = load_template_defaults("training_config", "egs3.TEMPLATE.asr") + assert ( + cfg.dataset._target_ == "espnet3.components.data.data_organizer.DataOrganizer" + ) + assert cfg.optimizer._target_ == "torch.optim.Adam" + assert cfg.scheduler._target_ == "espnet2.schedulers.warmup_lr.WarmupLR" + + +def test_load_template_defaults_infer_contains_expected_targets() -> None: + cfg = load_template_defaults("inference_config", "egs3.TEMPLATE.asr") + assert ( + cfg.provider._target_ + == "espnet3.systems.base.inference_provider.InferenceProvider" + ) + assert ( + cfg.runner._target_ == "espnet3.systems.base.inference_runner.InferenceRunner" + ) + + +def test_load_and_merge_config_user_overrides_template_defaults(tmp_path: Path) -> None: + user = tmp_path / "train_user.yaml" + user.write_text( + """ +exp_tag: user_train +optimizer: + _target_: torch.optim.SGD + lr: 0.1 +""".strip() + + "\n", + encoding="utf-8", + ) + + cfg = _load_and_merge_config( + user, + "training_config", + template_package="egs3.TEMPLATE.asr", + ) + + assert cfg is not None + # user config should override template default target + assert cfg.optimizer._target_ == "torch.optim.SGD" + assert cfg.optimizer.lr == 0.1 + # template defaults should still be present + assert ( + cfg.dataset._target_ == "espnet3.components.data.data_organizer.DataOrganizer" + ) + + +def test_load_and_merge_config_none_path_returns_none() -> None: + assert ( + _load_and_merge_config( + None, + "metrics_config", + template_package="egs3.TEMPLATE.asr", + ) + is None + ) diff --git a/test/espnet3/components/data/test_dataloader_builder.py b/test/espnet3/components/data/test_dataloader_builder.py index 3775efd5ea28..50d9a480a391 100644 --- a/test/espnet3/components/data/test_dataloader_builder.py +++ b/test/espnet3/components/data/test_dataloader_builder.py @@ -5,7 +5,6 @@ from omegaconf import OmegaConf from torch.utils.data import BatchSampler, Sampler -from espnet3.components.data.data_organizer import DataOrganizer, do_nothing from espnet3.components.data.dataloader import DataLoaderBuilder from espnet3.components.data.dataset import ShardedDataset from espnet3.utils.config_utils import load_config_with_defaults @@ -44,20 +43,16 @@ # | test_iter_factory_from_default_yaml_with_organizer | Uses a real YAML config with iter_factory and verifies batch structure | # noqa: E501 # | test_iter_factory_with_collate_fn | Confirms config-defined collate_fn takes precedence over argument | # noqa: E501 # -# Multiple Iterator Mode (Sharded Dataset) -# | Test Name | Description | # noqa: E501 -# |---------------------------------------|--------------------------------------------------------------------------| # noqa: E501 -# | test_multiple_iterator_shard_initialization | Checks correct shard is selected based on epoch | # noqa: E501 -# | test_multiple_iterator_epoch_shard_switching | Ensures different shard is selected as epoch changes | # noqa: E501 -# # Note: -# - DummyDataset and DummyShardedDataset are used to simulate real-world data layout -# - All `DataLoaderBuilder.build(mode)` modes are exercised: standard, iter_factory, -# and multiple_iterator +# - DummyDataset and DummyDatasetSameLength are used to simulate real-world data layout +# - All `DataLoaderBuilder.build(mode)` modes are exercised: standard, iter_factory DUMMY_DATASET_TARGET = ( "test.espnet3.components.data.test_dataloader_builder." "DummyDataset" ) +DUMMY_SAME_LENGTH_DATASET_TARGET = ( + "test.espnet3.components.data.test_dataloader_builder." "DummyDatasetSameLength" +) DUMMY_SAMPLER_TARGET = ( "test.espnet3.components.data.test_dataloader_builder." "DummySampler" ) @@ -67,6 +62,9 @@ DUMMY_SHARDED_DATASET_TARGET = ( "test.espnet3.components.data.test_dataloader_builder." "DummyShardedDataset" ) +DUMMY_MISSING_SHARD_TARGET = ( + "test.espnet3.components.data.test_dataloader_builder." "DummyMissingShardMethod" +) # -------- Dummy components for testing -------- @@ -120,13 +118,42 @@ def dummy_collate_fn(batch): return {"custom_collated": batch} -class DummyIterFactory: - def __init__(self, dataset, batches, **kwargs): - self.dataset = dataset - self.batches = list(batches) +class DummyShardedDataset(ShardedDataset): + def __init__( + self, shard_id: int = 0, num_shards: int = 2, world_shard_size: int = 1 + ): + self.shard_id = shard_id + self.num_shards = num_shards + self.world_shard_size = world_shard_size + self.data = [ + {"text": f"shard{shard_id}_sample0"}, + {"text": f"shard{shard_id}_sample1"}, + ] + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx] + + def shard(self, idx: int): + return DummyShardedDataset( + shard_id=idx, + num_shards=self.num_shards, + world_shard_size=self.world_shard_size, + ) + + +class DummyMissingShardMethod: + def __init__(self, num_shards=2, world_shard_size=1): + self.num_shards = num_shards + self.world_shard_size = world_shard_size + + def __len__(self): + return 0 - def build_iter(self, epoch, shuffle=False): - return list(self.batches) + def __getitem__(self, idx): + raise IndexError # -------- Config mocks -------- @@ -147,34 +174,75 @@ def make_standard_dataloader_config(sampler=None, batch_sampler=None, collate_fn return OmegaConf.create(config) +def make_dataset_config(dataset_target, dataset_kwargs=None): + dataset_kwargs = dataset_kwargs or {} + config = { + "_target_": "espnet3.components.data.data_organizer.DataOrganizer", + "train": [ + { + "name": "train_dummy", + "dataset": {"_target_": dataset_target, **dataset_kwargs}, + } + ], + "valid": [ + { + "name": "valid_dummy", + "dataset": {"_target_": dataset_target, **dataset_kwargs}, + } + ], + } + return OmegaConf.create(config) + + +def build_organizer(dataset_target, dataset_kwargs=None): + config = make_dataset_config(dataset_target, dataset_kwargs=dataset_kwargs) + return instantiate(config) + + +def build_builder(dataset, config, collate_fn, num_device, epoch): + return DataLoaderBuilder( + dataset=dataset, + config=config, + collate_fn=collate_fn, + num_device=num_device, + epoch=epoch, + ) + + # -------- Standard PyTorch DataLoader mode -------- def test_batch_sampler_only(): - dataset = DummyDataset() + organizer = build_organizer(DUMMY_DATASET_TARGET) config = make_standard_dataloader_config( batch_sampler={"_target_": DUMMY_BATCH_SAMPLER_TARGET} ) # We don't need batch size for batch sampler del config.dataloader.train.batch_size del config.dataloader.train.shuffle - builder = DataLoaderBuilder(dataset, config, collate_fn=None, num_device=1, epoch=0) + builder = build_builder( + organizer.train, config, collate_fn=None, num_device=1, epoch=0 + ) loader = builder.build("train") assert "DummyBatchSampler" in str(loader.batch_sampler.__class__) def test_sampler_only(): - dataset = DummyDataset() + organizer = build_organizer(DUMMY_DATASET_TARGET) config = make_standard_dataloader_config(sampler={"_target_": DUMMY_SAMPLER_TARGET}) - builder = DataLoaderBuilder(dataset, config, collate_fn=None, num_device=1, epoch=0) + builder = build_builder( + organizer.train, config, collate_fn=None, num_device=1, epoch=0 + ) loader = builder.build("train") assert "DummySampler" in str(loader.sampler.__class__) def test_collate_fn_none(): - dataset = DummyDatasetSameLength() + organizer = build_organizer(DUMMY_SAME_LENGTH_DATASET_TARGET) config = make_standard_dataloader_config() - builder = DataLoaderBuilder(dataset, config, collate_fn=None, num_device=1, epoch=0) + builder = build_builder( + organizer.train, config, collate_fn=None, num_device=1, epoch=0 + ) loader = builder.build("train") batch = next(iter(loader)) assert isinstance(batch, dict) @@ -184,25 +252,12 @@ def test_collate_fn_none(): def test_common_collate_fn(): from espnet2.train.collate_fn import CommonCollateFn - config = { - "train": [ - { - "name": "train_dummy", - "dataset": {"_target_": DUMMY_DATASET_TARGET}, - } - ], - } - config = OmegaConf.create(config) - dataset = DataOrganizer( - train=instantiate(config["train"]), - valid=instantiate(config["train"]), - preprocessor=do_nothing, - ) + organizer = build_organizer(DUMMY_DATASET_TARGET) config = make_standard_dataloader_config() collate_fn = CommonCollateFn(int_pad_value=-1) - dataset.train.use_espnet_collator = True - builder = DataLoaderBuilder( - dataset.train, config, collate_fn=collate_fn, num_device=1, epoch=0 + organizer.train.use_espnet_collator = True + builder = build_builder( + organizer.train, config, collate_fn=collate_fn, num_device=1, epoch=0 ) loader = builder.build("train") batch = next(iter(loader)) @@ -213,10 +268,10 @@ def test_common_collate_fn(): def test_custom_collate_fn(): - dataset = DummyDataset() + organizer = build_organizer(DUMMY_DATASET_TARGET) config = make_standard_dataloader_config() - builder = DataLoaderBuilder( - dataset, config, collate_fn=dummy_collate_fn, num_device=1, epoch=0 + builder = build_builder( + organizer.train, config, collate_fn=dummy_collate_fn, num_device=1, epoch=0 ) loader = builder.build("train") batch = next(iter(loader)) @@ -224,36 +279,25 @@ def test_custom_collate_fn(): def test_sampler_and_batch_sampler_conflict(): - dataset = DummyDataset() + organizer = build_organizer(DUMMY_DATASET_TARGET) config = make_standard_dataloader_config( sampler={"_target_": DUMMY_SAMPLER_TARGET}, batch_sampler={"_target_": DUMMY_BATCH_SAMPLER_TARGET}, ) - builder = DataLoaderBuilder(dataset, config, collate_fn=None, num_device=1, epoch=0) + builder = build_builder( + organizer.train, config, collate_fn=None, num_device=1, epoch=0 + ) with pytest.raises( AssertionError, match="Cannot specify both sampler and batch_sampler" ): - _ = builder._build_standard_dataloader(config.dataloader.train, mode="train") + _ = builder._build_standard_dataloader(config.dataloader.train) # -------- IterFactory Mode & YAML-based Integration -------- def test_iter_factory_from_default_yaml_with_organizer(tmp_path): - config = { - "train": [ - { - "name": "train_dummy", - "dataset": {"_target_": DUMMY_DATASET_TARGET}, - } - ], - } - config = OmegaConf.create(config) - dataset = DataOrganizer( - train=instantiate(config["train"]), - valid=instantiate(config["train"]), - preprocessor=do_nothing, - ) + organizer = build_organizer(DUMMY_DATASET_TARGET) yaml_text = """ dataloader: train: @@ -274,9 +318,9 @@ def test_iter_factory_from_default_yaml_with_organizer(tmp_path): config_path.write_text(yaml_text) cfg = load_config_with_defaults(str(config_path)) - dataset.train.use_espnet_collator = True - builder = DataLoaderBuilder( - dataset=dataset.train, + organizer.train.use_espnet_collator = True + builder = build_builder( + dataset=organizer.train, config=cfg, collate_fn=None, # Defined on config num_device=1, @@ -292,20 +336,7 @@ def test_iter_factory_from_default_yaml_with_organizer(tmp_path): def test_iter_factory_with_collate_fn(tmp_path): # For this case the collate_fn in configuration is used. # The collate_fn in DataloaderBuilder.__init__ is only used in standard dataloader. - config = { - "train": [ - { - "name": "train_dummy", - "dataset": {"_target_": DUMMY_DATASET_TARGET}, - } - ], - } - config = OmegaConf.create(config) - dataset = DataOrganizer( - train=instantiate(config["train"]), - valid=instantiate(config["train"]), - preprocessor=do_nothing, - ) + organizer = build_organizer(DUMMY_DATASET_TARGET) yaml_text = """ dataloader: train: @@ -326,9 +357,9 @@ def test_iter_factory_with_collate_fn(tmp_path): config_path.write_text(yaml_text) cfg = load_config_with_defaults(str(config_path)) - dataset.train.use_espnet_collator = True - builder = DataLoaderBuilder( - dataset=dataset.train, + organizer.train.use_espnet_collator = True + builder = build_builder( + dataset=organizer.train, config=cfg, collate_fn=dummy_collate_fn, # Defined on config num_device=1, @@ -342,148 +373,214 @@ def test_iter_factory_with_collate_fn(tmp_path): assert "audio_lengths" in batch[1] -@pytest.mark.parametrize("rank", [0, 1]) -def test_iter_factory_drops_tail_batches_for_ddp(monkeypatch, rank): - import espnet3.components.data.dataloader as dl - - # 5 total batches -> for world_size=2, drop 1 tail batch, keep 4 - def _fake_build_batch_sampler(**kwargs): - return [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]] - - monkeypatch.setattr(dl, "build_batch_sampler", _fake_build_batch_sampler) - monkeypatch.setattr(dl.torch.distributed, "get_world_size", lambda: 2) - monkeypatch.setattr(dl.torch.distributed, "get_rank", lambda: rank) - - config = OmegaConf.create( - { - "dataloader": { - "train": { - "iter_factory": { - "_target_": ( - "test.espnet3.components.data." - "test_dataloader_builder.DummyIterFactory" - ), - "batches": {"dummy": 1}, - } - } - } - } - ) - - dataset = DummyDataset() - builder = DataLoaderBuilder( - dataset=dataset, config=config, collate_fn=None, num_device=2, epoch=0 - ) - iterator = builder.build("train") - - assert len(iterator) == 2 - assert all(batch[0] != 8 for batch in iterator) - - -# --- Multiple Iterator Mode (Sharded Dataset) --- - - -# Dummy sharded dataset that returns shard-specific samples -class DummyShardedDataset(ShardedDataset): - def __init__(self, shard_id: int = None): - if shard_id is None: - self.samples = [] - else: - self.samples = [ - {"text": f"shard{shard_id}_sample0"}, - {"text": f"shard{shard_id}_sample1"}, - ] - self.shard_id = shard_id - - def __len__(self): - return len(self.samples) - - def __getitem__(self, idx): - return self.samples[idx] - - def shard(self, idx: int): - return DummyShardedDataset(idx) - - -@pytest.fixture -def dummy_multiple_iterator_dataset(tmp_path): - # YAML-style instantiation config for 3 shards - dataset_config = { - "train": [ - { - "name": "shard0", - "dataset": { - "_target_": DUMMY_SHARDED_DATASET_TARGET, - }, - }, - ], - "valid": [ - { - "name": "valid", - "dataset": { - "_target_": DUMMY_SHARDED_DATASET_TARGET, - }, - } - ], - } - return OmegaConf.create(dataset_config) - - -def make_multiple_iterator_config(num_shards: int, shuffle: bool): - return OmegaConf.create( - { - "dataloader": { - "train": { - "multiple_iterator": True, - "num_shards": num_shards, - "shuffle": shuffle, - "iter_factory": None, - "batch_size": 1, - "num_workers": 0, - } - }, - "num_device": 1, - } - ) - - -def test_multiple_iterator_shard_initialization(dummy_multiple_iterator_dataset): - organizer = DataOrganizer( - train=instantiate(dummy_multiple_iterator_dataset.train), - valid=instantiate(dummy_multiple_iterator_dataset.valid), - ) - config = make_multiple_iterator_config(num_shards=3, shuffle=False) - builder = DataLoaderBuilder( - dataset=organizer.train, - config=config, - collate_fn=None, - num_device=1, - epoch=0, - ) - loader = builder.build("train") - batch = next(iter(loader)) - assert "text" in batch - assert batch["text"][0].startswith("shard0_") - - -@pytest.mark.parametrize( - "epoch,expected_shard", [(0, "shard0"), (1, "shard1"), (2, "shard2")] -) -def test_multiple_iterator_epoch_shard_switching( - dummy_multiple_iterator_dataset, epoch, expected_shard -): - organizer = DataOrganizer( - train=instantiate(dummy_multiple_iterator_dataset["train"]), - valid=instantiate(dummy_multiple_iterator_dataset["valid"]), - ) - config = make_multiple_iterator_config(num_shards=3, shuffle=False) - builder = DataLoaderBuilder( - dataset=organizer.train, - config=config, - collate_fn=None, - num_device=1, - epoch=epoch, - ) - loader = builder.build("train") - batch = next(iter(loader)) - assert expected_shard in batch["text"][0] +# Commented out the tests for future bug fix. +# These tests are for the bug fixed in PR 14.5 (next PR) + +# @pytest.mark.parametrize("flag", [True, False]) +# def test_multiple_iterator_is_rejected(flag): +# organizer = build_organizer(DUMMY_DATASET_TARGET) +# config = make_standard_dataloader_config() +# config.dataloader.train.multiple_iterator = flag +# builder = build_builder( +# organizer.train, config, collate_fn=None, num_device=1, epoch=0 +# ) +# with pytest.raises(RuntimeError, match="multiple_iterator"): +# builder.build("train") + + +# def _collect_shard_ids(loader): +# shard_ids = set() +# for batch in loader: +# text = batch["text"][0] if isinstance(batch, dict) else batch[0]["text"] +# shard_ids.add(text.split("_")[0]) +# if len(shard_ids) >= 4: +# break +# return shard_ids + + +# def _first_shard_id(loader): +# batch = next(iter(loader)) +# text = batch["text"][0] if isinstance(batch, dict) else batch[0]["text"] +# return text.split("_")[0] + + +# def test_sharded_dataset_single_gpu_multiple_shards(): +# organizer = build_organizer( +# DUMMY_SHARDED_DATASET_TARGET, +# ataset_kwargs={"num_shards": 2, "world_shard_size": 1} +# ) +# config = make_standard_dataloader_config() +# config.dataloader.train.batch_size = 1 +# builder = build_builder( +# organizer.train, config, collate_fn=None, num_device=1, epoch=0 +# ) +# loader = builder.build("train") +# shard_ids = _collect_shard_ids(loader) +# assert shard_ids == {"shard0"} + + +# @pytest.mark.parametrize( +# "world_size,num_shards,rank,expected", +# [ +# (2, 2, 0, {"shard0"}), +# (2, 2, 1, {"shard1"}), +# (2, 4, 0, {"shard0"}), +# (2, 4, 1, {"shard1"}), +# ], +# ) +# def test_sharded_dataset_multi_gpu_assignment( +# monkeypatch, world_size, num_shards, rank, expected +# ): +# def _get_world_size(): +# return world_size + +# def _get_rank(): +# return rank + +# monkeypatch.setattr(torch.distributed, "get_world_size", _get_world_size) +# monkeypatch.setattr(torch.distributed, "get_rank", _get_rank) + +# organizer = build_organizer( +# DUMMY_SHARDED_DATASET_TARGET, +# dataset_kwargs={"num_shards": num_shards, "world_shard_size": world_size}, +# ) +# config = make_standard_dataloader_config() +# config.dataloader.train.batch_size = 1 +# builder = build_builder( +# organizer.train, +# config, +# collate_fn=None, +# num_device=world_size, +# epoch=0, +# ) +# loader = builder.build("train") +# shard_ids = _collect_shard_ids(loader) +# assert shard_ids == expected + + +# def test_sharded_dataset_multi_gpu_rotates_with_epoch(monkeypatch): +# def _get_world_size(): +# return 2 + +# def _get_rank(): +# return 0 + +# monkeypatch.setattr(torch.distributed, "get_world_size", _get_world_size) +# monkeypatch.setattr(torch.distributed, "get_rank", _get_rank) + +# organizer = build_organizer( +# DUMMY_SHARDED_DATASET_TARGET, +# dataset_kwargs={"num_shards": 4, "world_shard_size": 2} +# ) +# config = make_standard_dataloader_config() +# config.dataloader.train.batch_size = 1 + +# builder_epoch0 = build_builder( +# organizer.train, config, collate_fn=None, num_device=2, epoch=0 +# ) +# first_epoch0 = _first_shard_id(builder_epoch0.build("train")) + +# builder_epoch1 = build_builder( +# organizer.train, config, collate_fn=None, num_device=2, epoch=1 +# ) +# first_epoch1 = _first_shard_id(builder_epoch1.build("train")) + +# assert first_epoch0 == "shard0" +# assert first_epoch1 == "shard2" + + +# @pytest.mark.parametrize( +# "epoch,expected_rank0,expected_rank1", +# [ +# (0, "shard0", "shard1"), +# (1, "shard2", "shard3"), +# (2, "shard0", "shard1"), +# (3, "shard2", "shard3"), +# ], +# ) +# def test_sharded_dataset_multi_gpu_rotates_epochs_0_to_3( +# monkeypatch, epoch, expected_rank0, expected_rank1 +# ): +# def _get_world_size(): +# return 2 + +# monkeypatch.setattr(torch.distributed, "get_world_size", _get_world_size) + +# organizer = build_organizer( +# DUMMY_SHARDED_DATASET_TARGET, +# dataset_kwargs={"num_shards": 4, "world_shard_size": 2}, +# ) +# config = make_standard_dataloader_config() +# config.dataloader.train.batch_size = 1 + +# def _first_shard_for_rank(rank): +# monkeypatch.setattr(torch.distributed, "get_rank", lambda: rank) +# builder = build_builder( +# organizer.train, config, collate_fn=None, num_device=2, epoch=epoch +# ) +# return _first_shard_id(builder.build("train")) + +# assert _first_shard_for_rank(0) == expected_rank0 +# assert _first_shard_for_rank(1) == expected_rank1 + + +# @pytest.mark.parametrize("num_shards", [1, 3]) +# def test_sharded_dataset_invalid_shard_count(monkeypatch, num_shards): +# def _get_world_size(): +# return 2 + +# def _get_rank(): +# return 0 + +# monkeypatch.setattr(torch.distributed, "get_world_size", _get_world_size) +# monkeypatch.setattr(torch.distributed, "get_rank", _get_rank) + +# organizer = build_organizer( +# DUMMY_SHARDED_DATASET_TARGET, +# dataset_kwargs={"num_shards": num_shards, "world_shard_size": 2} +# ) +# config = make_standard_dataloader_config() +# config.dataloader.train.batch_size = 1 +# builder = build_builder( +# organizer.train, config, collate_fn=None, num_device=2, epoch=0 +# ) +# with pytest.raises( +# RuntimeError, +# match="num_shards must be divisible by world_size" +# ): +# builder.build("train") + + +# def test_sharded_dataset_missing_shard_method(): +# organizer = build_organizer( +# DUMMY_MISSING_SHARD_TARGET, +# dataset_kwargs={"num_shards": 2, "world_shard_size": 1} +# ) +# config = make_standard_dataloader_config() +# builder = build_builder( +# organizer.train, config, collate_fn=None, num_device=1, epoch=0 +# ) +# with pytest.raises(RuntimeError, match="shard\\(\\) is not implemented"): +# builder.build("train") + + +# def test_sharded_dataset_world_size_mismatch(monkeypatch): +# def _get_world_size(): +# return 2 + +# def _get_rank(): +# return 0 + +# monkeypatch.setattr(torch.distributed, "get_world_size", _get_world_size) +# monkeypatch.setattr(torch.distributed, "get_rank", _get_rank) + +# organizer = build_organizer( +# DUMMY_SHARDED_DATASET_TARGET, +# dataset_kwargs={"num_shards": 2, "world_shard_size": 1} +# ) +# config = make_standard_dataloader_config() +# builder = build_builder( +# organizer.train, config, collate_fn=None, num_device=2, epoch=0 +# ) +# with pytest.raises(RuntimeError, match="world_shard_size must match"): +# builder.build("train") diff --git a/test/espnet3/components/modeling/test_model_with_optim_scheduler.py b/test/espnet3/components/modeling/test_model_with_optim_scheduler.py index a8b56400f33b..ffd4e84ca6f0 100644 --- a/test/espnet3/components/modeling/test_model_with_optim_scheduler.py +++ b/test/espnet3/components/modeling/test_model_with_optim_scheduler.py @@ -1,5 +1,4 @@ import pytest -import torch import torch.nn as nn from omegaconf import OmegaConf @@ -16,6 +15,8 @@ # | test_multiple_optims_and_schedulers| Validates multiple optimizers and schedulers with param-based mapping | # noqa: E501 # | test_custom_scheduler_interval | Ensures scheduler interval is correctly set to "step" | # noqa: E501 # | test_reduce_on_plateau_with_config_adam | Tests ReduceLROnPlateau scheduler integration with manual stepping | # noqa: E501 +# | test_reduce_on_plateau_monitor_from_config | Ensures ReduceLROnPlateau gets monitor key from config | # noqa: E501 +# | test_val_scheduler_criterion_sets_epoch | Ensures val_scheduler_criterion switches to epoch | # noqa: E501 # # Invalid Configuration Tests # | Test Name | Description | # noqa: E501 @@ -38,34 +39,6 @@ def __init__(self): self.linear2 = nn.Linear(10, 5) -class ReduceLROnPlateauModel(ESPnetLightningModule): - def __init__(self, model, config): - super().__init__(model, config) - self.favorite_metric = 1.0 - - def configure_optimizers(self): - optimizer = torch.optim.Adam(self.parameters(), lr=0.01) - self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( - optimizer, patience=2, factor=0.5 - ) - return optimizer - - def optimizer_step( - self, - epoch_nb, - batch_nb, - optimizer, - ): - if batch_nb == 0: - self.scheduler.step(self.favorite_metric) - print( - f"metric: {self.favorite_metric}, best: {self.scheduler.best}, " - f"num_bad_epochs: {self.scheduler.num_bad_epochs}" - ) - optimizer.step() - optimizer.zero_grad() - - # ========== VALID CASES ========== @@ -131,7 +104,8 @@ def test_multiple_optims_and_schedulers(): model = ESPnetLightningModule(DummyModel(), config) out = model.configure_optimizers() assert hasattr(out["optimizer"], "optimizers") # HybridOptim - assert isinstance(out["lr_scheduler"]["scheduler"], list) + assert isinstance(out["lr_scheduler"], list) + assert all("scheduler" in item for item in out["lr_scheduler"]) def test_custom_scheduler_interval(): @@ -156,7 +130,7 @@ def test_custom_scheduler_interval(): assert out["lr_scheduler"]["interval"] == "step" -def test_reduce_on_plateau_with_config_adam(): +def test_reduce_on_plateau_monitor_from_config(): config = OmegaConf.create( { "optimizer": {"_target_": "torch.optim.Adam", "lr": 0.01}, @@ -165,6 +139,32 @@ def test_reduce_on_plateau_with_config_adam(): "patience": 1, "factor": 0.5, }, + "val_scheduler_criterion": "valid/loss", + "dataset": { + "_target_": "espnet3.components.data.data_organizer.DataOrganizer", + "train": [], + "valid": [], + }, + "dataloader": {"train": {}, "valid": {}}, + "num_device": 1, + } + ) + + model = ESPnetLightningModule(DummyModel(), config) + out = model.configure_optimizers() + assert out["lr_scheduler"]["interval"] == "epoch" + assert out["lr_scheduler"]["monitor"] == "valid/loss" + + +def test_val_scheduler_criterion_sets_epoch(): + config = OmegaConf.create( + { + "optimizer": {"_target_": "torch.optim.Adam", "lr": 0.01}, + "scheduler": { + "_target_": "torch.optim.lr_scheduler.StepLR", + "step_size": 5, + }, + "val_scheduler_criterion": "valid/acc", "dataset": { "_target_": "espnet3.components.data.data_organizer.DataOrganizer", "train": [], @@ -175,21 +175,10 @@ def test_reduce_on_plateau_with_config_adam(): } ) - model = DummyModel() - lit_model = ReduceLROnPlateauModel(model, config) - optimizer = lit_model.configure_optimizers() - - # mimic training loop - for epoch in range(3): - for batch in range(1): # we step scheduler at batch 0 - lit_model.favorite_metric = ( - 1.0 - 0.1 * epoch - ) # simulate val_loss decreasing - lit_model.optimizer_step( - epoch_nb=epoch, - batch_nb=batch, - optimizer=optimizer, - ) + model = ESPnetLightningModule(DummyModel(), config) + out = model.configure_optimizers() + assert out["lr_scheduler"]["interval"] == "epoch" + assert out["lr_scheduler"]["monitor"] == "valid/acc" # ========== INVALID CASES ========== diff --git a/test/espnet3/components/trainers/test_trainer.py b/test/espnet3/components/trainers/test_trainer.py index 351f08168a3b..0e8af0c92c79 100644 --- a/test/espnet3/components/trainers/test_trainer.py +++ b/test/espnet3/components/trainers/test_trainer.py @@ -213,6 +213,41 @@ def test_logger_variants( assert isinstance(wrapper.trainer.logger, expect_logger_type) +def test_collect_stats_then_train_does_not_mutate_trainer_config( + base_trainer_config, model_config, dummy_dataset_config, monkeypatch +): + model_config = OmegaConf.create(model_config) + model_config.dataset = dummy_dataset_config + trainer_config = OmegaConf.create(base_trainer_config) + trainer_config.logger = [ + { + "_target_": "lightning.pytorch.loggers.TensorBoardLogger", + "save_dir": "exp/tb", + "name": "tb", + } + ] + + original_logger = OmegaConf.to_container(trainer_config.logger, resolve=True) + original_keys = set(trainer_config.keys()) + + model = nn.Linear(10, 1) + lit = ESPnetLightningModule(model, model_config) + monkeypatch.setattr( + ESPnetLightningModule, "collect_stats", lambda *args, **kwargs: None + ) + + wrapper_stats = ESPnet3LightningTrainer( + model=lit, config=trainer_config, exp_dir=exp_dir + ) + wrapper_stats.collect_stats() + + assert set(trainer_config.keys()) == original_keys + assert hasattr(trainer_config, "logger") + assert ( + OmegaConf.to_container(trainer_config.logger, resolve=True) == original_logger + ) + + @pytest.mark.parametrize( "accelerator, expect_type", [ diff --git a/test/espnet3/systems/asr/test_asr_inference.py b/test/espnet3/systems/asr/test_asr_inference.py index 0f066248468c..fea68014da8c 100644 --- a/test/espnet3/systems/asr/test_asr_inference.py +++ b/test/espnet3/systems/asr/test_asr_inference.py @@ -7,13 +7,13 @@ def _output_fn(*, data, model_output, idx): - return {"uttid": data["uttid"], "hyp": model_output[0][0], "ref": data["text"]} + return {"utt_id": data["utt_id"], "hyp": model_output[0][0], "ref": data["text"]} def _batch_output_fn(*, data, model_output, idx): return [ { - "uttid": sample["uttid"], + "utt_id": sample["utt_id"], "hyp": output[0][0], "ref": sample["text"], } @@ -22,7 +22,7 @@ def _batch_output_fn(*, data, model_output, idx): def _param_output_fn(*, data, model_output, idx): - return {"uttid": data["uttid"], "hyp": model_output, "ref": "ref"} + return {"utt_id": data["utt_id"], "hyp": model_output, "ref": "ref"} class DummyProvider(InferenceProvider): @@ -105,7 +105,7 @@ def fake_instantiate(obj, device=None): def test_forward_returns_hyp_and_ref(): - dataset = [{"uttid": "utt1", "speech": "audio", "text": "ref"}] + dataset = [{"utt_id": "utt1", "speech": "audio", "text": "ref"}] class DummyModel: def __call__(self, speech): @@ -121,13 +121,31 @@ def __call__(self, speech): output_fn_path=output_path, ) - assert out == {"uttid": "utt1", "hyp": "hyp", "ref": "ref"} + assert out == {"utt_id": "utt1", "hyp": "hyp", "ref": "ref"} + + +def test_forward_without_output_fn_returns_raw_model_output(): + dataset = [{"utt_id": "utt1", "speech": "audio", "text": "ref"}] + + class DummyModel: + def __call__(self, speech): + assert speech == "audio" + return {"utt_id": "utt1", "hyp": "hyp"} + + out = InferenceRunner.forward( + 0, + dataset=dataset, + model=DummyModel(), + input_key="speech", + ) + + assert out == {"utt_id": "utt1", "hyp": "hyp"} def test_forward_batch_with_batched_inputs(): dataset = [ - {"uttid": "utt1", "speech": "audio1", "text": "ref1"}, - {"uttid": "utt2", "speech": "audio2", "text": "ref2"}, + {"utt_id": "utt1", "speech": "audio1", "text": "ref1"}, + {"utt_id": "utt2", "speech": "audio2", "text": "ref2"}, ] class DummyModel: @@ -145,15 +163,15 @@ def __call__(self, **inputs): ) assert out == [ - {"uttid": "utt1", "hyp": "hyp1", "ref": "ref1"}, - {"uttid": "utt2", "hyp": "hyp2", "ref": "ref2"}, + {"utt_id": "utt1", "hyp": "hyp1", "ref": "ref1"}, + {"utt_id": "utt2", "hyp": "hyp2", "ref": "ref2"}, ] def test_forward_batch_requires_batched_model(): dataset = [ - {"uttid": "utt1", "speech": "audio1", "text": "ref1"}, - {"uttid": "utt2", "speech": "audio2", "text": "ref2"}, + {"utt_id": "utt1", "speech": "audio1", "text": "ref1"}, + {"utt_id": "utt2", "speech": "audio2", "text": "ref2"}, ] class DummyModel: @@ -182,7 +200,7 @@ def test_forward_requires_fields(): output_fn_path=f"{__name__}._output_fn", ) - dataset = [{"uttid": "utt1", "speech": "audio"}] + dataset = [{"utt_id": "utt1", "speech": "audio"}] with pytest.raises(KeyError, match="text"): InferenceRunner.forward( 0, @@ -204,7 +222,7 @@ def test_inference_requires_provider_config(): "parallel": {"env": "local", "n_workers": 1}, } ) - with pytest.raises(RuntimeError, match="infer_config.provider must be set"): + with pytest.raises(RuntimeError, match="inference_config.provider must be set"): infer(cfg) @@ -215,7 +233,7 @@ def test_inference_params_affect_runner_forward(tmp_path, flip, expected): "inference_dir": str(tmp_path), "dataset": { "test": [{"name": "test"}], - "data": [{"uttid": "utt1", "speech": "s1"}], + "data": [{"utt_id": "utt1", "speech": "s1"}], }, "input_key": "speech", "output_fn": f"{__name__}._param_output_fn", diff --git a/test/espnet3/systems/asr/test_asr_transducer.py b/test/espnet3/systems/asr/test_asr_transducer.py new file mode 100644 index 000000000000..ff3e316a54b1 --- /dev/null +++ b/test/espnet3/systems/asr/test_asr_transducer.py @@ -0,0 +1,116 @@ +from argparse import Namespace + +import pytest + +from espnet3.systems.asr.transducer_task import ASRTransducerTask + + +def test_add_arguments(): + ASRTransducerTask.get_parser() + + +def test_add_arguments_help(): + parser = ASRTransducerTask.get_parser() + with pytest.raises(SystemExit): + parser.parse_args(["--help"]) + + +def test_main_help(): + with pytest.raises(SystemExit): + ASRTransducerTask.main(cmd=["--help"]) + + +def test_main_print_config(): + with pytest.raises(SystemExit): + ASRTransducerTask.main(cmd=["--print_config"]) + + +def test_main_with_no_args(): + with pytest.raises(SystemExit): + ASRTransducerTask.main(cmd=[]) + + +def test_print_config_and_load_it(tmp_path): + config_file = tmp_path / "config.yaml" + with config_file.open("w") as f: + ASRTransducerTask.print_config(f) + parser = ASRTransducerTask.get_parser() + parser.parse_args(["--config", str(config_file)]) + + +def get_dummy_namespace(): + return Namespace( + token_type="char", + token_list=["", "a", "b", "c", "", ""], + input_size=40, + frontend="frontend", + frontend_conf={"n_fft": 51, "win_length": 40, "hop_length": 16}, + specaug="specaug", + specaug_conf={"apply_time_warp": True, "time_mask_width_range": 4}, + normalize=None, + normalize_conf=None, + encoder_conf={ + "body_conf": [ + { + "block_type": "conformer", + "hidden_size": 16, + "linear_size": 8, + "conv_mod_kernel_size": 3, + } + ] + }, + decoder="stateless", + decoder_conf={"embed_size": 8}, + joint_network_conf={"joint_space_size": 8}, + model_conf={"transducer_weight": 1.0}, + init=None, + ) + + +def test_build_model(): + args = get_dummy_namespace() + + _ = ASRTransducerTask.build_model(args) + + with pytest.raises(RuntimeError): + args.token_list = -1 + + _ = ASRTransducerTask.build_model(args) + + +def test_build_collate_fn(): + args = get_dummy_namespace() + + _ = ASRTransducerTask.build_collate_fn(args, True) + + +@pytest.mark.parametrize("use_preprocessor", [True, False]) +def test_build_preprocess_fn(use_preprocessor): + args = get_dummy_namespace() + + new_args = { + "use_preprocessor": use_preprocessor, + "bpemodel": None, + "non_linguistic_symbols": args.token_list, + "cleaner": None, + "g2p": None, + } + args.__dict__.update(new_args) + + _ = ASRTransducerTask.build_preprocess_fn(args, True) + + +@pytest.mark.parametrize("inference", [True, False]) +def test_required_data_names(inference): + retval = ASRTransducerTask.required_data_names(True, inference) + + if inference: + assert retval == ("speech",) + else: + assert retval == ("speech", "text") + + +def test_optional_data_names(): + retval = ASRTransducerTask.optional_data_names(True, True) + + assert not retval diff --git a/test/espnet3/systems/asr/test_system.py b/test/espnet3/systems/asr/test_system.py index 9c09b401b070..9a76270349b0 100644 --- a/test/espnet3/systems/asr/test_system.py +++ b/test/espnet3/systems/asr/test_system.py @@ -17,7 +17,7 @@ def test_asr_system_create_dataset_invokes_helper(tmp_path, monkeypatch): "create_dataset": {"func": "dummy.dataset", "foo": "bar"}, } ) - system = ASRSystem(train_config=train_cfg) + system = ASRSystem(training_config=train_cfg) calls = {} def fake_fn(**kwargs): @@ -42,7 +42,7 @@ def test_asr_system_train_runs_tokenizer_then_train(tmp_path, monkeypatch): }, } ) - system = ASRSystem(train_config=train_cfg) + system = ASRSystem(training_config=train_cfg) calls = {} def fake_train_tokenizer(self): @@ -73,7 +73,7 @@ def test_asr_system_train_tokenizer_trains_sentencepiece(tmp_path, monkeypatch): }, } ) - system = ASRSystem(train_config=train_cfg) + system = ASRSystem(training_config=train_cfg) calls = {} def fake_builder(foo): diff --git a/test/espnet3/systems/base/test_base_system.py b/test/espnet3/systems/base/test_base_system.py index 24b9103d8697..910f1b75f78c 100644 --- a/test/espnet3/systems/base/test_base_system.py +++ b/test/espnet3/systems/base/test_base_system.py @@ -40,9 +40,9 @@ def fake_metric(cfg): monkeypatch.setattr(sysmod, "measure", fake_metric) system = BaseSystem( - train_config=train_cfg, - infer_config=infer_cfg, - measure_config=measure_cfg, + training_config=train_cfg, + inference_config=infer_cfg, + metrics_config=measure_cfg, ) assert system.exp_dir.is_dir() diff --git a/test/espnet3/systems/base/test_inference.py b/test/espnet3/systems/base/test_inference.py index c51e91b4936b..f2fd2029ad75 100644 --- a/test/espnet3/systems/base/test_inference.py +++ b/test/espnet3/systems/base/test_inference.py @@ -13,8 +13,8 @@ def dummy_output_fn(*, data, model_output, idx): class DummyProvider(InferenceProvider): - def __init__(self, infer_config, params): - super().__init__(infer_config) + def __init__(self, inference_config, params): + super().__init__(inference_config) self.params = params def build_dataset(self, _config): @@ -24,8 +24,8 @@ def build_dataset(self, _config): class DummyRunner(InferenceRunner): results = None - def __init__(self, provider, *, async_mode=False, results=None, **_kwargs): - super().__init__(provider, async_mode=async_mode) + def __init__(self, provider, *, async_mode=False, results=None, **kwargs): + super().__init__(provider, async_mode=async_mode, **kwargs) if results is not None: self._results = results else: @@ -42,8 +42,8 @@ def __call__(self, indices): class CaptureProvider(InferenceRunner): last_params = None - def __init__(self, infer_config, *, params): - super().__init__(infer_config) + def __init__(self, inference_config, *, params): + super().__init__(inference_config) CaptureProvider.last_params = params def build_dataset(self, _config): @@ -61,6 +61,9 @@ def __call__(self, _indices): idx_key = "idx" + def resolve_idx_key(self, _output): + return self.idx_key + def _read_scp(path: Path): return path.read_text(encoding="utf-8").splitlines() @@ -156,3 +159,118 @@ def test_inference_passes_provider_params(tmp_path, monkeypatch): "input_key": "speech", "output_fn_path": f"{__name__}.dummy_output_fn", } + + +def test_inference_without_output_fn_uses_model_output(tmp_path, monkeypatch): + cfg = OmegaConf.create( + { + "parallel": {"env": "local"}, + "inference_dir": str(tmp_path / "infer"), + "dataset": {"test": [{"name": "test_a"}]}, + "input_key": "speech", + "mock_dataset_length": 2, + "provider": {"_target_": f"{__name__}.DummyProvider"}, + "runner": {"_target_": f"{__name__}.DummyRunner"}, + } + ) + DummyRunner.results = [ + {"utt_id": "utt1", "hyp": "h1"}, + {"utt_id": "utt2", "hyp": "h2"}, + ] + + monkeypatch.setattr(inference_mod, "set_parallel", lambda arg: None) + + inference_mod.infer(cfg) + + assert _read_scp(tmp_path / "infer" / "test_a" / "hyp.scp") == [ + "utt1 h1", + "utt2 h2", + ] + + +def test_inference_without_idx_key_uses_default_utt_id(tmp_path, monkeypatch): + cfg = OmegaConf.create( + { + "parallel": {"env": "local"}, + "inference_dir": str(tmp_path / "infer"), + "dataset": {"test": [{"name": "test_a"}]}, + "input_key": "speech", + "output_fn": f"{__name__}.dummy_output_fn", + "mock_dataset_length": 1, + "provider": {"_target_": f"{__name__}.DummyProvider"}, + "runner": {"_target_": f"{__name__}.DummyRunner"}, + } + ) + DummyRunner.results = [{"utt_id": "utt1", "hyp": "h1"}] + + monkeypatch.setattr(inference_mod, "set_parallel", lambda arg: None) + + inference_mod.infer(cfg) + + assert _read_scp(tmp_path / "infer" / "test_a" / "hyp.scp") == ["utt1 h1"] + + +def test_inference_with_explicit_idx_key_override(tmp_path, monkeypatch): + cfg = OmegaConf.create( + { + "parallel": {"env": "local"}, + "inference_dir": str(tmp_path / "infer"), + "dataset": {"test": [{"name": "test_a"}]}, + "input_key": "speech", + "idx_key": "sample_id", + "mock_dataset_length": 1, + "provider": {"_target_": f"{__name__}.DummyProvider"}, + "runner": {"_target_": f"{__name__}.DummyRunner"}, + } + ) + DummyRunner.results = [{"sample_id": "sample-1", "hyp": "h1"}] + + monkeypatch.setattr(inference_mod, "set_parallel", lambda arg: None) + + inference_mod.infer(cfg) + + assert _read_scp(tmp_path / "infer" / "test_a" / "hyp.scp") == ["sample-1 h1"] + + +def test_inference_with_explicit_output_keys_filters_scp_outputs(tmp_path, monkeypatch): + cfg = OmegaConf.create( + { + "parallel": {"env": "local"}, + "inference_dir": str(tmp_path / "infer"), + "dataset": {"test": [{"name": "test_a"}]}, + "input_key": "speech", + "output_keys": ["hyp"], + "mock_dataset_length": 1, + "provider": {"_target_": f"{__name__}.DummyProvider"}, + "runner": {"_target_": f"{__name__}.DummyRunner"}, + } + ) + DummyRunner.results = [{"utt_id": "utt1", "hyp": "h1", "ref": "r1"}] + + monkeypatch.setattr(inference_mod, "set_parallel", lambda arg: None) + + inference_mod.infer(cfg) + + base = tmp_path / "infer" / "test_a" + assert (base / "hyp.scp").exists() + assert not (base / "ref.scp").exists() + + +def test_inference_without_output_fn_requires_utterance_id(tmp_path, monkeypatch): + cfg = OmegaConf.create( + { + "parallel": {"env": "local"}, + "inference_dir": str(tmp_path / "infer"), + "dataset": {"test": [{"name": "test_a"}]}, + "input_key": "speech", + "mock_dataset_length": 1, + "provider": {"_target_": f"{__name__}.DummyProvider"}, + "runner": {"_target_": f"{__name__}.DummyRunner"}, + } + ) + DummyRunner.results = [{"hyp": "h1"}] + + monkeypatch.setattr(inference_mod, "set_parallel", lambda arg: None) + + with pytest.raises(ValueError, match="sample identifier key"): + inference_mod.infer(cfg) diff --git a/test/espnet3/systems/base/test_inference_runner.py b/test/espnet3/systems/base/test_inference_runner.py index baa3346736a9..34d4cb3d6ab3 100644 --- a/test/espnet3/systems/base/test_inference_runner.py +++ b/test/espnet3/systems/base/test_inference_runner.py @@ -23,19 +23,41 @@ def test_validate_output_accepts_expected_keys(): runner._validate_output(output) +def test_resolve_idx_key_uses_default_utt_id(): + runner = DummyRunner(DummyProvider()) + + assert runner.resolve_idx_key({"utt_id": "u1", "hyp": "h", "ref": "r"}) == "utt_id" + + +def test_resolve_idx_key_uses_explicit_configured_key(): + runner = DummyRunner(DummyProvider(), idx_key="sample_id") + + assert ( + runner.resolve_idx_key({"sample_id": "u1", "hyp": "h", "ref": "r"}) + == "sample_id" + ) + + +def test_resolve_idx_key_rejects_missing_configured_key(): + runner = DummyRunner(DummyProvider(), idx_key="sample_id") + + with pytest.raises(ValueError, match="idx_key='sample_id'"): + runner.resolve_idx_key({"utt_id": "u1", "hyp": "h", "ref": "r"}) + + def test_validate_output_rejects_missing_or_extra_keys(): runner = DummyRunner(DummyProvider()) - with pytest.raises(ValueError, match="missing=\\['ref'\\]"): - runner._validate_output({"idx": 1, "hyp": "h"}) - runner._validate_output({"idx": 1, "hyp": "h", "ref": "r", "extra": 0}) + with pytest.raises(ValueError, match="idx_key='utt_id'"): + runner._validate_output({"hyp": "h"}) + runner._validate_output({"utt_id": "u1", "hyp": "h", "ref": "r", "extra": 0}) def test_validate_output_rejects_non_dict_and_idx_list(): runner = DummyRunner(DummyProvider()) with pytest.raises(TypeError, match="Expected dict output"): runner._validate_output("bad") - with pytest.raises(TypeError, match="'idx' must be a scalar"): - runner._validate_output({"idx": [1], "hyp": "h", "ref": "r"}) + with pytest.raises(TypeError, match="'utt_id' must be a single value"): + runner._validate_output({"utt_id": [1], "hyp": "h", "ref": "r"}) def test_call_async_returns_raw(monkeypatch): @@ -56,7 +78,7 @@ def fake_base_call(self, indices): ] monkeypatch.setattr(BaseRunner, "__call__", fake_base_call) - runner = DummyRunner(DummyProvider()) + runner = DummyRunner(DummyProvider(), idx_key="idx") assert runner([0, 1]) == [ {"idx": 0, "hyp": "h0", "ref": "r0"}, @@ -72,10 +94,8 @@ def test_call_returns_none(monkeypatch): def test_call_propagates_validation_error(monkeypatch): - monkeypatch.setattr( - BaseRunner, "__call__", lambda self, indices: [{"idx": 0, "hyp": "h"}] - ) + monkeypatch.setattr(BaseRunner, "__call__", lambda self, indices: [{"hyp": "h"}]) runner = DummyRunner(DummyProvider()) - with pytest.raises(ValueError, match="missing=\\['ref'\\]"): + with pytest.raises(ValueError, match="idx_key='utt_id'"): runner([0]) diff --git a/test/espnet3/systems/base/test_metric.py b/test/espnet3/systems/base/test_metric.py index 5781b64f7073..0c1d3c3f4da0 100644 --- a/test/espnet3/systems/base/test_metric.py +++ b/test/espnet3/systems/base/test_metric.py @@ -5,7 +5,7 @@ from omegaconf import OmegaConf from espnet3.components.metrics.abs_metric import AbsMetric -from espnet3.systems.base.metric import measure +from espnet3.systems.base.metric import _resolve_test_sets, measure from espnet3.utils.scp_utils import get_class_path @@ -82,6 +82,61 @@ def test_metric_uses_config_inputs_mapping(tmp_path): assert results[expected_key][test_name] == {"ok": True} +def test_metric_discovers_test_sets_from_inference_dir(tmp_path): + inference_dir = tmp_path / "infer" + for test_name in ("test_b", "test_a"): + task_dir = inference_dir / test_name + task_dir.mkdir(parents=True) + _write_scp(task_dir / "ref.scp", ["utt1 r1"]) + _write_scp(task_dir / "hyp.scp", ["utt1 h1"]) + + cfg = OmegaConf.create( + { + "inference_dir": str(inference_dir), + "metrics": [{"metric": {"_target_": f"{__name__}.DummyMetric"}}], + } + ) + + assert _resolve_test_sets(cfg) == ["test_a", "test_b"] + results = measure(cfg) + + expected_key = get_class_path(DummyMetric()) + assert set(results[expected_key]) == {"test_a", "test_b"} + assert results[expected_key]["test_a"] == {"count": 1} + assert results[expected_key]["test_b"] == {"count": 1} + + +def test_resolve_test_sets_prefers_metrics_config_dataset_over_inference_dir(tmp_path): + inference_dir = tmp_path / "infer" + for test_name in ("from_dir_a", "from_dir_b"): + (inference_dir / test_name).mkdir(parents=True) + + cfg = OmegaConf.create( + { + "inference_dir": str(inference_dir), + "dataset": {"test": [{"name": "from_cfg"}]}, + "metrics": [{"metric": {"_target_": f"{__name__}.DummyMetric"}}], + } + ) + + assert _resolve_test_sets(cfg) == ["from_cfg"] + + +def test_resolve_test_sets_ignores_hidden_directories(tmp_path): + inference_dir = tmp_path / "infer" + (inference_dir / ".ignored").mkdir(parents=True) + (inference_dir / "test_visible").mkdir(parents=True) + + cfg = OmegaConf.create( + { + "inference_dir": str(inference_dir), + "metrics": [{"metric": {"_target_": f"{__name__}.DummyMetric"}}], + } + ) + + assert _resolve_test_sets(cfg) == ["test_visible"] + + def test_metric_rejects_non_metric_instance(tmp_path): cfg = OmegaConf.create( { @@ -106,3 +161,16 @@ def test_metric_requires_inputs_when_metric_has_no_keys(tmp_path): with pytest.raises(ValueError, match="requires inputs in config"): measure(cfg) + + +def test_metric_requires_test_sets_from_config_or_inference_dir(tmp_path): + cfg = OmegaConf.create( + { + "inference_dir": str(tmp_path / "infer"), + "metrics": [{"metric": {"_target_": f"{__name__}.DummyMetric"}}], + } + ) + + (tmp_path / "infer").mkdir() + with pytest.raises(ValueError, match="No test sets found"): + measure(cfg) diff --git a/test/espnet3/utils/test_config.py b/test/espnet3/utils/test_config.py index 9b723da1acd4..45f1f70fed92 100644 --- a/test/espnet3/utils/test_config.py +++ b/test/espnet3/utils/test_config.py @@ -1,9 +1,18 @@ import os import pytest +from omegaconf import OmegaConf from yaml.parser import ParserError -from espnet3.utils.config_utils import load_config_with_defaults, load_line +from espnet3.utils.config_utils import ( + _build_config_path, + _ensure_target_convert_all, + load_and_merge_config, + load_config_with_defaults, + load_line, + load_template_defaults, + load_yaml, +) # =============================================================== # Test Case Summary for Config Utilities @@ -241,6 +250,28 @@ def test_defaults_removed_after_merge(write_yaml): assert "defaults" not in cfg +def test_load_yaml_relative_path_survives_merge(write_yaml): + write_yaml( + "conf/training.yaml", + """ +exp_dir: ./exp/train_debug +dataset_dir: ./data/mini_an4 +""", + ) + path = write_yaml( + "conf/inference.yaml", + """ +exp_dir: ${load_yaml:training.yaml,exp_dir} +dataset_dir: ${load_yaml:training.yaml,dataset_dir} +""", + ) + cfg = load_config_with_defaults(str(path)) + merged = OmegaConf.merge(OmegaConf.create({"provider": {"name": "dummy"}}), cfg) + + assert merged.exp_dir == "./exp/train_debug" + assert merged.dataset_dir == "./data/mini_an4" + + def test_missing_file_raises(tmp_path): """Test that loading a config with a missing file in `defaults` @@ -277,3 +308,138 @@ def test_invalid_yaml_raises(tmp_path): with pytest.raises(ParserError): load_config_with_defaults(str(main_path)) + + +def test_load_yaml_full_config(write_yaml): + path = write_yaml( + "config.yaml", + """ +foo: + bar: 123 +""", + ) + cfg = load_yaml(str(path)) + assert cfg.foo.bar == 123 + + +def test_load_yaml_nested_key(write_yaml): + path = write_yaml( + "config.yaml", + """ +foo: + bar: 123 +""", + ) + value = load_yaml(str(path), "foo.bar") + assert value == 123 + + +def test_load_yaml_missing_key_raises(write_yaml): + path = write_yaml("config.yaml", "foo: 1\n") + with pytest.raises(KeyError): + load_yaml(str(path), "missing.key") + + +def test_load_template_defaults_train(): + cfg = load_template_defaults("conf/training.yaml", "egs3.TEMPLATE.asr") + assert "dataset" in cfg + assert "exp_dir" in cfg + + +def test_load_and_merge_config_none(): + assert load_and_merge_config(None, "conf/metrics.yaml") is None + + +def test_load_and_merge_config_resolves_user_reference_to_template_value( + write_yaml, monkeypatch +): + template = write_yaml( + "template.yaml", + """ +exp_dir: ./exp/from_template +inference_dir: ${exp_dir}/inference +""", + ) + user = write_yaml( + "user.yaml", + """ +custom_dir: ${exp_dir}/custom +""", + ) + + monkeypatch.setattr( + "espnet3.utils.config_utils.load_template_defaults", + lambda _, __: load_config_with_defaults(str(template), resolve=False), + ) + + cfg = load_and_merge_config( + user, + "conf/training.yaml", + template_package="dummy.package", + ) + + assert cfg.exp_dir == "./exp/from_template" + assert cfg.inference_dir == "./exp/from_template/inference" + assert cfg.custom_dir == "./exp/from_template/custom" + + +def test_load_and_merge_config_resolves_template_reference_to_user_override( + write_yaml, monkeypatch +): + template = write_yaml( + "template.yaml", + """ +exp_dir: ./exp/from_template +custom_dir: ${exp_dir}/custom +""", + ) + user = write_yaml( + "user.yaml", + """ +exp_dir: ./exp/from_user +""", + ) + + monkeypatch.setattr( + "espnet3.utils.config_utils.load_template_defaults", + lambda _, __: load_config_with_defaults(str(template), resolve=False), + ) + + cfg = load_and_merge_config( + user, + "conf/training.yaml", + template_package="dummy.package", + ) + + assert cfg.exp_dir == "./exp/from_user" + assert cfg.custom_dir == "./exp/from_user/custom" + + +def test_ensure_target_convert_all_nested(): + cfg = OmegaConf.create( + { + "a": {"_target_": "package.A", "x": 1}, + "b": [{"_target_": "package.B", "y": 2}, {"z": 3}], + "c": {"nested": {"_target_": "package.C"}}, + } + ) + _ensure_target_convert_all(cfg) + assert cfg.a._convert_ == "all" + assert cfg.b[0]._convert_ == "all" + assert cfg.c.nested._convert_ == "all" + assert "_convert_" not in cfg.b[1] + + +def test_build_config_path_adds_yaml(tmp_path): + path = _build_config_path(tmp_path, "train") + assert path == tmp_path / "train.yaml" + + +def test_build_config_path_keeps_yaml(tmp_path): + path = _build_config_path(tmp_path, "train.yaml") + assert path == tmp_path / "train.yaml" + + +def test_build_config_path_rejects_suffix(tmp_path): + with pytest.raises(ValueError): + _build_config_path(tmp_path, "train.txt") diff --git a/test/espnet3/utils/test_stages.py b/test/espnet3/utils/test_stages.py index bb295cb17930..764790327722 100644 --- a/test/espnet3/utils/test_stages.py +++ b/test/espnet3/utils/test_stages.py @@ -1,3 +1,4 @@ +import argparse import logging import pytest @@ -9,6 +10,7 @@ class DummySystem: def __init__(self): self.calls = [] + self.stage_log_dirs = {"default": None} def stage_a(self): self.calls.append("a") @@ -30,7 +32,11 @@ def test_resolve_stages_subset_preserves_stage_order(): def test_run_stages_dry_run_skips_execution(caplog): system = DummySystem() with caplog.at_level(logging.INFO): - run_stages(system, ["stage_a", "stage_b"], dry_run=True) + run_stages( + system, + ["stage_a", "stage_b"], + args=argparse.Namespace(dry_run=True), + ) assert system.calls == [] assert "[DRY RUN] would run stage: stage_a" in caplog.text @@ -64,3 +70,38 @@ def stage_a(self): system = CrashSystem() with pytest.raises(ValueError, match="boom"): run_stages(system, ["stage_a"]) + + +def test_run_stages_writes_stage_logs(tmp_path): + class LoggingSystem: + def __init__(self, log_dir): + self.log_dir = log_dir + self.stage_log_dirs = {"default": log_dir} + + def stage_a(self): + pass + + def stage_b(self): + pass + + system = LoggingSystem(tmp_path) + logger = logging.getLogger("espnet3.test_stages") + root = logging.getLogger() + previous_level = root.level + root.setLevel(logging.INFO) + try: + run_stages(system, ["stage_a", "stage_b"], log=logger) + finally: + root.setLevel(previous_level) + + stage_a_log = tmp_path / "stage_a.log" + stage_b_log = tmp_path / "stage_b.log" + assert stage_a_log.exists() + assert stage_b_log.exists() + + stage_a_text = stage_a_log.read_text(encoding="utf-8") + stage_b_text = stage_b_log.read_text(encoding="utf-8") + assert "stage: stage_a" in stage_a_text + assert "stage: stage_b" not in stage_a_text + assert "stage: stage_b" in stage_b_text + assert "stage: stage_a" not in stage_b_text