Nida v0.5 (research preview): weights, 16 voices, code, model card

c0bce59 verified 23 days ago

9.49 kB

"""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}"

	"""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}"