"""Isolated ``uv`` environments for the Phase 1 data-pipeline model workers. Same mechanism as :mod:`nida.bench.envs` (this module only *describes* environments; build them with :func:`nida.bench.envs.ensure_env`, which writes the freeze file and the marker): one environment per model family because their dependency stacks conflict (NeMo, FunASR/ModelScope, transformers, panns-inference/librosa). Python and the torch pins are taken from the Phase 0 ``hf`` environment, so every worker environment runs the same ``torch==2.8.0`` (cu128) build. ``EnvSpec.workers`` here are module names under ``nida.data.workers`` (see :func:`worker_module_name`), not under ``nida.bench.workers``. Version pins (PyPI, looked up 2026-10-08): * ``panns-inference==0.1.1`` (latest, 2023-03-26) with ``torchlibrosa==0.1.0`` (latest). * ``nemo_toolkit[asr]==3.0.0`` (latest, 2026-08-07; ``torch>=2.6.0`` so the torch 2.8.0 wheel is kept; the build's torch guard would catch a replacement). The model card asks for ``Cython`` and ``packaging`` to be installed first. System packages ``libsndfile1`` and ``ffmpeg`` must exist (present on Colab). * ``funasr==1.4.16`` (latest, 2026-09-18). * ``speechbrain==1.1.1`` (2026-08-27; the latest, 1.1.2, was published on 2026-10-14 and is not pinned yet: one day old). It requires ``torch>=2.1.0`` (the 2.8.0 wheel is kept), and its dependency ``hyperpyyaml==1.2.3`` (pinned explicitly) caps ``ruamel.yaml<0.19`` (0.19.x breaks HyperPyYAML with ``'Loader' object has no attribute 'max_depth'``), which the resolver respects. Importing ``speechbrain`` needs ``torchaudio`` (in the torch pins). Checked locally on CPU with torch and torchaudio 2.11.0 (not with the cu128 2.8.0 pair: unverified). * ``transformers==4.57.6``: **not** the latest (5.19.0, 2026-10-06). That is a new major version never exercised here; 4.57.6 is what the Phase 0 ``hf`` environment already ran with torch 2.8.0. Unpinned helpers (numpy, soundfile, huggingface-hub, librosa, matplotlib) follow the bench environments; the exact resolved versions are recorded in each environment's ``freeze.txt``. The speech-enhancement arms (:mod:`nida.data.enhance`) have their own mapping, :data:`ENHANCE_ENV_SPECS`, because their pins deliberately differ from the shared stack above (Python 3.11 for ``dfn``, numpy 1 vs 2, no torch pin for ``uvr``). They are the pins of ``scripts/enhance_envs.sh`` (exercised for the A/B listening test, 2026-10-13), copied unchanged; ``tests/test_data_enhance.py`` compares both. Same ``EnvSpec`` mechanism, so ``ensure_env`` builds them; :func:`data_env_for_worker` looks in both mappings. """ from __future__ import annotations from collections.abc import Mapping from types import MappingProxyType from nida.bench.envs import ENV_SPECS as BENCH_ENV_SPECS from nida.bench.envs import EnvSpec __all__ = [ "DATA_ENV_SPECS", "DATA_WORKER_PACKAGE", "ENHANCE_ENV_SPECS", "data_env_for_worker", "worker_module_name", ] #: Package whose modules are the Phase 1 workers. DATA_WORKER_PACKAGE = "nida.data.workers" _BASE = BENCH_ENV_SPECS["hf"] # python + torch pins shared with Phase 0 DATA_ENV_SPECS: Mapping[str, EnvSpec] = MappingProxyType( { "data_sed": EnvSpec( name="data_sed", python=_BASE.python, torch_packages=_BASE.torch_packages, torch_index_url=_BASE.torch_index_url, packages=( "panns-inference==0.1.1", "torchlibrosa==0.1.0", "librosa", "matplotlib", "numpy", "soundfile", ), workers=("sed_panns",), ), "data_diar": EnvSpec( name="data_diar", python=_BASE.python, torch_packages=_BASE.torch_packages, torch_index_url=_BASE.torch_index_url, packages=( "Cython", "packaging", "nemo_toolkit[asr]==3.0.0", # Unpinned, the resolver picks transformers 4.12 + tokenizers 0.10 (no wheel # for py3.12, Rust build fails on Colab); pin the Phase 0 transformers. "transformers==4.57.6", "huggingface-hub<1.0", "numpy", "soundfile", ), workers=("diar_sortformer",), ), "data_ctc": EnvSpec( name="data_ctc", python=_BASE.python, torch_packages=_BASE.torch_packages, torch_index_url=_BASE.torch_index_url, packages=("transformers==4.57.6", "huggingface-hub", "numpy", "soundfile"), workers=("ctc_emissions",), ), "data_emotion": EnvSpec( name="data_emotion", python=_BASE.python, torch_packages=_BASE.torch_packages, torch_index_url=_BASE.torch_index_url, packages=( "funasr==1.4.16", "transformers==4.57.6", # same resolver trap as data_diar "huggingface-hub<1.0", "numpy", "soundfile", ), workers=("emotion_e2v",), ), # Speaker similarity (evaluation, scripts/speaker_sim.py): SpeechBrain ECAPA-TDNN. # scipy: SpeechBrain depends on it and the script resamples with it. "data_spk": EnvSpec( name="data_spk", python=_BASE.python, torch_packages=_BASE.torch_packages, torch_index_url=_BASE.torch_index_url, packages=( "speechbrain==1.1.1", # The version SpeechBrain 1.1.1 was checked with; it caps ruamel.yaml<0.19. "hyperpyyaml==1.2.3", "huggingface-hub", "numpy", "scipy", "soundfile", ), workers=("spk_ecapa",), ), # S9 latents: the B2 VAE stack (torch + torchaudio for resampling), without the # Audiobox quality model. "data_vae": EnvSpec( name="data_vae", python=_BASE.python, torch_packages=BENCH_ENV_SPECS["audio"].torch_packages, torch_index_url=_BASE.torch_index_url, packages=("numpy", "pydantic", "soundfile", "huggingface-hub"), workers=("vae_encode", "vae_decode"), ), } ) #: scripts/enhance_envs.sh, copied: ``torch==2.8.0`` + ``torchaudio==2.8.0`` is the pair of the #: sidon/dfn venvs (PyPI's Linux 2.8.0 wheel is the CUDA 12.8 build, as is the cu128 index #: wheel installed here). ``uvr`` pins no torch: ``audio-separator[gpu]==0.47.0`` resolves its #: own (torch 2.14.1 and onnxruntime-gpu 1.31.0, CUDA 13 builds, NVIDIA driver >= 580; this #: combination was NOT exercised on a GPU: unverified, and to be pinned (torch, onnxruntime-gpu) #: from the first GPU run's freeze file) and ``ensure_env`` skips its torch step for an empty #: ``torch_packages``. ``[gpu]`` is the extra the script picks when ``nvidia-smi`` exists; if #: onnxruntime cannot activate the CUDA provider the worker aborts (it never runs on the CPU #: silently) and the A/B script's ``ENH_UVR_EXTRA=cpu`` variant is the fallback. ENHANCE_ENV_SPECS: Mapping[str, EnvSpec] = MappingProxyType( { "data_enh_uvr": EnvSpec( name="data_enh_uvr", python="3.12", torch_packages=(), torch_index_url="", packages=( "audio-separator[gpu]==0.47.0", "audioread==3.1.0", "numpy==2.5.4", "scipy==1.18.1", "soundfile==0.14.0", ), workers=("enhance_uvr",), ), "data_enh_sidon": EnvSpec( name="data_enh_sidon", python="3.12", torch_packages=_BASE.torch_packages, torch_index_url=_BASE.torch_index_url, packages=( "transformers==5.19.0", "huggingface_hub==2.2.0", "numpy==2.5.4", "scipy==1.18.1", "soundfile==0.14.0", ), workers=("enhance_sidon",), ), "data_enh_dfn": EnvSpec( name="data_enh_dfn", python="3.11", # deepfilterlib 0.5.6 has wheels for cp38..cp311 only torch_packages=_BASE.torch_packages, torch_index_url=_BASE.torch_index_url, packages=( "deepfilternet==0.5.6", "numpy==1.26.4", "scipy==1.17.1", "soundfile==0.14.0", ), workers=("enhance_dfn",), ), } ) def data_env_for_worker(worker: str) -> EnvSpec: """Return the environment spec that runs the Phase 1 ``worker``. Raises: ValueError: if no data (or enhancement) environment lists ``worker``. """ specs = (*DATA_ENV_SPECS.values(), *ENHANCE_ENV_SPECS.values()) for spec in specs: if worker in spec.workers: return spec known = sorted(w for spec in specs for w in spec.workers) raise ValueError(f"unknown data worker {worker!r}; known workers: {known}") def worker_module_name(worker: str) -> str: """Dotted module to run with ``python -m`` for ``worker``. Raises: ValueError: if ``worker`` is not a known data worker. """ data_env_for_worker(worker) # validates the name return f"{DATA_WORKER_PACKAGE}.{worker}"