SpeechSynthesisDataset returns speaker_ids

lhotse/dataset/speech_synthesis.py

@@ -23,6 +23,7 @@
             'audio_lens': (B, ) int tensor
             'features_lens': (B, ) int tensor
             'tokens_lens': (B, ) int tensor
+            'speakers': List[str] of len B (optional) # if return_spk_ids is True
         }
     """
 
@@ -34,13 +35,15 @@
         feature_transforms: Union[Sequence[Callable], Callable] = None,
         add_eos: bool = True,
         add_bos: bool = True,
+        return_spk_ids: bool = False,
     ) -> None:
         super().__init__()
 
         self.cuts = cuts
         self.token_collater = TokenCollater(cuts, add_eos=add_eos, add_bos=add_bos)
         self.cut_transforms = ifnone(cut_transforms, [])
         self.feature_input_strategy = feature_input_strategy
+        self.return_spk_ids = return_spk_ids
 
         if feature_transforms is None:
             feature_transforms = []
@@ -65,15 +68,18 @@
             features = transform(features)
 
         tokens, tokens_lens = self.token_collater(cuts)
-
-        return {
+        batch = {
             "audio": audio,
             "features": features,
             "tokens": tokens,
             "audio_lens": audio_lens,
             "features_lens": features_lens,
             "tokens_lens": tokens_lens,
         }
+        if self.return_spk_ids:
+            batch["speakers"] = [cut.supervisions[0].speaker for cut in cuts]
+
+        return batch
 
 
 def validate_for_tts(cuts: CutSet) -> None: