Fix loading of MagpieTTS and add test for all 3 versions

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=' ',

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', '<absent>')}\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.<name>`: "
+        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]+ \|"

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()