test-uv-pipeline / ultravox_pipeline.py
farzadab's picture
Upload UltravoxPipeline
d0ed04b verified
raw
history blame
3.93 kB
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",
)