use phoneme tokenizer for edit speech

Signed-off-by: Paarth Neekhara <[email protected]>
paarthneekhara · Dec 1, 2023 · 981c34b · 981c34b
1 parent 28ebdc8
commit 981c34b
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/nemo/collections/nlp/data/language_modeling/megatron/t5_speechlm_dataset.py b/nemo/collections/nlp/data/language_modeling/megatron/t5_speechlm_dataset.py
@@ -647,6 +647,11 @@ def _get_tokens(self, doc, field, field_data):
                 instruction_tokens = self._get_text_tokens("Phoneme TTS")
                 field_tokens = self._get_phoneme_tokens(_text.replace("Phoneme TTS ", ""))
                 field_tokens = instruction_tokens + field_tokens
+            elif _text.startswith("Edit Speech"):
+                # Always use phoneme tokenizer for edit speech
+                instruction_tokens = self._get_text_tokens("Edit Speech")
+                field_tokens = self._get_phoneme_tokens(_text.replace("Edit Speech ", ""))
+                field_tokens = instruction_tokens + field_tokens
             else:
                 field_tokens = self._get_text_tokens(field_data.strip(" "))  # list of ids
         elif doc[f"{field}_type"] == 'SPEECH':