Update ultravox_processing.py

ultravox_processing.py

@@ -163,11 +163,11 @@ class UltravoxProcessor(transformers.ProcessorMixin):
         chunked_audio_values: List[torch.Tensor] = []
         chunked_audio_lens: List[int] = []
         is_continuation_list: List[bool] = []
-        batch_size: List[int] = []
+        num_chunks: List[int] = []
         context_size = self.audio_context_size or audio_values.shape[-1]
 
         for i in range(audio_values.shape[0]):  # iterate over the batch
-            batch_size.append(int(np.ceil(audio_lens[i] / context_size)))
+            num_chunks.append(int(np.ceil(audio_lens[i] / context_size)))
             for offset in range(0, audio_lens[i], context_size):
                 is_continuation = offset > 0
                 chunk = audio_values[i, :, offset : offset + context_size]
@@ -193,7 +193,10 @@ class UltravoxProcessor(transformers.ProcessorMixin):
                 is_continuation_list, dtype=torch.bool, device=audio_values.device
             ),
             "audio_batch_size": torch.tensor(
-                batch_size, device=audio_values.device
+                [len(chunked_audio_values)], device=audio_values.device
+            ),
+            "audio_num_chunks": torch.tensor(
+                num_chunks, dtype=torch.int64, device=audio_values.device
             ),
         }