add phoneme control augmentation to multiturn dataloader (#16117)

* add phoneme control augmentation to multiturn dataloader

Signed-off-by: paarthneekhara <paarth.n@gmail.com>

* Apply suggestions from code review

Co-authored-by: Jason <jasoli@nvidia.com>
Signed-off-by: Jason <jasoli@nvidia.com>

---------

Signed-off-by: paarthneekhara <paarth.n@gmail.com>
Signed-off-by: Jason <jasoli@nvidia.com>
Co-authored-by: Jason <jasoli@nvidia.com>

Author: paarthneekhara

examples/tts/conf/magpietts/easy_magpietts_lhotse_multiturn.yaml

@@ -93,6 +93,14 @@ model:
   phoneme_corruption_timestep_ratio: 0.15
   phoneme_corruption_unk_mode_prob: 0.5
   phoneme_corruption_type: "repeat_skip_unk" # "repeat_skip_unk" or "complete_channel"
+  enable_phoneme_text_input: false # Enable inline IPA spans in text as <bop>...<eop>; requires phoneme_tokenizer.
+  partial_phoneme_text_prob: 0.0 # Training-only probability of using precomputed IPA alignments.
+  partial_phoneme_portion_min: 0.25 # Minimum portion of aligned words to phonemize in a selected sample.
+  partial_phoneme_portion_max: 0.75 # Maximum portion of aligned words to phonemize in a selected sample.
+  phoneme_text_bop_marker: "<bop>"
+  phoneme_text_eop_marker: "<eop>"
+  ignore_phoneme_languages: [] # Languages for which missing IPA phoneme fields are allowed during training.
+  add_language_to_context_text: false # Prefix context text with language metadata for multilingual conditioning.
   phoneme_turn_dropout_batch_prob: 0.0  # prob of applying turn dropout to a sample
   phoneme_turn_dropout_turn_prob: 0.0  # prob of dropping each phoneme turn within a sample
   phoneme_turn_max_words_to_drop: 0  # turns with <= this many words keep phoneme tokens as pad_id

nemo/collections/tts/data/text_to_speech_dataset_lhotse_multiturn.py

@@ -29,9 +29,13 @@
 from nemo.collections.speechlm2.parts.precision import fp32_precision
 from nemo.collections.tts.data.text_to_speech_dataset_lhotse import setup_tokenizers
 from nemo.collections.tts.parts.utils.tts_dataset_utils import (
+    _sample_probability_range,
     beta_binomial_prior_distribution,
+    has_phoneme_text_spans,
     normalize_volume,
+    partially_phonemize_text,
     stack_tensors,
+    tokenize_text_with_phoneme_spans,
 )
 from nemo.core.classes.common import safe_instantiate
 from nemo.utils import logging
@@ -150,6 +154,13 @@ def __init__(
         text_context_remapping_prob: float = 0.0,
         phoneme_tokenizer_config: DictConfig = None,
         ignore_phoneme_languages: List[str] = None,
+        enable_phoneme_text_input: bool = False,
+        text_phoneme_token_offset: int = None,
+        partial_phoneme_text_prob: float = 0.0,
+        partial_phoneme_portion_min: float = 0.25,
+        partial_phoneme_portion_max: float = 0.75,
+        phoneme_text_bop_marker: str = "<bop>",
+        phoneme_text_eop_marker: str = "<eop>",
         add_language_to_context_text: bool = False,
         source_sample_rate: int = 16000,
         input_roles: List[str] = ["user", "User"],
@@ -183,6 +194,13 @@ def __init__(
         self.text_context_remapping_prob = text_context_remapping_prob
         self.phoneme_tokenizer_config = phoneme_tokenizer_config
         self.ignore_phoneme_languages = ignore_phoneme_languages or []
+        self.enable_phoneme_text_input = enable_phoneme_text_input
+        self.text_phoneme_token_offset = text_phoneme_token_offset
+        self.partial_phoneme_text_prob = partial_phoneme_text_prob
+        self.partial_phoneme_portion_min = partial_phoneme_portion_min
+        self.partial_phoneme_portion_max = partial_phoneme_portion_max
+        self.phoneme_text_bop_marker = phoneme_text_bop_marker
+        self.phoneme_text_eop_marker = phoneme_text_eop_marker
         self.add_language_to_context_text = add_language_to_context_text
 
         self.source_sample_rate = source_sample_rate
@@ -313,6 +331,16 @@ def _collate_text_channels(self, cuts: CutSet, batch_tokenizer_names: list[str])
             eos_id=self.eos_id,
             bos_id=self.bos_id,
             interruption_token_id=self.interruption_token_id,
+            phoneme_tokenizer=self.phoneme_tokenizer,
+            enable_phoneme_text_input=self.enable_phoneme_text_input,
+            text_phoneme_token_offset=self.text_phoneme_token_offset,
+            partial_phoneme_text_prob=self.partial_phoneme_text_prob,
+            partial_phoneme_portion_min=self.partial_phoneme_portion_min,
+            partial_phoneme_portion_max=self.partial_phoneme_portion_max,
+            phoneme_text_bop_marker=self.phoneme_text_bop_marker,
+            phoneme_text_eop_marker=self.phoneme_text_eop_marker,
+            ignore_phoneme_languages=self.ignore_phoneme_languages,
+            apply_partial_phoneme_text=self.dataset_type == 'train',
         )
         source_tokens, source_token_lens = collate_token_channel(
             cuts,
@@ -768,6 +796,16 @@ def collate_token_channel(
     eos_id: int = None,
     bos_id: int = None,
     interruption_token_id: int = None,
+    phoneme_tokenizer=None,
+    enable_phoneme_text_input: bool = False,
+    text_phoneme_token_offset: int = None,
+    partial_phoneme_text_prob: float = 0.0,
+    partial_phoneme_portion_min: float = 0.25,
+    partial_phoneme_portion_max: float = 0.75,
+    phoneme_text_bop_marker: str = "<bop>",
+    phoneme_text_eop_marker: str = "<eop>",
+    ignore_phoneme_languages: list[str] = None,
+    apply_partial_phoneme_text: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Build and collate token channels aligned to the audio frame grid."""
     tokens = []
@@ -786,6 +824,16 @@ def collate_token_channel(
                 interruption_token_id,
                 add_text_bos,
                 tok_name,
+                phoneme_tokenizer,
+                enable_phoneme_text_input,
+                text_phoneme_token_offset,
+                partial_phoneme_text_prob,
+                partial_phoneme_portion_min,
+                partial_phoneme_portion_max,
+                phoneme_text_bop_marker,
+                phoneme_text_eop_marker,
+                ignore_phoneme_languages,
+                apply_partial_phoneme_text,
             )
         )
     token_lens = torch.tensor([len(tt) for tt in tokens])
@@ -831,22 +879,64 @@ def build_token_channel(
     interruption_token_id: int = -4,
     add_text_bos: bool = True,
     tokenizer_name: str = "english_phoneme",
+    phoneme_tokenizer=None,
+    enable_phoneme_text_input: bool = False,
+    text_phoneme_token_offset: int = None,
+    partial_phoneme_text_prob: float = 0.0,
+    partial_phoneme_portion_min: float = 0.25,
+    partial_phoneme_portion_max: float = 0.75,
+    phoneme_text_bop_marker: str = "<bop>",
+    phoneme_text_eop_marker: str = "<eop>",
+    ignore_phoneme_languages: list[str] = None,
+    apply_partial_phoneme_text: bool = False,
 ) -> torch.Tensor:
 
     total = compute_num_frames(cut.duration, frame_length, cut.sampling_rate)
     tokens = torch.ones(total, dtype=torch.long) * pad_id
 
     for supervision in cut.supervisions:
         if supervision.speaker in roles:
-            text = supervision.text
-
-            if hasattr(tokenizer, "encode"):
-                try:
-                    raw_ids = tokenizer.encode(text=text, tokenizer_name=tokenizer_name)
-                except TypeError:
-                    raw_ids = tokenizer.encode(text)
-            else:
-                raw_ids = tokenizer.text_to_ids(text)
+            # TODO: Current multi-turn datasets do not contain the normalized_text field so check will always default to the else branch. Will need to evaluate whether it makes sense to keep this in future.
+            # This code path is used for both multi-turn and single-turn datasets.
+            text = supervision.normalized_text if supervision.has_custom("normalized_text") else supervision.text
+            text_for_tokens = text
+            language = cut.lang if cut.has_custom("lang") else supervision.language
+            if (
+                apply_partial_phoneme_text
+                and enable_phoneme_text_input
+                and partial_phoneme_text_prob > 0.0
+                and language not in (ignore_phoneme_languages or [])
+                and supervision.has_custom("ipa_alignment")
+                and not has_phoneme_text_spans(
+                    text,
+                    bop_marker=phoneme_text_bop_marker,
+                    eop_marker=phoneme_text_eop_marker,
+                )
+                and random.random() < partial_phoneme_text_prob
+            ):
+                sampled_portion = _sample_probability_range(
+                    "partial_phoneme_portion",
+                    partial_phoneme_portion_min,
+                    partial_phoneme_portion_max,
+                )
+                text_for_tokens = partially_phonemize_text(
+                    text=text,
+                    ipa_alignment=supervision.ipa_alignment,
+                    partial_phoneme_portion=sampled_portion,
+                    full_ipa_text=_get_supervision_ipa_text(supervision),
+                    bop_marker=phoneme_text_bop_marker,
+                    eop_marker=phoneme_text_eop_marker,
+                )
+            raw_ids = tokenize_text_with_phoneme_spans(
+                text_tokenizer=tokenizer,
+                text_str=text_for_tokens,
+                tokenizer_name=tokenizer_name,
+                enable_phoneme_text_input=enable_phoneme_text_input,
+                phoneme_tokenizer=phoneme_tokenizer,
+                text_phoneme_token_offset=text_phoneme_token_offset,
+                bop_marker=phoneme_text_bop_marker,
+                eop_marker=phoneme_text_eop_marker,
+            )
 
             if add_text_bos:
                 text_ids = torch.as_tensor([bos_id] + raw_ids + [eos_id])

nemo/collections/tts/models/easy_magpietts.py

@@ -1975,6 +1975,13 @@ def get_lhotse_dataloader(self, dataset_cfg, mode='train') -> torch.utils.data.D
                 tokenizer_config=self.cfg.text_tokenizers,
                 phoneme_tokenizer_config=self.cfg.get("phoneme_tokenizer", None),
                 ignore_phoneme_languages=self.cfg.get("ignore_phoneme_languages", []),
+                enable_phoneme_text_input=self.enable_phoneme_text_input,
+                text_phoneme_token_offset=self.text_phoneme_token_offset,
+                partial_phoneme_text_prob=self.partial_phoneme_text_prob if mode == 'train' else 0.0,
+                partial_phoneme_portion_min=self.partial_phoneme_portion_min,
+                partial_phoneme_portion_max=self.partial_phoneme_portion_max,
+                phoneme_text_bop_marker=self.phoneme_text_bop_marker,
+                phoneme_text_eop_marker=self.phoneme_text_eop_marker,
                 add_language_to_context_text=self.add_language_to_context_text,
                 source_sample_rate=self.sample_rate,
                 input_roles=["user", "User"],

tests/collections/tts/data/test_magpietts_dataset_lhotse.py

@@ -47,6 +47,23 @@
     pass
 
 
+class _FakeIPATokenizer:
+    pad = 0
+    bos_token_id = 1
+    eos_token_id = 2
+
+    def encode(self, text):
+        return [7, 8] if text else []
+
+
+class _FakeTextTokenizer:
+    tokens = list(range(100))
+    pad = 0
+
+    def encode(self, text, tokenizer_name):
+        return [10 + len(text)]
+
+
 def _seed_everything():
     random.seed(42)
     np.random.seed(42)
@@ -140,6 +157,7 @@ def _multiturn_cutset():
             text="hello",
             language="en",
             speaker="assistant",
+            custom={"ipa": "həloʊ", "ipa_alignment": [[0, 5, "hello", "həloʊ"]]},
         ),
         SupervisionSegment(
             id="turn-user-1",
@@ -214,6 +232,45 @@ def test_single_turn_dataset_uses_bpe_and_cached_codes(self):
         assert batch["context_text_tokens_lens"].item() > 0
         assert batch["has_text_context"].tolist() == [True]
 
+    def test_multiturn_pronunciation_control_only_changes_target_turns(self):
+        _seed_everything()
+        kwargs = _dataset_kwargs()
+        kwargs.update(
+            {
+                "codec_model_input_sample_rate": CODEC_MODEL_INPUT_SAMPLE_RATE,
+                "frame_stacking_factor": FRAME_STACKING_FACTOR,
+                "source_sample_rate": SAMPLE_RATE,
+                "input_roles": ["user"],
+                "output_roles": ["assistant"],
+                "add_text_bos": False,
+                "use_text_conditioning_tokenizer": False,
+                "enable_phoneme_text_input": True,
+                "partial_phoneme_text_prob": 1.0,
+                "partial_phoneme_portion_min": 1.0,
+                "partial_phoneme_portion_max": 1.0,
+            }
+        )
+        dataset = MagpieTTSLhotseMultiturnDataset(**kwargs)
+        dataset.text_tokenizer = _FakeTextTokenizer()
+        dataset.phoneme_tokenizer = _FakeIPATokenizer()
+        dataset.bos_id = len(dataset.text_tokenizer.tokens)
+        dataset.eos_id = dataset.bos_id + 1
+        dataset.cfg_unk_token_id = dataset.bos_id + 2
+        dataset.interruption_token_id = dataset.bos_id + 3
+        dataset.pad_id = dataset.text_tokenizer.pad
+        dataset.text_phoneme_token_offset = dataset.bos_id + 4
+
+        batch = dataset[_multiturn_cutset()]
+
+        target_tokens = batch["text"][0]
+        source_tokens = batch["source_tokens"][0]
+        assert target_tokens[15:17].tolist() == [
+            dataset.text_phoneme_token_offset + 7,
+            dataset.text_phoneme_token_offset + 8,
+        ]
+        assert torch.all(source_tokens < dataset.text_phoneme_token_offset)
+        assert target_tokens[25].item() == dataset.interruption_token_id
+
     def test_multiturn_dataset_uses_bpe_and_cached_codes(self):
         _seed_everything()
         kwargs = _dataset_kwargs()