Update ultravox_processing.py

ultravox_processing.py

@@ -113,7 +113,7 @@ class UltravoxProcessor(transformers.ProcessorMixin):
             tokenizer.eos_token is not None
         ), "The tokenizer has no EOS token. Cannot recover."
         self.vocab = tokenizer.get_vocab()
-        self.audio_replacement = tokenizer.eos_token
+        self.audio_token_replacement = tokenizer.eos_token
         if tokenizer.pad_token_id is None:
             tokenizer.pad_token_id = tokenizer.eos_token_id
 
@@ -188,7 +188,7 @@ class UltravoxProcessor(transformers.ProcessorMixin):
                 )
                 is_continuation_list.append(is_continuation)
 
-        return {
+        data = {
             "audio_values": torch.stack(chunked_audio_values, dim=0),
             "audio_lens": torch.tensor(
                 chunked_audio_lens, dtype=torch.int64, device=audio_values.device
@@ -199,12 +199,12 @@ class UltravoxProcessor(transformers.ProcessorMixin):
             "audio_batch_size": torch.tensor(
                 [len(chunked_audio_values)], device=audio_values.device
             ),
-            "audio_num_chunks": (
-                torch.tensor(num_chunks, dtype=torch.int64, device=audio_values.device)
-                if include_audio_num_chunks
-                else None
-            ),
         }
+        if include_audio_num_chunks:
+            data["audio_num_chunks"] = torch.tensor(
+                num_chunks, dtype=torch.int64, device=audio_values.device
+            )
+        return data
 
     def __call__(
         self,
@@ -327,7 +327,7 @@ class UltravoxProcessor(transformers.ProcessorMixin):
             split_input_ids = tokenized_parts["input_ids"]
             input_ids: List[int] = []
 
-            audio_replacement_token_id = self.vocab[self.audio_replacement]
+            audio_token_replacement_token_id = self.vocab[self.audio_token_replacement]
 
             for i, token_len in enumerate(data.get("audio_token_len", [])):
                 if not audio_is_continuation[i]:
@@ -341,7 +341,7 @@ class UltravoxProcessor(transformers.ProcessorMixin):
 
                 audio_token_start_idx.append(len(input_ids))
 
-                input_ids.extend([audio_replacement_token_id] * token_len)
+                input_ids.extend([audio_token_replacement_token_id] * token_len)
 
             # Include any tokens after the last audio.
             placeholder_index += 1