Add phoneme prediction on exportable .json and .csv files

Signed-off-by: Edresson Casanova <edresson1@gmail.com>

Author: Edresson

examples/tts/magpietts_inference.py

@@ -190,6 +190,8 @@ def _enrich_filewise_metrics_with_manifest(filewise_metrics: list, manifest_path
             "rank_local_idx",
             "audio_filepath",
             "context_audio_filepath",
+            "predicted_phoneme_text",
+            "predicted_phoneme_tokens",
         ]:
             if key in record and key not in new_row:
                 new_row[key] = record[key]
@@ -249,6 +251,8 @@ def turn_sort_key(r):
         eou_type_turns = [r.get("eou_type") for r in turns]
         eou_trailing_duration_turns = [r.get("eou_trailing_duration") for r in turns]
         eou_trail_rms_ratio_turns = [r.get("eou_trail_rms_ratio") for r in turns]
+        predicted_phoneme_text_turns = [r.get("predicted_phoneme_text", "") for r in turns]
+        predicted_phoneme_tokens_turns = [r.get("predicted_phoneme_tokens", []) for r in turns]
 
         grouped_rows.append(
             {
@@ -280,6 +284,8 @@ def turn_sort_key(r):
                 "eou_type_turns": eou_type_turns,
                 "eou_trailing_duration_turns": eou_trailing_duration_turns,
                 "eou_trail_rms_ratio_turns": eou_trail_rms_ratio_turns,
+                "predicted_phoneme_text_turns": predicted_phoneme_text_turns,
+                "predicted_phoneme_tokens_turns": predicted_phoneme_tokens_turns,
                 "reference_text": [r.get("gt_text", "") for r in turns],
                 "asr_hyp": [r.get("pred_text", "") for r in turns],
                 "pred_audio_paths": [r.get("pred_audio_filepath", "") for r in turns],
@@ -327,6 +333,8 @@ def _write_grouped_multiturn_filewise_metrics_csv(csv_path: str, grouped_rows: l
         "eou_type_turns",
         "eou_trailing_duration_turns",
         "eou_trail_rms_ratio_turns",
+        "predicted_phoneme_text_turns",
+        "predicted_phoneme_tokens_turns",
         "target_audio_path",
         "context_audio_path",
         "pred_audio_paths",

nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py

@@ -89,6 +89,8 @@ def strip_text_annotations_from_text(text: str) -> str:
     'pred_context_ssim',
     'pred_gt_esim',
     'pred_gt_ems',
+    'predicted_phoneme_text',
+    'predicted_phoneme_tokens',
     'pred_text',
     'gt_text',
     'gt_audio_filepath',
@@ -690,6 +692,9 @@ def evaluate_dir(
             'total_gen_audio_seconds': file_duration,
             'predicted_codes_path': codes_file_lists[ridx] if has_codes else None,
         }
+        for manifest_debug_key in ['predicted_phoneme_text', 'predicted_phoneme_tokens']:
+            if manifest_debug_key in record:
+                metric_row[manifest_debug_key] = record[manifest_debug_key]
         if with_emotion_metrics:
             metric_row['pred_gt_esim'] = pred_gt_esim
             metric_row['pred_gt_ems'] = pred_gt_ems

nemo/collections/tts/modules/magpietts_inference/inference.py

@@ -1167,6 +1167,40 @@ def _ensure_codec_silence_codes(self) -> torch.Tensor:
 
         return self.model._codec_sil_codes_buffer.to(self.model.device).long()
 
+    def _decode_phoneme_prediction_slice(self, state, start_step: int, end_step: int) -> Tuple[List[int], str]:
+        """Decode predicted phoneme-channel tokens for one generated turn."""
+        phoneme_tokenizer = getattr(self.model, "phoneme_tokenizer", None)
+        phoneme_predictions = getattr(state, "all_phoneme_predictions", None)
+        if phoneme_tokenizer is None or not phoneme_predictions or end_step <= start_step:
+            return [], ""
+
+        phoneme_slice = phoneme_predictions[start_step:end_step]
+        if len(phoneme_slice) == 0:
+            return [], ""
+
+        phoneme_tensor = torch.stack(phoneme_slice, dim=-1)  # (B, S, T)
+        tokens = phoneme_tensor[0].detach().cpu().T.reshape(-1).tolist()
+
+        special_tokens = {
+            getattr(phoneme_tokenizer, "bos_token_id", None),
+            getattr(phoneme_tokenizer, "eos_token_id", None),
+            getattr(phoneme_tokenizer, "pad_token_id", None),
+            getattr(phoneme_tokenizer, "pad", None),
+        }
+        special_tokens = {int(token) for token in special_tokens if token is not None}
+        tokens = [int(token) for token in tokens if int(token) >= 0 and int(token) not in special_tokens]
+
+        if not tokens:
+            return [], ""
+
+        try:
+            phoneme_text = phoneme_tokenizer.decode(tokens)
+        except Exception as exc:
+            logging.warning(f"Could not decode predicted phoneme tokens for multiturn debug export: {exc}")
+            phoneme_text = ""
+
+        return tokens, phoneme_text
+
     def _run_multiturn_generation(self, batch: Dict[str, Any]):
         model = self.model
         device = model.device
@@ -1237,6 +1271,7 @@ def _run_multiturn_generation(self, batch: Dict[str, Any]):
             )
 
             turn_frame_ranges = []
+            turn_phoneme_outputs = []
             decode_start_frame = 0
             max_decoder_steps = params.max_decoder_steps
 
@@ -1357,6 +1392,8 @@ def _run_multiturn_generation(self, batch: Dict[str, Any]):
                         user_audio_channel_embedding=user_audio_channel_embedding,
                     )
 
+                turn_phoneme_start_step = len(getattr(state, "all_phoneme_predictions", []))
+
                 for i in range(delay_tokens):
                     user_step_emb = warmup_user_audio[:, i] if warmup_user_audio is not None else None
                     state.finished.zero_()
@@ -1405,7 +1442,20 @@ def _run_multiturn_generation(self, batch: Dict[str, Any]):
                 state.audio_prediction_end_idx.fill_(-1)
                 state.finished.zero_()
                 turn_end_frame = sum(p.size(-1) for p in state.all_predictions)
+                turn_phoneme_end_step = len(getattr(state, "all_phoneme_predictions", []))
+                predicted_phoneme_tokens, predicted_phoneme_text = self._decode_phoneme_prediction_slice(
+                    state,
+                    turn_phoneme_start_step,
+                    turn_phoneme_end_step,
+                )
                 turn_frame_ranges.append((turn_id, turn_start_frame, turn_end_frame))
+                turn_phoneme_outputs.append(
+                    {
+                        "turn_id": int(turn_id),
+                        "predicted_phoneme_text": predicted_phoneme_text,
+                        "predicted_phoneme_tokens": predicted_phoneme_tokens,
+                    }
+                )
 
             codec_sil_codes = self._ensure_codec_silence_codes()
             bos_id = getattr(model, "audio_bos_id", -1)
@@ -1428,7 +1478,7 @@ def _run_multiturn_generation(self, batch: Dict[str, Any]):
 
             finalize_output = model.streaming_finalize(state, use_inference_mode=True)
 
-        return finalize_output, turn_frame_ranges, decode_start_frame, generated_codes
+        return finalize_output, turn_frame_ranges, decode_start_frame, generated_codes, turn_phoneme_outputs
 
     @staticmethod
     def _save_code_slice(
@@ -1637,9 +1687,15 @@ def _run_multiturn_user_audio_inference(
                 continue
 
             start_time = time.time()
-            output, turn_frame_ranges, decode_start_frame, generated_codes = self._run_multiturn_generation(batch)
+            output, turn_frame_ranges, decode_start_frame, generated_codes, turn_phoneme_outputs = (
+                self._run_multiturn_generation(batch)
+            )
             elapsed = time.time() - start_time
 
+            turn_phoneme_outputs_by_turn_id = {
+                int(item.get("turn_id", -1)): item for item in turn_phoneme_outputs
+            }
+
             predicted_audio = output.audio.float().detach().cpu()
             predicted_audio_lens = output.audio_len.int().detach().cpu()
             full_len = int(predicted_audio_lens[0].item())
@@ -1712,6 +1768,8 @@ def _run_multiturn_user_audio_inference(
                     description=f"target audio fallback/context for sample_idx={sample_idx}, turn_id={turn_id}",
                 )
 
+                phoneme_debug = turn_phoneme_outputs_by_turn_id.get(int(turn_id), {})
+
                 turn_manifest_records.append(
                     {
                         "audio_filepath": f"target_audio_{item_idx}.wav",
@@ -1720,6 +1778,8 @@ def _run_multiturn_user_audio_inference(
                         "speaker": str(sample_idx),
                         "source_sample_idx": sample_idx,
                         "turn_id": int(turn_id),
+                        "predicted_phoneme_text": phoneme_debug.get("predicted_phoneme_text", ""),
+                        "predicted_phoneme_tokens": phoneme_debug.get("predicted_phoneme_tokens", []),
                     }
                 )
                 logging.info(