diff --git a/nemo/collections/common/tokenizers/text_to_speech/tts_tokenizers.py b/nemo/collections/common/tokenizers/text_to_speech/tts_tokenizers.py index a571415f3123..42c51dd4c1e9 100644 --- a/nemo/collections/common/tokenizers/text_to_speech/tts_tokenizers.py +++ b/nemo/collections/common/tokenizers/text_to_speech/tts_tokenizers.py @@ -19,8 +19,11 @@ import warnings from abc import ABC, abstractmethod from contextlib import contextmanager -from typing import List, Optional, Union +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Mapping, Optional, Union +from omegaconf import open_dict +from packaging.version import InvalidVersion, Version from tokenizers import Tokenizer from transformers import PreTrainedTokenizerBase @@ -44,14 +47,128 @@ from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec from nemo.utils import logging -CASELESS_SCRIPT_TOKENIZER_TARGETS = frozenset( - { - 'nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.HindiCharsTokenizer', - 'nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.ArabicCharsTokenizer', - } -) +_TOKENIZER_MODULE = 'nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers' +_HINDI_CHARS_TOKENIZER_TARGET = f'{_TOKENIZER_MODULE}.HindiCharsTokenizer' +_ARABIC_CHARS_TOKENIZER_TARGET = f'{_TOKENIZER_MODULE}.ArabicCharsTokenizer' +_IPA_TOKENIZER_TARGET = f'{_TOKENIZER_MODULE}.IPATokenizer' DEFAULT_CHARSET_VERSION = 2 +DEFAULT_PUNCT_VERSION = 2 +DEFAULT_LOCALE_SPECIFIC_PUNCT = True + + +@dataclass(frozen=True) +class VersionedTokenizerField: + """A tokenizer default that changed after checkpoints had already been trained and released. + + A config that leaves ``field`` unset is ambiguous: it is either an archive that predates the field + (whose vocabulary was built with ``legacy``) or a config written since (which means ``current``). + ``changed_in`` is what dates it, against the config's own ``nemo_version`` stamp. + + Attributes: + field: The config key this entry resolves. + current: Value in force today. Must *be* the tokenizer's own default, not a copy of it. + legacy: Value in force before ``changed_in``. + changed_in: NeMo release that introduced ``current``. + applies_to: Whether the field is meaningful for a given tokenizer config node. Must exclude + tokenizers that do not accept the argument, and configs that already fix the vocabulary by + other means (an explicit ``chars`` makes ``charset_version`` inert). + """ + + field: str + current: Any + legacy: Any + changed_in: str + applies_to: Callable[[Mapping], bool] + + +# Adding a versioned field, or flipping a default, is one new row. Both PRs that introduced these +# (#15567, #15614) landed during 2.8.0rc0, each shipping the current default from day one. pt-BR is the +# only locale with released checkpoints built on DEFAULT_PUNCTUATION instead of its own punctuation. +# Test unset-ness with ``.get(...) is None``, never ``not in``: an explicit ``chars: null`` must read as +# unset, since that is how the tokenizers themselves gate their version branches. +VERSIONED_TOKENIZER_FIELDS = ( + VersionedTokenizerField( + field='charset_version', + current=DEFAULT_CHARSET_VERSION, + legacy=1, + changed_in="2.8.0", + applies_to=lambda cfg: cfg.get('_target_') in (_HINDI_CHARS_TOKENIZER_TARGET, _ARABIC_CHARS_TOKENIZER_TARGET) + and cfg.get('chars') is None, + ), + VersionedTokenizerField( + field='punct_version', + current=DEFAULT_PUNCT_VERSION, + legacy=1, + changed_in="2.8.0", + applies_to=lambda cfg: cfg.get('_target_') == _HINDI_CHARS_TOKENIZER_TARGET + and cfg.get('non_default_punct_list') is None, + ), + VersionedTokenizerField( + field='locale_specific_punct', + current=DEFAULT_LOCALE_SPECIFIC_PUNCT, + legacy=False, + changed_in="2.8.0", + applies_to=lambda cfg: cfg.get('_target_') == _IPA_TOKENIZER_TARGET + and cfg.get('locale') == "pt-BR" + and cfg.get('non_default_punct_list') is None, + ), +) + + +def _predates_nemo_release(cfg_nemo_version: Optional[str], changed_in: str) -> bool: + """Whether a config stamped ``cfg_nemo_version`` was authored before ``changed_in``. + + An unstamped config was hand-authored for fresh training, not serialized by ``ModelPT``, so it gets + today's defaults. Comparison is on ``Version.release`` so "2.8.0rc0" reads as 2.8.0 rather than as + older than it (PEP 440 orders pre-releases first). That resolves the entire 2.8.0rc0 window to the + current defaults even though the fields landed mid-window -- genuinely undecidable from the stamp, + tie-broken toward current because ArabicCharsTokenizer was added by the very commit that added + ``charset_version`` (so for Arabic this reading is always right), and because guessing legacy + downgrades new training silently while guessing current fails loudly. Released checkpoints from + that window pin their fields anyway, so none of them reach here. + """ + if cfg_nemo_version is None: + return False + try: + return Version(str(cfg_nemo_version)).release < Version(changed_in).release + except InvalidVersion: + logging.warning(f"Unparseable nemo_version={cfg_nemo_version!r}; assuming current tokenizer defaults.") + return False + + +def resolve_versioned_tokenizer_defaults(tokenizer_config, cfg_nemo_version: Optional[str] = None) -> Dict[str, Any]: + """Pin every unset versioned field of one tokenizer config node, in place. + + Pinning the values rather than leaning on the class defaults is what makes a vocabulary survive a + ``.nemo`` round-trip, and leaves the config unambiguous for downstream ``setup_tokenizers`` calls. + + Returns the fields that were backfilled, mapped to the values chosen for them. + """ + if tokenizer_config.get('_target_') is None: + return {} + + backfilled, legacy = {}, {} + for spec in VERSIONED_TOKENIZER_FIELDS: + if tokenizer_config.get(spec.field) is not None or not spec.applies_to(tokenizer_config): + continue + is_legacy = _predates_nemo_release(cfg_nemo_version, spec.changed_in) + value = spec.legacy if is_legacy else spec.current + with open_dict(tokenizer_config): + tokenizer_config[spec.field] = value + backfilled[spec.field] = value + if is_legacy: + legacy[spec.field] = value + + # Only the legacy direction is worth reporting: it silently reproduces an older vocabulary. + if legacy: + settings = ", ".join(f"{field}={value}" for field, value in legacy.items()) + logging.warning( + f"{tokenizer_config['_target_'].rsplit('.', 1)[-1]}: assuming {settings} from " + f"nemo_version={cfg_nemo_version}, which predates those fields. This reproduces the vocabulary " + f"the checkpoint was trained with; set them explicitly to opt into the current defaults." + ) + return backfilled class BaseTokenizer(ABC): @@ -429,21 +546,21 @@ class HindiCharsTokenizer(BaseCharsTokenizer): def __init__( self, chars=None, - charset_version=2, + charset_version=DEFAULT_CHARSET_VERSION, punct=True, apostrophe=True, add_blank_at=None, pad_with_space=False, non_default_punct_list=None, - punct_version=2, + punct_version=DEFAULT_PUNCT_VERSION, text_preprocessing_func=any_locale_text_preprocessing, ): if chars is None: if charset_version == 1: warnings.warn( - "HindiCharsTokenizer charset_version=1 (case='mixed' + ascii_lowercase) is deprecated " - "and will be removed in a future release. " - "Migrate to charset_version=2 (case='upper' + ascii_letters) and retrain.", + "HindiCharsTokenizer charset_version=1 (case='mixed' + ascii_lowercase) is frozen: it is " + "retained so that checkpoints trained on it stay loadable, and receives no further " + "changes. New models should use charset_version=2 (case='upper' + ascii_letters).", DeprecationWarning, stacklevel=2, ) @@ -457,8 +574,8 @@ def __init__( if non_default_punct_list is None: if punct_version == 1: warnings.warn( - "HindiCharsTokenizer: punct_version=1 uses DEFAULT_PUNCTUATION without dandas " - "and will be removed in a future release. Migrate to punct_version=2.", + "HindiCharsTokenizer: punct_version=1 uses DEFAULT_PUNCTUATION without dandas. It is frozen " + "so that checkpoints trained on it stay loadable; new models should use punct_version=2.", DeprecationWarning, stacklevel=2, ) @@ -543,7 +660,7 @@ class ArabicCharsTokenizer(BaseCharsTokenizer): def __init__( self, chars=None, - charset_version=2, + charset_version=DEFAULT_CHARSET_VERSION, punct=True, apostrophe=True, add_blank_at=None, @@ -554,9 +671,9 @@ def __init__( if chars is None: if charset_version == 1: warnings.warn( - "ArabicCharsTokenizer charset_version=1 (case='mixed' + ascii_letters) is deprecated " - "and will be removed in a future release. " - "Migrate to charset_version=2 (case='upper' + ascii_letters) and retrain.", + "ArabicCharsTokenizer charset_version=1 (case='mixed' + ascii_letters) is frozen: it is " + "retained so that checkpoints trained on it stay loadable (the released v2607 MagpieTTS " + "pins it), and receives no further changes. New models should use charset_version=2.", DeprecationWarning, stacklevel=2, ) @@ -935,7 +1052,7 @@ def __init__( locale="en-US", punct=True, non_default_punct_list=None, - locale_specific_punct=True, + locale_specific_punct=DEFAULT_LOCALE_SPECIFIC_PUNCT, fixed_vocab=None, *, space=' ', diff --git a/nemo/collections/tts/data/text_to_speech_dataset_lhotse.py b/nemo/collections/tts/data/text_to_speech_dataset_lhotse.py index c9baaaa96324..33575d0710a4 100644 --- a/nemo/collections/tts/data/text_to_speech_dataset_lhotse.py +++ b/nemo/collections/tts/data/text_to_speech_dataset_lhotse.py @@ -20,14 +20,13 @@ import torch from lhotse import CutSet from lhotse.dataset.collation import collate_matrices, collate_vectors -from omegaconf import DictConfig, open_dict +from omegaconf import DictConfig from transformers import AutoTokenizer, T5Tokenizer from nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers import ( - CASELESS_SCRIPT_TOKENIZER_TARGETS, - DEFAULT_CHARSET_VERSION, AggregatedTTSTokenizer, IPABPETokenizer, + resolve_versioned_tokenizer_defaults, ) from nemo.collections.tts.parts.utils.tts_dataset_utils import ( beta_binomial_prior_distribution, @@ -40,9 +39,23 @@ from nemo.utils import logging -def setup_tokenizers(all_tokenizers_config, mode='train'): - # Being used in both model and worker_init_fn, so it is defined here - # Returns two tokenizers: one for TTS transcript and one for conditioning text (if needed) +def setup_tokenizers(all_tokenizers_config, mode='train', cfg_nemo_version=None): + """Instantiate the aggregated TTS transcript tokenizer described by ``all_tokenizers_config``. + + Being used in both model and worker_init_fn, so it is defined here. + + Args: + all_tokenizers_config: The ``text_tokenizers`` config node. Mutated in place so that the + versioned tokenizer defaults it resolved to are persisted into any ``.nemo`` saved later. + mode: 'train' or 'test'. 'test' forces phoneme probability to 1.0 where supported. + cfg_nemo_version: The enclosing model config's ``nemo_version``, used to date any versioned + tokenizer field the config leaves unset (see ``resolve_versioned_tokenizer_defaults``). + Only a model's ``__init__`` needs it: that call pins the resolved values into the config, + so later calls -- dataloader setup, worker_init_fn -- can leave it None. + + Returns: + An ``AggregatedTTSTokenizer`` over every configured tokenizer. + """ tokenizers = [] tokenizer_names = [] for tokenizer_name in all_tokenizers_config: @@ -55,44 +68,12 @@ def setup_tokenizers(all_tokenizers_config, mode='train'): text_tokenizer_kwargs = {} if "g2p" in tokenizer_config: text_tokenizer_kwargs["g2p"] = safe_instantiate(tokenizer_config.g2p) - # Ensure locale_specific_punct is persisted so it survives .nemo save/restore. - # New training for locales with extended punctuation should use the full set (True). - if ( - hasattr(tokenizer_config, '_target_') - and tokenizer_config._target_ - == "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPATokenizer" - and tokenizer_config.get('locale', None) == "pt-BR" - and not hasattr(tokenizer_config, 'non_default_punct_list') - and not hasattr(tokenizer_config, 'locale_specific_punct') - ): - with open_dict(tokenizer_config): - tokenizer_config.locale_specific_punct = True - # Persist punct_version=2 for HindiCharsTokenizer so .nemo save/restore - # always uses the expanded punctuation set (with dandas). - if ( - hasattr(tokenizer_config, '_target_') - and tokenizer_config._target_ - == "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.HindiCharsTokenizer" - and not hasattr(tokenizer_config, 'punct_version') - ): - with open_dict(tokenizer_config): - tokenizer_config.punct_version = 2 + resolve_versioned_tokenizer_defaults(tokenizer_config, cfg_nemo_version) tokenizer = safe_instantiate(tokenizer_config, **text_tokenizer_kwargs) # TODO @xueyang: is it really necessary to set phone probability to 1.0 for test mode? if mode == 'test' and hasattr(tokenizer, "set_phone_prob"): tokenizer.set_phone_prob(1.0) - # Persist charset_version so it's saved in .nemo archives and - # update_config_for_inference can distinguish old checkpoints - # (missing charset_version → v1) from new ones. - if ( - hasattr(tokenizer_config, '_target_') - and tokenizer_config._target_ in CASELESS_SCRIPT_TOKENIZER_TARGETS - and not hasattr(tokenizer_config, 'charset_version') - ): - with open_dict(all_tokenizers_config): - tokenizer_config.charset_version = DEFAULT_CHARSET_VERSION - tokenizers.append(tokenizer) tokenizer_names.append(tokenizer_name) @@ -101,6 +82,34 @@ def setup_tokenizers(all_tokenizers_config, mode='train'): return aggregated_tokenizer +def check_text_embedding_matches_tokenizer(state_dict, *, text_embedding, tokenizer, model_cfg) -> None: + """Fail actionably when a checkpoint's text embedding disagrees with the rebuilt tokenizer. + + This is the symptom of every tokenizer-versioning mistake: the vocabulary is rebuilt with different + character or punctuation sets than the checkpoint was trained with, every token ID shifts, and + ``load_state_dict`` reports an opaque size mismatch. Naming the per-tokenizer token counts and the + fields to pin turns that into something a user can act on. A no-op when either side is absent -- + the CAS-encoder MagpieTTS variant has no text embedding table. + """ + ckpt_weight = state_dict.get('text_embedding.weight', None) + if ckpt_weight is None or text_embedding is None: + return + if ckpt_weight.shape[0] == text_embedding.num_embeddings: + return + + raise RuntimeError( + f"MagpieTTS tokenizer/checkpoint mismatch: the checkpoint's text_embedding has " + f"{ckpt_weight.shape[0]} rows but this config builds {text_embedding.num_embeddings}. The text " + f"tokenizer vocabulary was not rebuilt the way this checkpoint was trained.\n" + f" tokens per tokenizer: {getattr(tokenizer, 'num_tokens_per_tokenizer', 'n/a')}\n" + f" config nemo_version: {model_cfg.get('nemo_version', '')}\n" + f"This is usually a versioned tokenizer field resolving differently than at training time. Pin " + f"the values the checkpoint was trained with explicitly under `model.text_tokenizers.`: " + f"`charset_version` (Hindi/Arabic char tokenizers), `punct_version` (Hindi), or " + f"`locale_specific_punct` (pt-BR IPA). See `VERSIONED_TOKENIZER_FIELDS` in tts_tokenizers.py." + ) + + def check_speaker_format(item: str): # enforce the format as example like "| Language:en Dataset:HiFiTTS Speaker:9136_other |". pattern = r"\| Language:\w+ Dataset:[\w\d\W]+ Speaker:[\w\d\W]+ \|" diff --git a/nemo/collections/tts/models/easy_magpietts_inference.py b/nemo/collections/tts/models/easy_magpietts_inference.py index 5cc228f26ef9..580d00bcbdb8 100644 --- a/nemo/collections/tts/models/easy_magpietts_inference.py +++ b/nemo/collections/tts/models/easy_magpietts_inference.py @@ -26,7 +26,10 @@ from transformers import AutoConfig, AutoModelForCausalLM from nemo.collections.audio.parts.utils.transforms import resample -from nemo.collections.tts.data.text_to_speech_dataset_lhotse import setup_tokenizers +from nemo.collections.tts.data.text_to_speech_dataset_lhotse import ( + check_text_embedding_matches_tokenizer, + setup_tokenizers, +) from nemo.collections.tts.models import AudioCodecModel from nemo.collections.tts.modules import transformer_2501 from nemo.collections.tts.modules.audio_codec_modules import VectorQuantizerIndexConverter @@ -332,6 +335,8 @@ def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None): self.tokenizer = setup_tokenizers( all_tokenizers_config=cfg.text_tokenizers, mode='train', + # Read before super().__init__, which stamps the *current* version into configs that lack one. + cfg_nemo_version=cfg.get('nemo_version', None), ) num_tokens_tokenizer = len(self.tokenizer.tokens) @@ -604,6 +609,12 @@ def state_dict(self, destination=None, prefix='', keep_vars=False): return state_dict def load_state_dict(self, state_dict, strict=True): + check_text_embedding_matches_tokenizer( + state_dict, + text_embedding=getattr(self, 'text_embedding', None), + tokenizer=self.tokenizer, + model_cfg=self.cfg, + ) if not strict: super().load_state_dict(state_dict, strict=False) modules_to_skip = self._get_state_dict_keys_to_exclude() diff --git a/nemo/collections/tts/models/magpietts.py b/nemo/collections/tts/models/magpietts.py index 79018231c6bf..cc54a19f8cc7 100644 --- a/nemo/collections/tts/models/magpietts.py +++ b/nemo/collections/tts/models/magpietts.py @@ -33,7 +33,11 @@ from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict from torch import nn from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config -from nemo.collections.tts.data.text_to_speech_dataset_lhotse import MagpieTTSLhotseDataset, setup_tokenizers +from nemo.collections.tts.data.text_to_speech_dataset_lhotse import ( + MagpieTTSLhotseDataset, + check_text_embedding_matches_tokenizer, + setup_tokenizers, +) from nemo.collections.tts.losses.aligner_loss import ForwardSumLoss from nemo.collections.tts.losses.moe_loss import MoEAuxiliaryLoss, compute_expert_usage from nemo.collections.tts.models import AudioCodecModel @@ -434,6 +438,8 @@ def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None): self.tokenizer = setup_tokenizers( all_tokenizers_config=cfg.text_tokenizers, mode='train', + # Read before super().__init__, which stamps the *current* version into configs that lack one. + cfg_nemo_version=cfg.get('nemo_version', None), ) num_tokens_tokenizer = len(self.tokenizer.tokens) @@ -1071,6 +1077,13 @@ def load_state_dict(self, state_dict, strict=True): (N, T*D) and reconstructed to (N, T, D) at inference time using stored T and D dimensions. """ state_dict = self.update_ckpt(state_dict) + # `text_embedding` is absent on the CAS-encoder variant, which has no such table to compare. + check_text_embedding_matches_tokenizer( + state_dict, + text_embedding=getattr(self, 'text_embedding', None), + tokenizer=self.tokenizer, + model_cfg=self.cfg, + ) # Check if checkpoint has baked context embedding (nn.Embedding format) has_baked_embedding_in_ckpt = 'baked_context_embedding.weight' in state_dict diff --git a/nemo/collections/tts/modules/magpietts_inference/utils.py b/nemo/collections/tts/modules/magpietts_inference/utils.py index 337f3df1d0b3..8a28b15bdea7 100644 --- a/nemo/collections/tts/modules/magpietts_inference/utils.py +++ b/nemo/collections/tts/modules/magpietts_inference/utils.py @@ -28,7 +28,6 @@ import torch from omegaconf import DictConfig, OmegaConf, open_dict -from nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers import CASELESS_SCRIPT_TOKENIZER_TARGETS from nemo.collections.tts.models import EasyMagpieTTSInferenceModel, MagpieTTSModel from nemo.utils import logging @@ -152,55 +151,6 @@ def validate(self) -> None: ) -def _migrate_charset_version(model_cfg: DictConfig) -> None: - """Pin charset_version=1 for Hindi/Arabic tokenizers in old checkpoints. - - New models have ``charset_version`` persisted by ``setup_tokenizers()``. - Old checkpoints lack it, so without this migration the new default (v2) - would silently change the token-to-ID mapping and break the model. - - Must be called inside ``open_dict(model_cfg)``. - """ - if not hasattr(model_cfg, 'text_tokenizers'): - return - for tok_name in model_cfg.text_tokenizers: - tok_cfg = model_cfg.text_tokenizers[tok_name] - if hasattr(tok_cfg, '_target_') and tok_cfg._target_ in CASELESS_SCRIPT_TOKENIZER_TARGETS: - if not hasattr(tok_cfg, 'charset_version'): - tok_cfg.charset_version = 1 - - -def _migrate_tokenizer_punctuation(model_cfg: DictConfig) -> None: - """Backfill punctuation fields for tokenizers that predate them. - - Old checkpoints were trained with DEFAULT_PUNCTUATION only. Without these - migrations, restoring those checkpoints would pick up expanded defaults - from new code, adding extra punctuation tokens and breaking the vocabulary. - - Handles: - - pt-BR IPATokenizer: sets locale_specific_punct=False (old default had no guillemets/quotes). - - HindiCharsTokenizer: sets punct_version=1 (old default had no dandas). - """ - if not hasattr(model_cfg, 'text_tokenizers'): - return - for tok_name in model_cfg.text_tokenizers: - tok_cfg = model_cfg.text_tokenizers[tok_name] - if not hasattr(tok_cfg, '_target_'): - continue - if ( - tok_cfg._target_ == "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPATokenizer" - and tok_cfg.get('locale', None) == "pt-BR" - and not hasattr(tok_cfg, 'non_default_punct_list') - and not hasattr(tok_cfg, 'locale_specific_punct') - ): - tok_cfg.locale_specific_punct = False - if ( - tok_cfg._target_ == "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.HindiCharsTokenizer" - and not hasattr(tok_cfg, 'punct_version') - ): - tok_cfg.punct_version = 1 - - def update_config_for_inference( model_cfg: DictConfig, codecmodel_path: Optional[str], @@ -223,8 +173,9 @@ def update_config_for_inference( """ model_cfg.codecmodel_path = codecmodel_path - _migrate_tokenizer_punctuation(model_cfg) - _migrate_charset_version(model_cfg) + # The versioned tokenizer fields (charset_version / punct_version / locale_specific_punct) need no + # migration here: the model's __init__ resolves them from the config's own ``nemo_version`` stamp + # and pins the result, and every path below constructs the model from this config afterwards. # Update text tokenizer paths for backward compatibility if hasattr(model_cfg, 'text_tokenizer'): diff --git a/tests/collections/tts/models/test_magpietts_tokenizer_versioning.py b/tests/collections/tts/models/test_magpietts_tokenizer_versioning.py new file mode 100644 index 000000000000..2f25f0eca13b --- /dev/null +++ b/tests/collections/tts/models/test_magpietts_tokenizer_versioning.py @@ -0,0 +1,472 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Unit tests for the versioned tokenizer defaults that MagpieTTS configs carry. + +``HindiCharsTokenizer`` / ``ArabicCharsTokenizer`` (``charset_version``, ``punct_version``) and the +pt-BR ``IPATokenizer`` (``locale_specific_punct``) each gained a second character/punctuation set after +models had already been trained and released. A config that omits those fields is therefore ambiguous: +it is either an archive that predates them (and whose vocabulary was built with the v1 values) or a +fresh training config that should get today's defaults. These tests pin both readings: + +1. ``setup_tokenizers`` uses the current defaults for fresh configs, and the pre-versioning values only + for configs stamped with a release that predates the fields -- so new training is never silently + downgraded to v1. +2. The ``nemo_version`` stamp that ``ModelPT`` writes into every config it saves is what dates the + config, and ``VERSIONED_TOKENIZER_FIELDS`` records the release each default changed in. +3. Whatever the resolved values are, they are written back into the config, so a model saved today + restores to the same token-to-ID mapping no matter how the class defaults evolve later. +""" + +import inspect +from unittest.mock import MagicMock, patch + +import hydra +import pytest +import torch +from omegaconf import OmegaConf, open_dict +from packaging.version import Version + +from nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers import ( + VERSIONED_TOKENIZER_FIELDS, + resolve_versioned_tokenizer_defaults, +) +from nemo.collections.tts.data.text_to_speech_dataset_lhotse import ( + check_text_embedding_matches_tokenizer, + setup_tokenizers, +) +from nemo.core.classes import ModelPT + +_HINDI_CHARS = "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.HindiCharsTokenizer" +_ARABIC_CHARS = "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.ArabicCharsTokenizer" +_IPA = "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPATokenizer" + +# A release that predates the versioned fields (both v2512 and v2602 are stamped exactly this) and one +# that postdates them. "2.8.0rc0" is deliberately the *current* side: see ``_predates_nemo_release``. +LEGACY_VERSION = "2.6.0rc0" +CURRENT_VERSION = "2.8.0rc0" + + +def _tokenizers_cfg(**tokenizer_fields): + """A single-entry ``text_tokenizers`` config for the Hindi character tokenizer.""" + return OmegaConf.create({"hindi_chartokenizer": {"_target_": _HINDI_CHARS, **tokenizer_fields}}) + + +class _TokenizerVersionModel(ModelPT): + """Minimal stand-in for MagpieTTS that reproduces only its tokenizer/embedding sizing. + + It mirrors the two lines that matter -- resolve the tokenizer defaults according to the config's + provenance, then size the text embedding from the resulting vocabulary -- so a ``save_to`` / + ``restore_from`` round-trip exercises the real vocabulary-drift failure without building a codec, + encoders, or downloading a released archive. + """ + + def __init__(self, cfg, trainer=None): + self.tokenizer = setup_tokenizers( + all_tokenizers_config=cfg.text_tokenizers, + cfg_nemo_version=cfg.get('nemo_version', None), + ) + super().__init__(cfg=cfg, trainer=trainer) + self.text_embedding = torch.nn.Embedding(len(self.tokenizer.tokens) + 2, 4) + + def setup_training_data(self, train_data_config): + self._train_dl = None + + def setup_validation_data(self, val_data_config): + self._validation_dl = None + + def setup_test_data(self, test_data_config): + self._test_dl = None + + @classmethod + def list_available_models(cls): + return [] + + +class TestSetupTokenizersDefaults: + """Covers the plumbing from ``setup_tokenizers`` down to the resolver. Which value each stamp + resolves to is pinned by ``TestVersionDating``, and what each value *means* for the resulting + vocabulary by ``test_tts_tokenizers.py`` at the tokenizer-class level.""" + + @pytest.mark.unit + def test_version_reaches_the_tokenizer_config(self): + """The one line of plumbing between the model constructor and the resolver.""" + cfg = _tokenizers_cfg() + + setup_tokenizers(cfg, cfg_nemo_version=LEGACY_VERSION) + + assert cfg.hindi_chartokenizer.charset_version == 1 + assert cfg.hindi_chartokenizer.punct_version == 1 + + @pytest.mark.unit + @pytest.mark.parametrize("cfg_nemo_version", [CURRENT_VERSION, LEGACY_VERSION]) + def test_explicit_values_are_never_overridden(self, cfg_nemo_version): + """An explicitly configured version wins over both defaults -- that is what pins v2607.""" + cfg = _tokenizers_cfg(charset_version=1, punct_version=2) + + setup_tokenizers(cfg, cfg_nemo_version=cfg_nemo_version) + + assert cfg.hindi_chartokenizer.charset_version == 1 + assert cfg.hindi_chartokenizer.punct_version == 2 + + +class TestResolveVersionedTokenizerDefaults: + # The current direction for these two targets is covered by + # ``test_current_backfill_matches_the_tokenizer_class_default``, which asserts something stronger. + @pytest.mark.unit + def test_arabic_charset_version_dates_back_to_v1(self): + cfg = OmegaConf.create({"_target_": _ARABIC_CHARS}) + + resolve_versioned_tokenizer_defaults(cfg, LEGACY_VERSION) + + assert cfg.charset_version == 1 + + @pytest.mark.unit + def test_pt_br_locale_specific_punct_dates_back_to_off(self): + cfg = OmegaConf.create({"_target_": _IPA, "locale": "pt-BR"}) + + resolve_versioned_tokenizer_defaults(cfg, LEGACY_VERSION) + + assert cfg.locale_specific_punct is False + + @pytest.mark.unit + def test_non_default_punct_list_suppresses_pt_br_backfill(self): + """An explicit punctuation list already fixes the vocabulary; adding the flag would fight it.""" + cfg = OmegaConf.create({"_target_": _IPA, "locale": "pt-BR", "non_default_punct_list": [".", ","]}) + + resolve_versioned_tokenizer_defaults(cfg, LEGACY_VERSION) + + assert "locale_specific_punct" not in cfg + + @pytest.mark.unit + @pytest.mark.parametrize( + "cfg_fields, field, expected", + [ + ({"chars": None}, "charset_version", 1), + ({"non_default_punct_list": None}, "punct_version", 1), + ({"charset_version": None}, "charset_version", 1), + ], + ids=["null-chars", "null-punct-list", "null-field-itself"], + ) + def test_explicit_null_counts_as_unset(self, cfg_fields, field, expected): + """``chars: null`` must not read as "specified". + + The tokenizers gate their version branches on ``is None``, so a null suppressing the backfill + would leave the branch running against the class default rather than the dated value. + """ + cfg = OmegaConf.create({"_target_": _HINDI_CHARS, **cfg_fields}) + + resolve_versioned_tokenizer_defaults(cfg, LEGACY_VERSION) + + assert cfg[field] == expected + + @pytest.mark.unit + def test_other_ipa_locales_are_untouched(self): + """Only pt-BR's punctuation set diverged from DEFAULT_PUNCTUATION, so only it is backfilled.""" + cfg = OmegaConf.create({"_target_": _IPA, "locale": "es-ES"}) + + resolve_versioned_tokenizer_defaults(cfg, LEGACY_VERSION) + + assert "locale_specific_punct" not in cfg + + @pytest.mark.unit + @pytest.mark.parametrize("extra_cfg", [{}, {"_target_": None}], ids=["absent", "null"]) + def test_config_without_target_is_ignored(self, extra_cfg): + """HuggingFace tokenizer nodes carry no resolvable ``_target_`` and must pass through untouched.""" + cfg = OmegaConf.create({"pretrained_model": "google-t5/t5-small", **extra_cfg}) + before = OmegaConf.to_container(cfg) + + resolve_versioned_tokenizer_defaults(cfg, LEGACY_VERSION) + + assert OmegaConf.to_container(cfg) == before + + @pytest.mark.unit + @pytest.mark.parametrize( + "target, field, extra_cfg", + [ + (_HINDI_CHARS, "charset_version", {}), + (_HINDI_CHARS, "punct_version", {}), + (_ARABIC_CHARS, "charset_version", {}), + (_IPA, "locale_specific_punct", {"locale": "pt-BR"}), + ], + ) + def test_current_backfill_matches_the_tokenizer_class_default(self, target, field, extra_cfg): + """What gets persisted for a current config must equal what the tokenizer would have chosen itself. + + ``VERSIONED_TOKENIZER_FIELDS`` and the tokenizer signatures read the same ``DEFAULT_*`` constants, + so this holds by construction today. It is kept as the guard against someone re-introducing a + literal on either side, which would silently pin every newly trained model a version behind. + """ + class_default = inspect.signature(hydra.utils.get_class(target)).parameters[field].default + cfg = OmegaConf.create({"_target_": target, **extra_cfg}) + + resolve_versioned_tokenizer_defaults(cfg, CURRENT_VERSION) + + assert cfg[field] == class_default + + @pytest.mark.unit + @pytest.mark.parametrize( + "cfg_fields, cfg_nemo_version, should_warn", + [ + ({}, LEGACY_VERSION, True), + ({"_target_": _IPA, "locale": "pt-BR"}, LEGACY_VERSION, True), + ({}, CURRENT_VERSION, False), + ({"punct_version": 1, "charset_version": 1}, LEGACY_VERSION, False), + ], + ids=["legacy-hindi", "legacy-pt-br", "current", "already-explicit"], + ) + def test_legacy_backfill_is_never_silent(self, cfg_fields, cfg_nemo_version, should_warn): + """Falling back to a pre-versioning vocabulary must say so. + + The tokenizers cannot be relied on for this: the pt-BR path emits nothing at all, and the + Hindi/Arabic ``DeprecationWarning`` is swallowed by Python's default filters outside pytest. + Warning only on the legacy direction keeps ordinary training runs quiet. + """ + cfg = OmegaConf.create({"_target_": _HINDI_CHARS, **cfg_fields}) + + with patch("nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.logging.warning") as mock_warning: + resolve_versioned_tokenizer_defaults(cfg, cfg_nemo_version) + + assert mock_warning.called is should_warn + if should_warn: + # The message has to name the field, otherwise it is not actionable. + assert any(field in mock_warning.call_args.args[0] for field in ("punct_version", "locale_specific_punct")) + + +class TestVersionDating: + """How a config's ``nemo_version`` stamp maps onto the versioned fields.""" + + @pytest.mark.unit + @pytest.mark.parametrize( + "cfg_nemo_version, expect_legacy", + [ + ("2.6.0rc0", True), # v2512 and v2602 are both stamped exactly this + ("2.7.3", True), + ("2.8.0rc0", False), # v2607; the rc window is broken toward the current defaults + ("2.8.0", False), + ("3.1.0", False), + (None, False), # hand-authored fresh training config + ("not-a-version", False), + ], + ids=["v2512-v2602", "2.7.3", "v2607-rc", "2.8.0", "3.1.0", "unstamped", "garbage"], + ) + def test_stamp_decides_the_backfilled_value(self, cfg_nemo_version, expect_legacy): + cfg = OmegaConf.create({"_target_": _HINDI_CHARS}) + + resolve_versioned_tokenizer_defaults(cfg, cfg_nemo_version) + + assert cfg.charset_version == (1 if expect_legacy else 2) + + @pytest.mark.unit + @pytest.mark.parametrize("spec", VERSIONED_TOKENIZER_FIELDS, ids=lambda s: s.field) + def test_every_field_documents_a_parseable_release(self, spec): + """``changed_in`` is compared against real stamps, so it has to parse and differ from legacy.""" + assert Version(spec.changed_in).release + assert spec.current != spec.legacy + + +class _StopAtTokenizerSetup(Exception): + """Raised from the patched ``setup_tokenizers`` to end ``__init__`` once the call site has run.""" + + +def _mock_codec(): + """An AudioCodecModel stand-in with the numeric attributes ``__init__`` reads before tokenizer setup.""" + codec = MagicMock() + codec.sample_rate = 22050 + codec.output_sample_rate = 22050 + codec.samples_per_frame = 1024 + codec.num_codebooks = 8 + codec.codebook_size = 1000 + return codec + + +class TestProductionCallSites: + """The wiring in the real model constructors, which is the whole fix. + + Without this, mutating either call site (deleting the kwarg, or negating it) leaves the entire TTS + unit suite green while reintroducing the v2602 restore failure -- the tests below are the only thing + that fails on such a mutation, because every other test drives ``setup_tokenizers`` directly. + + Each constructor is stopped at the tokenizer-setup call rather than run to completion, so no codec, + encoders, or downloads are needed. + """ + + @staticmethod + def _captured_kwargs(model_cls, module_path, cfg_extra, codec_attr): + cfg = OmegaConf.create( + { + "codecmodel_path": "nvidia/fake-codec", + "text_tokenizers": _tokenizers_cfg(), + **cfg_extra, + } + ) + captured = {} + with ( + patch(f"{module_path}.AudioCodecModel") as mock_codec, + patch(f"{module_path}.setup_tokenizers") as mock_setup, + ): + getattr(mock_codec, codec_attr).return_value = _mock_codec() + + def _record(*_args, **kwargs): + captured.update(kwargs) + raise _StopAtTokenizerSetup() + + mock_setup.side_effect = _record + with pytest.raises(_StopAtTokenizerSetup): + model_cls(cfg=cfg) + return captured + + @pytest.mark.unit + @pytest.mark.parametrize( + "cfg_extra, expected", + [({}, None), ({"nemo_version": LEGACY_VERSION}, LEGACY_VERSION)], + ids=["fresh", "serialized"], + ) + def test_magpietts_passes_config_nemo_version(self, cfg_extra, expected): + from nemo.collections.tts.models.magpietts import MagpieTTSModel + + captured = self._captured_kwargs( + MagpieTTSModel, "nemo.collections.tts.models.magpietts", cfg_extra, "from_pretrained" + ) + + assert captured["cfg_nemo_version"] == expected + + @pytest.mark.unit + @pytest.mark.parametrize( + "cfg_extra, expected", + [({}, None), ({"nemo_version": LEGACY_VERSION}, LEGACY_VERSION)], + ids=["fresh", "serialized"], + ) + def test_easy_magpietts_passes_config_nemo_version(self, cfg_extra, expected): + from nemo.collections.tts.models.easy_magpietts_inference import EasyMagpieTTSInferenceModel + + captured = self._captured_kwargs( + EasyMagpieTTSInferenceModel, + "nemo.collections.tts.models.easy_magpietts_inference", + cfg_extra, + "restore_from", + ) + + assert captured["cfg_nemo_version"] == expected + + +class TestTextEmbeddingMismatchError: + """The diagnostic a user actually hits when a vocabulary is rebuilt the wrong way.""" + + @staticmethod + def _model_and_state_dict(ckpt_rows): + model = MagicMock() + model.num_tokens_per_tokenizer = {"hindi_chartokenizer": 191} + return {"text_embedding.weight": torch.zeros(ckpt_rows, 4)}, model + + @pytest.mark.unit + def test_matching_sizes_pass(self): + state_dict, tokenizer = self._model_and_state_dict(100) + + check_text_embedding_matches_tokenizer( + state_dict, text_embedding=torch.nn.Embedding(100, 4), tokenizer=tokenizer, model_cfg=OmegaConf.create({}) + ) + + @pytest.mark.unit + def test_mismatch_names_the_fields_to_pin(self): + state_dict, tokenizer = self._model_and_state_dict(193) + + with pytest.raises(RuntimeError) as excinfo: + check_text_embedding_matches_tokenizer( + state_dict, + text_embedding=torch.nn.Embedding(148, 4), + tokenizer=tokenizer, + model_cfg=OmegaConf.create({"nemo_version": LEGACY_VERSION}), + ) + + message = str(excinfo.value) + assert "193" in message and "148" in message # both sides of the mismatch + assert LEGACY_VERSION in message # the stamp that drove the decision + assert "hindi_chartokenizer" in message # which tokenizer to look at + for field in ("charset_version", "punct_version", "locale_specific_punct"): + assert field in message # what to set + + @pytest.mark.unit + @pytest.mark.parametrize( + "state_dict, text_embedding", + [({}, torch.nn.Embedding(10, 4)), ({"text_embedding.weight": torch.zeros(10, 4)}, None)], + ids=["no-weight-in-ckpt", "cas-encoder-variant-has-no-table"], + ) + def test_absent_embedding_is_not_an_error(self, state_dict, text_embedding): + """The CAS-encoder MagpieTTS variant has no text_embedding to compare; it must not trip here.""" + check_text_embedding_matches_tokenizer( + state_dict, text_embedding=text_embedding, tokenizer=MagicMock(), model_cfg=OmegaConf.create({}) + ) + + +class TestNemoRoundTrip: + """End-to-end ``save_to``/``restore_from`` coverage of the vocabulary-drift failure.""" + + @staticmethod + def _save_legacy_archive(tmp_path): + """Write a .nemo that looks like a pre-versioning release (v2512/v2602). + + Such archives were trained with the v1 charset/punctuation but their configs name neither, so + the versioned fields are stripped back out after the model is built. The ``nemo_version`` stamp + is what remains to date them -- both released archives carry exactly ``2.6.0rc0``. Setting it + before ``super().__init__`` is also what a real restore does: ``ModelPT`` only stamps a config + that has no version yet, so the original release's stamp survives every later save. + """ + model = _TokenizerVersionModel( + OmegaConf.create({"text_tokenizers": _tokenizers_cfg(), "nemo_version": LEGACY_VERSION}) + ) + model.text_embedding = torch.nn.Embedding( + len(setup_tokenizers(_tokenizers_cfg(), cfg_nemo_version=LEGACY_VERSION).tokens) + 2, 4 + ) + with open_dict(model.cfg): + del model.cfg.text_tokenizers.hindi_chartokenizer.charset_version + del model.cfg.text_tokenizers.hindi_chartokenizer.punct_version + path = str(tmp_path / "legacy.nemo") + model.save_to(path) + return path, model.text_embedding.num_embeddings + + @pytest.mark.unit + def test_archive_without_version_fields_restores_with_v1_vocab(self, tmp_path): + """Regression: a released checkpoint that predates the fields must not pick up the v2 charsets. + + Picking them up rebuilds a vocabulary of a different size than the checkpoint was trained with + (v2 collapses case, so it is the *smaller* of the two), and ``restore_from`` then dies on a + text_embedding size mismatch -- exactly how v2602 broke. + """ + path, expected_num_embeddings = self._save_legacy_archive(tmp_path) + + restored = _TokenizerVersionModel.restore_from(path, map_location="cpu") + + assert restored.text_embedding.num_embeddings == len(restored.tokenizer.tokens) + 2 + assert restored.text_embedding.num_embeddings == expected_num_embeddings + + @pytest.mark.unit + def test_newly_trained_model_round_trips_on_current_defaults(self, tmp_path): + """A model trained today keeps its v2 vocabulary through a save/restore cycle. + + The archive is stamped with the current release, so nothing here depends on the dating rule: + this pins that the versions persisted at save time are what decide the vocabulary, which is + what makes every future release self-describing regardless of how the class defaults move. + """ + model = _TokenizerVersionModel(OmegaConf.create({"text_tokenizers": _tokenizers_cfg()})) + assert model.cfg.text_tokenizers.hindi_chartokenizer.charset_version == 2 + path = str(tmp_path / "current.nemo") + model.save_to(path) + + restored = _TokenizerVersionModel.restore_from(path, map_location="cpu") + + assert restored.cfg.text_tokenizers.hindi_chartokenizer.charset_version == 2 + assert restored.cfg.text_tokenizers.hindi_chartokenizer.punct_version == 2 + assert restored.text_embedding.num_embeddings == model.text_embedding.num_embeddings diff --git a/tests/e2e_nightly/L2_Model_Support_nvidia__magpie_tts_multilingual_357m.sh b/tests/e2e_nightly/L2_Model_Support_nvidia__magpie_tts_multilingual_357m.sh index da7fbaff1597..a7d316a0e4a2 100644 --- a/tests/e2e_nightly/L2_Model_Support_nvidia__magpie_tts_multilingual_357m.sh +++ b/tests/e2e_nightly/L2_Model_Support_nvidia__magpie_tts_multilingual_357m.sh @@ -14,4 +14,5 @@ coverage run -a --data-file=/workspace/.coverage --source=/workspace/nemo \ -m pytest \ "tests/e2e_nightly/test_model_support_nvidia__magpie_tts_multilingual_357m.py" \ + --with_downloads \ -v diff --git a/tests/e2e_nightly/test_model_support_nvidia__magpie_tts_multilingual_357m.py b/tests/e2e_nightly/test_model_support_nvidia__magpie_tts_multilingual_357m.py index 2ad062a5f102..0216732a34eb 100644 --- a/tests/e2e_nightly/test_model_support_nvidia__magpie_tts_multilingual_357m.py +++ b/tests/e2e_nightly/test_model_support_nvidia__magpie_tts_multilingual_357m.py @@ -14,13 +14,17 @@ """Functional tests for nvidia/magpie_tts_multilingual_357m.""" +import gc import os import pytest import torch +from huggingface_hub import hf_hub_download MODEL_NAME = "nvidia/magpie_tts_multilingual_357m" NEMO_FILE = "nvidia__magpie_tts_multilingual_357m.nemo" +HF_NEMO_FILE = "magpie_tts_multilingual_357m.nemo" +MODEL_RELEASES = ("v2512", "v2602", "v2607") MODEL_DIR = os.environ.get( "NEMO_MODEL_SUPPORT_DIR", @@ -41,6 +45,45 @@ def _load_model(): return _model +@pytest.mark.with_downloads +@pytest.mark.parametrize("revision", MODEL_RELEASES) +def test_model_release_restore(revision): + """Ensure every published MagpieTTS release remains loadable by current NeMo code. + + The three releases cover each way the versioned tokenizer fields can present themselves: + + - v2512 (stamped nemo_version 2.6.0rc0) has no version-sensitive tokenizer at all. + - v2602 (2.6.0rc0) carries a ``HindiCharsTokenizer`` naming neither ``charset_version`` nor + ``punct_version``, so both must be dated back to v1 from the stamp. This is the restore that + broke and prompted the fix. + - v2607 (2.8.0rc0) pins ``locale_specific_punct=false`` (pt-BR) and ``charset_version=1`` (Arabic) + explicitly, so the stamp must not override them. + + A vocabulary rebuilt from anything but the values a release was trained with shifts every token ID, + which surfaces during ``restore_from`` as a text embedding size mismatch. + """ + from nemo.collections.tts.models import MagpieTTSModel + + filepath = hf_hub_download(repo_id=MODEL_NAME, filename=HF_NEMO_FILE, revision=revision) + model = MagpieTTSModel.restore_from(filepath, map_location="cpu") + + assert model is not None + # Independent restatement of the sizing in MagpieTTSModel.__init__, so this still catches drift if + # the check inside load_state_dict is ever removed. All three releases set + # legacy_text_conditioning=false; legacy models exclude the context-text tokens from this table. + expected_tokens = len(model.tokenizer.tokens) + if model.legacy_text_conditioning: + expected_tokens -= model.tokenizer.num_tokens_per_tokenizer[model.text_conditioning_tokenizer_name] + assert model.text_embedding.num_embeddings == expected_tokens + 2, ( + f"{revision}: text embedding has {model.text_embedding.num_embeddings} rows but the tokenizer " + f"built {expected_tokens} usable tokens (+2 for BOS/EOS) -- the restored tokenizer config does " + f"not match the one this release was trained with." + ) + + del model + gc.collect() + + def test_model_init(): model = _load_model() assert model is not None @@ -68,11 +111,9 @@ def test_model_training_step(): model.train() d = _DEVICE - # Derive safe upper bounds for token IDs from the model's embedding tables. - # text_embedding vocab = num_tokens_tokenizer + 2 (BOS/EOS). - text_vocab_size = model.text_embedding.num_embeddings # audio_codes must be strictly less than codebook_size (special tokens are # appended *after* codebook_size inside process_batch / add_special_tokens). + # Text IDs need no such bound here: the batch below uses the fixed safe ID 1. audio_token_max = model.codebook_size - 1 B = 1