Upload UltravoxPipeline

config.json

@@ -25,6 +25,24 @@
     ]
   },
   "audio_token_index": 32000,
+  "custom_pipelines": {
+    "ultravox-pipeline": {
+      "default": {
+        "model": {
+          "pt": [
+            "fixie-ai/ultravox-v0.2",
+            "main"
+          ]
+        }
+      },
+      "impl": "ultravox_pipeline.UltravoxPipeline",
+      "pt": [
+        "UltravoxModel"
+      ],
+      "tf": [],
+      "type": "multimodal"
+    }
+  },
   "hidden_size": 4096,
   "ignore_index": -100,
   "initializer_range": 0.02,

ultravox_pipeline.py

@@ -0,0 +1,111 @@
+import logging
+from typing import Any, Dict, List, Optional
+
+import transformers
+
+# We must use relative import in this directory to allow uploading to HF Hub
+from . import ultravox_model
+from . import ultravox_processing
+
+
+class UltravoxPipeline(transformers.Pipeline):
+    def __init__(
+        self,
+        model: ultravox_model.UltravoxModel,
+        tokenizer: Optional[transformers.PreTrainedTokenizerBase] = None,
+        audio_processor: Optional[transformers.ProcessorMixin] = None,
+        **kwargs
+    ):
+        if tokenizer is None:
+            tokenizer = transformers.AutoTokenizer.from_pretrained(
+                model.config._name_or_path
+            )
+
+        if audio_processor is None:
+            audio_processor = transformers.Wav2Vec2Processor.from_pretrained(
+                model.config.audio_model_id
+            )
+
+        self.processor = ultravox_processing.UltravoxProcessor(
+            audio_processor, tokenizer=tokenizer, stack_factor=model.config.stack_factor
+        )
+
+        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"]
+        return {}, generation_kwargs, {}
+
+    def preprocess(self, inputs: Dict[str, Any]):
+        if "turns" in inputs:
+            turns = inputs["turns"]
+        else:
+            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 = [{"role": "user", "content": prompt}]
+
+        text = self.processor.tokenizer.apply_chat_template(turns, tokenize=False)
+
+        # TODO: allow text-only mode?
+        assert "audio" in inputs, "Audio input is required"
+
+        if "sampling_rate" not in inputs:
+            logging.warning(
+                "No sampling rate provided, using default of 16kHz. We highly recommend providing the correct sampling rate."
+            )
+
+        return self.processor(
+            text=text,
+            audio=inputs["audio"],
+            sampling_rate=inputs.get("sampling_rate", 16000),
+        )
+
+    def _forward(
+        self,
+        model_inputs: Dict[str, Any],
+        temperature: Optional[float] = None,
+        max_new_tokens: Optional[int] = None,
+        repetition_penalty: float = 1.1,
+    ) -> List[int]:
+        temperature = temperature or None
+        do_sample = temperature is not None
+
+        terminators = [self.tokenizer.eos_token_id]
+        if "<|eot_id|>" in self.tokenizer.added_tokens_encoder:
+            terminators.append(self.tokenizer.convert_tokens_to_ids("<|eot_id|>"))
+
+        input_len = model_inputs["input_ids"].shape[1]
+
+        outputs = self.model.generate(
+            **model_inputs,
+            do_sample=do_sample,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            repetition_penalty=repetition_penalty,
+            eos_token_id=terminators
+        )
+        return outputs[0][input_len:]
+
+    def postprocess(self, model_outputs) -> str:
+        output_text = self.tokenizer.decode(model_outputs, skip_special_tokens=True)
+        return output_text
+
+
+transformers.pipeline
+transformers.pipelines.PIPELINE_REGISTRY.register_pipeline(
+    "ultravox-pipeline",  # TODO: make it broader later on
+    pipeline_class=UltravoxPipeline,
+    pt_model=ultravox_model.UltravoxModel,
+    default={"pt": ("fixie-ai/ultravox-v0.2", "main")},
+    type="multimodal",
+)