From 673ccf51229ddd44ad22166fa4eb636159c3c25a Mon Sep 17 00:00:00 2001 From: Jason Date: Tue, 28 Jul 2026 11:21:07 -0400 Subject: [PATCH 1/2] Fix loading of MagpieTTS and add test for all 3 versions (#15941) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix legacy MagpieTTS text embedding restore Migrate checkpoints that contain conditioning-token rows while preserving BOS/EOS weights. Signed-off-by: Jason * Test MagpieTTS release checkpoint restores Signed-off-by: Jason * Fix v2602 restore without version metadata Signed-off-by: Jason * revert some code changes Signed-off-by: Jason * minor Signed-off-by: Jason * remove unused import Signed-off-by: Jason * Fix MagpieTTS tokenizer defaults for restore vs fresh training Configs that omit charset_version/punct_version/locale_specific_punct are ambiguous: either a released archive predating those fields (v1 vocabulary) or a config for a new run (current defaults). Forcing either one broke the other — v2602 failed to restore, or new training silently fell back to v1. Tell them apart via the nemo_version ModelPT stamps into saved configs, plus AppState's restore flag, read at the top of __init__ before the nested codec restore clears it. setup_tokenizers' use_legacy_defaults goes back to False. Also dedupes the backfill into persist_versioned_tokenizer_defaults(), which requires moving the charset_version backfill ahead of safe_instantiate() — it previously ran after, so the tokenizer never saw the value. That was harmless while both defaults were 2, but is fatal once the value can be 1. Verified against all three releases: on main, v2602 fails with a text_embedding size mismatch (2362 trained vs 2317 rebuilt, exactly the Hindi v1->v2 charset shrink), while v2512 and v2607 are unaffected. All three restore after the fix. Signed-off-by: Xuesong Yang <1646669+XuesongYang@users.noreply.github.com> * Decide legacy tokenizer defaults from config origin, not restore state Being restored does not make a checkpoint legacy, so drop the AppState().is_model_being_restored term (review feedback). It was redundant — every .nemo carries nemo_version — and it was what forced the check to run at the top of __init__, ahead of the nested codec restore that clears the flag. What remains asks only about the config: every NeMo that knows these fields writes them, so a serialized config lacking them must predate them. Renamed is_restored_model_config -> predates_versioned_tokenizer_fields to say that, and moved the call next to setup_tokenizers now that ordering no longer matters. All three releases still restore unchanged. Signed-off-by: Xuesong Yang <1646669+XuesongYang@users.noreply.github.com> * Cover the real MagpieTTS call sites and warn on legacy fallback Review follow-up, three fixes. Mutation testing showed the fix itself was untested: inverting or deleting use_legacy_defaults at magpietts.py:442 and easy_magpietts_inference.py:339 reintroduces the v2602 restore failure while the whole TTS unit suite stays green, because the suite drove a test-local ModelPT that re-implemented the wiring rather than exercising it. TestProductionCallSites runs the real constructors, stopped at the tokenizer-setup call so no codec or downloads are needed. Both mutations now fail. The docstring claimed the legacy fallback announces itself with a DeprecationWarning. It does not: IPATokenizer is silent when locale_specific_punct=False drops pt-BR's 11 extra marks, and the Hindi/Arabic warning is swallowed by Python's default filters outside pytest. Make the claim true rather than delete it — persist_versioned_tokenizer_defaults now logs the fields it backfilled, legacy direction only, so ordinary training stays quiet. Also drop the stale 'Must be called inside open_dict(model_cfg)'' from _migrate_versioned_tokenizer_fields, which unlocks each tokenizer node itself now. All three releases still restore unchanged. Signed-off-by: Xuesong Yang <1646669+XuesongYang@users.noreply.github.com> * Make the tokenizer version constants the source of truth DEFAULT_CHARSET_VERSION was never read inside tts_tokenizers.py — both tokenizers spelled their default as a literal 2, so the constant only looked authoritative. Add DEFAULT_PUNCT_VERSION, have the signatures read both, and assert the non-legacy backfill equals the class default read off inspect.signature, so a future version bump cannot leave new models pinned a version behind. The validation branches keep their literals: they define the supported set, not the default. Signed-off-by: Xuesong Yang <1646669+XuesongYang@users.noreply.github.com> * Date tokenizer defaults from nemo_version instead of config shape A config missing charset_version / punct_version / locale_specific_punct is ambiguous: pre-2.8.0 archives were trained on the v1 values, newer ones on the current defaults. Decide it from the config's nemo_version stamp rather than from whether the config was serialized at all, which mislabelled every modern checkpoint as legacy. - Add VERSIONED_TOKENIZER_FIELDS: one row per field, recording the release its default changed in, next to the DEFAULT_* constants it reads. - setup_tokenizers takes cfg_nemo_version; the resolved values are pinned into the config, so dataloader and worker_init_fn call sites stay version-agnostic and update_config_for_inference needs no migration. - Split CHARSET/PUNCT_VERSIONED_TOKENIZER_TARGETS out of CASELESS_SCRIPT_TOKENIZER_TARGETS: which tokenizers accept an argument is an API fact, not a linguistic one. - Raise an actionable error when a checkpoint's text embedding disagrees with the rebuilt tokenizer, naming the fields to pin. - Describe charset_version=1 / punct_version=1 as frozen for checkpoint compatibility rather than pending removal; v2607 pins charset_version=1. Verified against the released v2512 / v2602 / v2607 configs. Signed-off-by: Xuesong Yang <1646669+XuesongYang@users.noreply.github.com> * refine docstring Signed-off-by: Xuesong Yang <1646669+XuesongYang@users.noreply.github.com> * refine tests and codes Signed-off-by: Xuesong Yang <1646669+XuesongYang@users.noreply.github.com> * removed imported but unused lib. Signed-off-by: Xuesong Yang <1646669+XuesongYang@users.noreply.github.com> --------- Signed-off-by: Jason Signed-off-by: Xuesong Yang <1646669+XuesongYang@users.noreply.github.com> Co-authored-by: Xuesong Yang <1646669+XuesongYang@users.noreply.github.com> --- .../text_to_speech/tts_tokenizers.py | 155 +++++- .../tts/data/text_to_speech_dataset_lhotse.py | 87 ++-- .../tts/models/easy_magpietts_inference.py | 14 +- nemo/collections/tts/models/magpietts.py | 15 +- .../tts/modules/magpietts_inference/utils.py | 55 +- .../test_magpietts_tokenizer_versioning.py | 472 ++++++++++++++++++ ...rt_nvidia__magpie_tts_multilingual_357m.sh | 1 + ...rt_nvidia__magpie_tts_multilingual_357m.py | 47 +- 8 files changed, 731 insertions(+), 115 deletions(-) create mode 100644 tests/collections/tts/models/test_magpietts_tokenizer_versioning.py 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 c0ef45d9a7a9..750f840b26e7 100644 --- a/nemo/collections/tts/models/easy_magpietts_inference.py +++ b/nemo/collections/tts/models/easy_magpietts_inference.py @@ -25,7 +25,11 @@ from torch import nn from transformers import AutoConfig, AutoModelForCausalLM -from nemo.collections.tts.data.text_to_speech_dataset_lhotse import setup_tokenizers +from nemo.collections.audio.parts.utils.transforms import resample +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 @@ -331,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) @@ -603,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 e5f2b642033c..9aa690add12c 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 @@ -442,6 +446,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) @@ -1055,6 +1061,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 8b41ec1cd158..1f5c4c4d1689 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 From ded346a7bd3479d51b6ae93e8c64c72c54df08d1 Mon Sep 17 00:00:00 2001 From: Jason Date: Tue, 28 Jul 2026 13:49:16 -0700 Subject: [PATCH 2/2] fix import introduct during conflict resolution Signed-off-by: Jason --- nemo/collections/tts/models/easy_magpietts_inference.py | 1 - 1 file changed, 1 deletion(-) diff --git a/nemo/collections/tts/models/easy_magpietts_inference.py b/nemo/collections/tts/models/easy_magpietts_inference.py index 750f840b26e7..414c14f11e97 100644 --- a/nemo/collections/tts/models/easy_magpietts_inference.py +++ b/nemo/collections/tts/models/easy_magpietts_inference.py @@ -25,7 +25,6 @@ from torch import nn 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 ( check_text_embedding_matches_tokenizer, setup_tokenizers,