Update ultravox_pipeline.py

ultravox_pipeline.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Any, Dict, List, Optional
 
+import numpy as np
 import transformers
 
 # We must use relative import in this directory to allow uploading to HF Hub
@@ -41,27 +42,30 @@ class UltravoxPipeline(transformers.Pipeline):
         super().__init__(model=model, tokenizer=tokenizer, **kwargs)
 
     def _sanitize_parameters(self, **kwargs):
-        generation_kwargs = {}
-        if "temperature" in kwargs:
-            generation_kwargs["temperature"] = kwargs["temperature"]
-        if "max_new_tokens" in kwargs:
-            generation_kwargs["max_new_tokens"] = kwargs["max_new_tokens"]
-        if "repetition_penalty" in kwargs:
-            generation_kwargs["repetition_penalty"] = kwargs["repetition_penalty"]
+        generation_keys = ["temperature", "max_new_tokens", "repetition_penalty"]
+        generation_kwargs = {k: kwargs[k] for k in kwargs if k in generation_keys}
         return {}, generation_kwargs, {}
 
     def preprocess(self, inputs: Dict[str, Any]):
-        if "turns" in inputs:
-            turns = inputs["turns"]
-        else:
-            turns = []
-
-        if not turns or turns[-1]["role"] != "user":
+        turns: list = inputs.get("turns", [])
+
+        audio = inputs.get("audio", None)
+        # Convert to float32 if needed.
+        if isinstance(audio, np.ndarray):
+            if audio.dtype == np.float64:
+                audio = audio.astype(np.float32)
+            elif audio.dtype == np.int16:
+                audio = audio.astype(np.float32) / np.float32(32768.0)
+            elif audio.dtype == np.int32:
+                audio = audio.astype(np.float32) / np.float32(2147483648.0)
+
+        if audio is not None and (len(turns) == 0 or turns[-1]["role"] != "user"):
             prompt = inputs.get("prompt", "<|audio|>")
             if "<|audio|>" not in prompt:
                 logging.warning(
                     "Prompt does not contain '<|audio|>', appending '<|audio|>' to the end of the prompt."
                 )
+
                 prompt += " <|audio|>"
             turns.append({"role": "user", "content": prompt})
 
@@ -69,17 +73,14 @@ class UltravoxPipeline(transformers.Pipeline):
             turns, add_generation_prompt=True, tokenize=False
         )
 
-        # TODO: allow text-only mode?
-        assert "audio" in inputs, "Audio input is required"
-
-        if "sampling_rate" not in inputs:
+        if "sampling_rate" not in inputs and audio is not None:
             logging.warning(
                 "No sampling rate provided, using default of 16kHz. We highly recommend providing the correct sampling rate."
             )
 
         output = self.processor(
             text=text,
-            audio=inputs["audio"],
+            audio=audio,
             sampling_rate=inputs.get("sampling_rate", 16000),
         )
         if "audio_values" in output:
@@ -123,4 +124,4 @@ transformers.pipelines.PIPELINE_REGISTRY.register_pipeline(
     pipeline_class=UltravoxPipeline,
     pt_model=transformers.AutoModel,
     type="multimodal",
-)
+)