Update modeling_speech_encoder.py

modeling_speech_encoder.py

@@ -11,7 +11,24 @@ import torchaudio
 from transformers import PreTrainedModel
 
 from .configuration_speech_encoder import SpeechEncoderConfig
-from .collater_utils import wrap_bos_eos
+
+
+def wrap_bos_eos(
+    units: torch.Tensor,
+    durations: torch.Tensor,
+    f0: torch.Tensor | None,
+    dense_features: torch.Tensor,
+    bos: torch.Tensor,
+    eos: torch.Tensor,
+):
+    # bos/eos are 1-element tensors on the right device/dtype
+    one = durations.new_ones(1)
+    units = torch.cat([bos.to(units.device), units, eos.to(units.device)], dim=0)
+    durations = torch.cat([one, durations, one], dim=0)
+    if f0 is not None:
+        # pad f0 with edge values
+        f0 = torch.cat([f0[:1], f0, f0[-1:]], dim=0)
+    return units, durations, f0, dense_features
 
 
 # ----------------------------