Upload UltravoxPipeline

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

chat_template.jinja

@@ -0,0 +1,103 @@
+[gMASK]<sop>
+{%- if tools -%}
+<|system|>
+# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{% for tool in tools %}
+{{ tool | tojson(ensure_ascii=False) }}
+{% endfor %}
+</tools>
+
+For each function call, output the function name and arguments within the following XML format:
+<tool_call>{function-name}
+<arg_key>{arg-key-1}</arg_key>
+<arg_value>{arg-value-1}</arg_value>
+<arg_key>{arg-key-2}</arg_key>
+<arg_value>{arg-value-2}</arg_value>
+...
+</tool_call>{%- endif -%}
+{%- macro visible_text(content) -%}
+    {%- if content is string -%}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping -%}
+        {%- for item in content -%}
+            {%- if item is mapping and item.type == 'text' -%}
+                {{- item.text }}
+            {%- elif item is string -%}
+                {{- item }}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- else -%}
+        {{- content }}
+    {%- endif -%}
+{%- endmacro -%}
+{%- set ns = namespace(last_user_index=-1) %}
+{%- for m in messages %}
+    {%- if m.role == 'user' %}
+        {% set ns.last_user_index = loop.index0 -%}
+    {%- endif %}
+{%- endfor %}
+{% for m in messages %}
+{%- if m.role == 'user' -%}<|user|>
+{{ visible_text(m.content) }}
+{{- '/nothink' if (enable_thinking is defined and not enable_thinking and not visible_text(m.content).endswith("/nothink")) else '' -}}
+{%- elif m.role == 'assistant' -%}
+<|assistant|>
+{%- set reasoning_content = '' %}
+{%- set content = visible_text(m.content) %}
+{%- if m.reasoning_content is string %}
+    {%- set reasoning_content = m.reasoning_content %}
+{%- else %}
+    {%- if '</think>' in content %}
+        {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+        {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+    {%- endif %}
+{%- endif %}
+{%- if loop.index0 > ns.last_user_index and reasoning_content -%}
+{{ '\n<think>' + reasoning_content.strip() +  '</think>'}}
+{%- else -%}
+{{ '\n<think></think>' }}
+{%- endif -%}
+{%- if content.strip() -%}
+{{ '\n' + content.strip() }}
+{%- endif -%}
+{% if m.tool_calls %}
+{% for tc in m.tool_calls %}
+{%- if tc.function %}
+    {%- set tc = tc.function %}
+{%- endif %}
+{{ '\n<tool_call>' + tc.name }}
+{% set _args = tc.arguments %}
+{% for k, v in _args.items() %}
+<arg_key>{{ k }}</arg_key>
+<arg_value>{{ v | tojson(ensure_ascii=False) if v is not string else v }}</arg_value>
+{% endfor %}
+</tool_call>{% endfor %}
+{% endif %}
+{%- elif m.role == 'tool' -%}
+{%- if m.content is string -%}
+{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+    {{- '<|observation|>' }}
+{%- endif %}
+{{- '\n<tool_response>\n' }}
+{{- m.content }}
+{{- '\n</tool_response>' }}
+{%- else -%}
+<|observation|>{% for tr in m.content %}
+
+<tool_response>
+{{ tr.output if tr.output is defined else tr }}
+</tool_response>{% endfor -%}
+{% endif -%}
+{%- elif m.role == 'system' -%}
+<|system|>
+{{ visible_text(m.content) }}
+{%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    <|assistant|>{{- '\n<think></think>' if (enable_thinking is defined and not enable_thinking) else '' -}}
+{%- endif -%}

config.json

@@ -0,0 +1,92 @@
+{
+  "architectures": [
+    "UltravoxModel"
+  ],
+  "attn_implementation": null,
+  "audio_config": {
+    "_name_or_path": "openai/whisper-large-v3-turbo",
+    "activation_dropout": 0.0,
+    "activation_function": "gelu",
+    "apply_spec_augment": false,
+    "architectures": [
+      "WhisperForConditionalGeneration"
+    ],
+    "attention_dropout": 0.0,
+    "bos_token_id": 50257,
+    "classifier_proj_size": 256,
+    "d_model": 1280,
+    "decoder_attention_heads": 20,
+    "decoder_ffn_dim": 5120,
+    "decoder_layerdrop": 0.0,
+    "decoder_layers": 4,
+    "decoder_start_token_id": 50258,
+    "dropout": 0.0,
+    "dtype": "float16",
+    "encoder_attention_heads": 20,
+    "encoder_ffn_dim": 5120,
+    "encoder_layerdrop": 0.0,
+    "encoder_layers": 32,
+    "eos_token_id": 50257,
+    "init_std": 0.02,
+    "mask_feature_length": 10,
+    "mask_feature_min_masks": 0,
+    "mask_feature_prob": 0.0,
+    "mask_time_length": 10,
+    "mask_time_min_masks": 2,
+    "mask_time_prob": 0.05,
+    "max_source_positions": 1500,
+    "max_target_positions": 448,
+    "median_filter_width": 7,
+    "model_type": "whisper",
+    "num_hidden_layers": 32,
+    "num_mel_bins": 128,
+    "pad_token_id": 50257,
+    "scale_embedding": false,
+    "use_cache": true,
+    "use_weighted_layer_sum": false,
+    "vocab_size": 51866
+  },
+  "audio_latency_block_size": null,
+  "audio_model_id": null,
+  "audio_token_index": 151365,
+  "auto_map": {
+    "AutoConfig": "ultravox_config.UltravoxConfig",
+    "AutoModel": "ultravox_model.UltravoxModel"
+  },
+  "custom_pipelines": {
+    "ultravox-pipeline": {
+      "impl": "ultravox_pipeline.UltravoxPipeline",
+      "pt": [
+        "AutoModel"
+      ],
+      "tf": [],
+      "type": "multimodal"
+    }
+  },
+  "dtype": "bfloat16",
+  "hidden_logits_layers": null,
+  "hidden_size": 4096,
+  "ignore_index": -100,
+  "initializer_range": 0.02,
+  "llm_only_training": false,
+  "model_type": "ultravox",
+  "moe_capture_layers": [
+    5,
+    25,
+    45,
+    65,
+    85
+  ],
+  "norm_init": 0.4,
+  "num_llm_layers": -1,
+  "num_projector_layers": 4,
+  "pad_token_id": 151329,
+  "projector_act": "swiglu",
+  "projector_ln_mid": true,
+  "stack_factor": 8,
+  "text_model_id": "zai-org/GLM-4.6",
+  "text_model_quantization": null,
+  "transformers_version": "4.56.2",
+  "use_fsdp": true,
+  "vocab_size": 151552
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": [
+    151329,
+    151336,
+    151338
+  ],
+  "pad_token_id": 151329,
+  "transformers_version": "4.56.2"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd9c5309e305c27c6620149fdec942a3c917afa2c6ef27409aeecda264550d79
+size 1474609120

special_tokens_map.json

@@ -0,0 +1,25 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "<|audio|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96e1ca1945519d47f9b5e3d8f0eeb3c2955ccdcf1a27cc062c44e8507fa4ab59
+size 19970886

tokenizer_config.json

@@ -0,0 +1,316 @@
+{
+  "added_tokens_decoder": {
+    "151329": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151330": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151331": {
+      "content": "[gMASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151332": {
+      "content": "[sMASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151333": {
+      "content": "<sop>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151334": {
+      "content": "<eop>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151335": {
+      "content": "<|system|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151336": {
+      "content": "<|user|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151337": {
+      "content": "<|assistant|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151338": {
+      "content": "<|observation|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151339": {
+      "content": "<|begin_of_image|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151340": {
+      "content": "<|end_of_image|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151341": {
+      "content": "<|begin_of_video|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151342": {
+      "content": "<|end_of_video|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151343": {
+      "content": "<|begin_of_audio|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151344": {
+      "content": "<|end_of_audio|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151345": {
+      "content": "<|begin_of_transcription|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151346": {
+      "content": "<|end_of_transcription|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151347": {
+      "content": "<|code_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151348": {
+      "content": "<|code_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151349": {
+      "content": "<|code_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151350": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151351": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151352": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151353": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151354": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151355": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151356": {
+      "content": "<arg_key>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151357": {
+      "content": "</arg_key>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151358": {
+      "content": "<arg_value>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151359": {
+      "content": "</arg_value>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151360": {
+      "content": "/nothink",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151361": {
+      "content": "<|begin_of_box|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151362": {
+      "content": "<|end_of_box|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151363": {
+      "content": "<|image|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151364": {
+      "content": "<|video|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151365": {
+      "content": "<|audio|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|audio|>"
+  ],
+  "auto_map": {
+    "AutoProcessor": "ultravox_processing.UltravoxProcessor"
+  },
+  "clean_up_tokenization_spaces": false,
+  "do_lower_case": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 128000,
+  "pad_token": "<|endoftext|>",
+  "padding_side": "right",
+  "processor_class": "UltravoxProcessor",
+  "remove_space": false,
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}

ultravox_config.py

@@ -0,0 +1,200 @@
+import dataclasses
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+import transformers
+from transformers import BitsAndBytesConfig
+from transformers import HqqConfig
+
+@dataclasses.dataclass
+class LoraConfigSimplified:
+    """
+    Low Rank Approximation (LoRA) configuration.
+
+    Used for language and audio models separately.
+    """
+
+    # The rank of the approximation
+    r: int = 0
+    lora_alpha: float = 8
+    target_modules: Optional[List[str]] = dataclasses.field(
+        default_factory=lambda: ["k_proj", "q_proj", "linear_k", "linear_q"]
+    )
+    # A list of module names regex patterns to unfreeze. Only used if r == 0.
+    unfreeze_layers: Optional[List[str]] = None
+
+
+class LossMaskType(str, Enum):
+    """Type of loss mask to use."""
+
+    LAST_ASSISTANT = "last_assistant"
+    """This applies the loss mask up until the last assistant token"""
+    ALL = "all"  # This does not work with KL loss
+    """No loss mask, all inputs are used for loss"""
+    AFTER_AUDIO = "after_audio"
+    """Applies the loss mask up until the audio token"""
+
+
+class UltravoxConfig(transformers.PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`UltravoxForConditionalGeneration`]. It is used to instantiate an
+    Ultravox model according to the specified arguments, defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        audio_config (`WhisperConfig`,  *optional*):
+            Custom audio config or dict
+        text_config (`Union[AutoConfig, dict]`, *optional*):
+            The config object of the text backbone. Can be any of `LlamaConfig` or `MistralConfig`.
+        ignore_index (`int`, *optional*, defaults to -100):
+            The ignore index for the loss function.
+        audio_token_index (`int`, *optional*, defaults to 32000):
+            The audio token index to encode the audio prompt.
+        stack_factor (`int`, *optional*, defaults to 8):
+            Audio downsampling factor for the multimodal projector.
+        norm_init (`float`, *optional*, defaults to 0.4):
+            The initialization value for the layer normalization.
+        projector_act (`str`, *optional*, defaults to `"swiglu"`):
+            The activation function used by the multimodal projector.
+        text_model_lora_config (`LoraConfigSimplified`, *optional*):
+            The LoRA configuration for finetuning the text model.
+        audio_model_lora_config (`LoraConfigSimplified`, *optional*):
+            The LoRA configuration for finetuning the audio model.
+        audio_latency_block_size (`int`, *optional*, defaults to `None`):
+            The latency block size for simulating audio streaming.
+        attn_implementation (`str`, *optional*, defaults to "eager"):
+            The attention implementation to use. Choose from one of "eager", "flash_attention_2", or "sdpa".
+        use_fsdp (`bool`, *optional*, defaults to False):
+            If set to True, Ultravox will use FSDP for distributed training.
+            Configuration for kernel acceleration backends (unsloth, liger, etc).
+
+
+    Example:
+
+    ```python
+    >>> from transformers import UltravoxModel, WhisperConfig, UltravoxConfig, LlamaConfig
+
+    >>> # Initializing an audio encoder config
+    >>> audio_config = WhisperConfig()
+
+    >>> # Initializing a Llama config
+    >>> text_config = LlamaConfig()
+
+    >>> # Initializing a default configuration
+    >>> configuration = UltravoxConfig(audio_config, text_config)
+
+    >>> # Initializing a completely untrained model from the configuration
+    >>> model = UltravoxModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+
+    >>> # Initialize a model from pretrained checkpoints and random projector weights
+    >>> config = UltravoxConfig(audio_model_id="openai/whisper-tiny", text_model_id="meta-llama/Llama-2-7b-chat-hf")
+    ```"""
+
+    model_type = "ultravox"
+    is_composition = False
+
+    def __init__(
+        self,
+        audio_config: dict[str, Any] | transformers.PretrainedConfig | None = None,
+        text_config: dict[str, Any] | transformers.PretrainedConfig | None = None,
+        audio_model_id: str | None = None,
+        text_model_id: str | None = None,
+        llm_only_training: bool = False,
+        ignore_index: int = -100,
+        audio_token_index: int | None = None,
+        hidden_size: int = 4096,
+        stack_factor: int = 8,
+        norm_init: float = 0.4,
+        projector_act: str = "swiglu",
+        projector_ln_mid: bool = False,  # defaults to False for compatibility with v0.4.1 and below
+        num_projector_layers: int = 4,
+        text_model_lora_config: LoraConfigSimplified | None = None,
+        audio_model_lora_config: LoraConfigSimplified | None = None,
+        num_llm_layers: int = -1,
+        audio_latency_block_size: int | None = None,
+        attn_implementation: str = "eager",
+        use_fsdp: bool = False,
+        **kwargs,
+    ):
+        self.ignore_index = ignore_index
+
+        self.audio_model_id = audio_model_id
+        self.text_model_id = text_model_id
+
+        self.audio_token_index = audio_token_index
+
+        self.hidden_size = hidden_size
+        self.stack_factor = stack_factor
+        self.norm_init = norm_init
+        self.projector_act = projector_act
+        self.projector_ln_mid = projector_ln_mid
+        self.num_projector_layers = num_projector_layers
+        if text_model_id is not None:
+            text_config = transformers.AutoConfig.from_pretrained(text_model_id)
+        else:
+            text_config = text_config or {}
+            if isinstance(text_config, dict):
+                text_config = transformers.CONFIG_MAPPING[
+                    text_config.get("model_type", "llama")
+                ](**text_config)
+
+        if audio_model_id is not None:
+            audio_config = transformers.AutoConfig.from_pretrained(audio_model_id)
+        else:
+            audio_config = audio_config or {}
+            if isinstance(audio_config, dict):
+                audio_config = transformers.CONFIG_MAPPING[
+                    audio_config.get("model_type", "whisper")
+                ](**audio_config)
+
+        self.text_config = text_config
+        self.audio_config = audio_config
+
+        self.llm_only_training = llm_only_training
+        self.text_model_lora_config = (
+            text_model_lora_config
+            if isinstance(text_model_lora_config, dict)
+            else dataclasses.asdict(text_model_lora_config or LoraConfigSimplified())
+        )
+        self.audio_model_lora_config = (
+            audio_model_lora_config
+            if isinstance(audio_model_lora_config, dict)
+            else dataclasses.asdict(audio_model_lora_config or LoraConfigSimplified())
+        )
+        self.audio_latency_block_size = audio_latency_block_size
+
+        if hasattr(text_config, "text_config"):
+            text_config.vocab_size = text_config.text_config.vocab_size
+            text_config.hidden_size = text_config.text_config.hidden_size
+
+        self.vocab_size = text_config.vocab_size
+
+        self.initializer_range = text_config.initializer_range
+
+        self.attn_implementation = attn_implementation
+
+        self.use_fsdp = use_fsdp
+        self.num_llm_layers = num_llm_layers
+
+        super().__init__(**kwargs)
+
+    def to_diff_dict(self) -> Dict[str, Any]:
+        diff_dict = super().to_diff_dict()
+
+        # remove text_config and audio_config if text_model_id and audio_model_id are present
+        if self.text_model_id is not None:
+            diff_dict.pop("text_config", None)
+        elif "text_config" in diff_dict:
+            diff_dict["text_config"].pop("_attn_implementation_autoset", None)
+
+        if self.audio_model_id is not None:
+            diff_dict.pop("audio_config", None)
+        elif "audio_config" in diff_dict:
+            diff_dict["audio_config"].pop("_attn_implementation_autoset", None)
+
+        return diff_dict

ultravox_model.py

@@ -0,0 +1,1038 @@
+import logging
+import re
+from typing import Any, Dict, Generator, Optional, Set, Tuple, TypeVar, Union
+
+import contextlib
+
+import copy
+import accelerate
+import peft
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+import transformers.activations
+import transformers.modeling_outputs
+import transformers.models
+from transformers.generation.utils import GenerationMixin
+from transformers.models.whisper import modeling_whisper as whisper
+
+# from ultravox.model import kernel_integrations
+# from ultravox.model import model_loading_utils
+# from ultravox.utils import device_helpers
+
+# We must use relative import in this directory to allow uploading to HF Hub
+# Even "from . import X" pattern doesn't work (undocumented and unclear why)
+from .ultravox_config import UltravoxConfig
+from .ultravox_processing import UltravoxProcessor
+
+FROM_PRETRAINED_KWARGS = {}
+SHARED_PRETRAINED_KWARGS = [
+    "tp_plan",
+    "device_map",
+    "torch_dtype",
+    "attn_implementation",
+    "use_flash_attention_2",
+]
+
+class UltravoxModel(transformers.LlamaPreTrainedModel, GenerationMixin):
+    """
+    The Ultravox model which consists of an audio encoder and a language model.
+
+    Audio input is processed by the audio encoder, then every `stack_factor` frames are stacked together and
+    projected to the language model's embedding space using a few linear layers.
+    The text is embedded by the language model as usual and then the audio and text embeddings are merged together.
+
+    A special token `<|audio|>` is used to indicate the start of the audio embeddings in the merged embeddings.
+
+    Parameters:
+        config: Model configuration class with all the parameters of the model.
+    """
+
+    config_class = UltravoxConfig
+    config: UltravoxConfig  # for type hinting
+    # Usually we load encoder and LLM weights from a pretrained model separately, so they are allowed to be missing
+    _keys_to_ignore_on_load_missing = ["audio_tower.*", "language_model.*"]
+    # Since we have kwargs in forward, we need to set this to False, otherwise grad_accum_steps will cause incorrect train loss to be reported
+    # see https://github.com/huggingface/transformers/issues/35856 and https://github.com/huggingface/trl/pull/2615/files
+    accepts_loss_kwargs = False
+
+    def __init__(self, config: UltravoxConfig):
+        super().__init__(config)
+        self._register_load_state_dict_pre_hook(self._pre_load_state_dict_hook)
+
+        self.keep_params: Set[str] = set()
+        self.vocab_size = config.vocab_size
+
+        if not config.llm_only_training:
+            self.audio_tower = self._create_audio_tower(config)
+            self.multi_modal_projector = self._create_multi_modal_projector(config)
+            self.audio_tower_context_length = self.audio_tower.max_context_length
+
+        self.language_model = self._create_language_model(config)
+
+        if self.language_model._tied_weights_keys is not None:
+            self._tied_weights_keys = [
+                f"language_model.{k}" for k in self.language_model._tied_weights_keys
+            ]
+
+        # Determine no_split_modules dynamically to use with FSDP auto_wrap policy.
+        # This would be something like ["LlamaDecoderLayer"] as we don't split audio encoder layers.
+        # FSDP throws an error if some of the layer types are not found in the model, and they need to be filted out.
+        # 1. Get the names the language model *wants* to keep intact
+        candidate_names = set(
+            getattr(self.language_model, "_no_split_modules", []) or []
+        )
+        # 2. Names that actually exist in the current model
+        present_names = {m.__class__.__name__ for m in self.modules()}
+        # 3. Keep only those that are both requested and present
+        self._no_split_modules = list(candidate_names & present_names)
+
+        self.post_init()
+
+    def _init_weights(self, module):
+        if module is self:
+            if self.config.text_model_id is not None:
+                self.language_model = self._create_language_model(self.config)
+            if self.config.audio_model_id is not None:
+                self.audio_tower = self._create_audio_tower(self.config)
+        elif module in self.language_model.modules():
+            pass
+        elif module in self.audio_tower.modules():
+            pass
+        else:
+            super()._init_weights(module)
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        global FROM_PRETRAINED_KWARGS
+        FROM_PRETRAINED_KWARGS = {
+            k: v for k, v in kwargs.items() if k in SHARED_PRETRAINED_KWARGS
+        }
+        model = super().from_pretrained(*args, **kwargs)
+        FROM_PRETRAINED_KWARGS = {}
+        return model
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def set_decoder(self, decoder):
+        self.language_model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.language_model.get_decoder()
+
+    def tie_weights(self):
+        return self.language_model.tie_weights()
+
+    def _setup_cache(
+        self, cache_cls, max_batch_size: int, max_cache_len: Optional[int] = None
+    ):
+        self.language_model._setup_cache(cache_cls, max_batch_size, max_cache_len)
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        return self.language_model._reorder_cache(past_key_values, beam_idx)
+
+    def resize_token_embeddings(
+        self,
+        new_num_tokens: Optional[int] = None,
+        pad_to_multiple_of: Optional[int] = None,
+    ) -> nn.Embedding:
+        model_embeds = self.language_model.resize_token_embeddings(
+            new_num_tokens, pad_to_multiple_of
+        )
+        # update vocab size
+        self.config.text_config.vocab_size = model_embeds.num_embeddings
+        self.config.vocab_size = model_embeds.num_embeddings
+        self.vocab_size = model_embeds.num_embeddings
+        return model_embeds
+
+    def _audio_iter(
+        self, audio_batch_size: torch.Tensor
+    ) -> Generator[Tuple[int, int], None, None]:
+        """
+        Iterate over the audio batch size and yield the batch index and audio index of each audio item.
+
+        Args:
+            audio_batch_size: A tensor of shape (B,) where B is the batch size.
+
+        Returns:
+            A generator that yields a tuple of (start index, length) for each audio item.
+        """
+        audio_index = 0
+        for i_b, batch_count in enumerate(audio_batch_size):
+            for _ in range(batch_count):
+                yield i_b, audio_index
+                audio_index += 1
+
+    def _select_embedings(
+        self,
+        inputs_embeds: torch.Tensor,
+        start_idx: torch.Tensor,
+        lengths: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Select a contiguous slice per batch starting at `start_idx[b]` with
+        length `lengths[b]`, returned in a compact, front-aligned tensor.
+        Any positions in the output that correspond to padding are zeroed out.
+
+        Supports both 3D tensors (B, T, D) and 2D tensors (B, T).
+        """
+        B = inputs_embeds.size(0)
+        T = inputs_embeds.size(1)
+        max_length = int(lengths.max().item())
+        if max_length == 0:
+            # Return an empty slice with correct rank
+            if inputs_embeds.dim() == 3:
+                return inputs_embeds.new_zeros((B, 0, inputs_embeds.size(2)))
+            else:
+                return inputs_embeds.new_zeros((B, 0), dtype=inputs_embeds.dtype)
+
+        # --- Create indices to gather ---
+        idx = torch.arange(
+            max_length, device=inputs_embeds.device, dtype=start_idx.dtype
+        )  # (Lmax,)
+        pos = start_idx.unsqueeze(1) + idx.unsqueeze(0)  # (B, Lmax)
+        # Clamp to prevent out-of-bounds gather, we will mask the invalid values later
+        pos = pos.clamp_(0, T - 1)
+
+        # --- Create mask for valid output positions ---
+        mask = idx.unsqueeze(0) < lengths.unsqueeze(1)  # (B, Lmax)
+
+        # --- Gather and mask ---
+        if inputs_embeds.dim() == 3:
+            D = inputs_embeds.size(2)
+            gathered = inputs_embeds.gather(
+                1, pos.unsqueeze(-1).expand(B, max_length, D)
+            )
+            # Zero out the padded values
+            gathered = gathered * mask.unsqueeze(-1)
+            return gathered
+        elif inputs_embeds.dim() == 2:
+            gathered = inputs_embeds.gather(1, pos)
+            # Zero out the padded values
+            gathered = gathered * mask
+            return gathered
+        else:
+            raise ValueError(
+                f"_select_embedings expects 2D or 3D tensors, got {inputs_embeds.dim()}D"
+            )
+
+    def _decoder_layers(self):
+        """Return decoder blocks across architectures (LLaMA/GLM/etc.)."""
+        lm = self.language_model
+        candidates = [
+            getattr(getattr(lm, "model", None), "layers", None),
+            getattr(getattr(lm, "transformer", None), "layers", None),
+            getattr(lm, "layers", None),
+        ]
+        for c in candidates:
+            if c is not None and hasattr(c, "__len__") and len(c) > 0:
+                return c
+        raise AttributeError(
+            "Could not locate decoder layers on language_model. Tried .model.layers, .transformer.layers, and .layers"
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        audio_values: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        audio_token_start_idx: Optional[torch.Tensor] = None,
+        audio_lens: Optional[torch.Tensor] = None,
+        audio_token_len: Optional[torch.Tensor] = None,
+        audio_batch_size: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None,
+        # the alt_* fields are needed for KL divergence loss
+        alt_input_ids: Optional[torch.Tensor] = None,
+        alt_attention_mask: Optional[torch.Tensor] = None,
+        alt_labels: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> transformers.modeling_outputs.CausalLMOutputWithPast:
+        """
+        Forward pass for the Ultravox model.
+
+        `input_ids` are the tokenized text input. They are embedded by the language model as usual.
+        `audio_values` are processed by the audio encoder and then every `stack_factor` frames are stacked together and
+        projected to the language model's embedding space using a few linear layers.
+        The audio and text embeddings are merged together. A special token `<|audio|>` is used to indicate the start
+        of the audio embeddings in the merged embeddings.
+
+        Args:
+            input_ids: The tokenized text input.
+            audio_values: The processed audio values.
+            inputs_embeds: The embeddings for the input tokens.
+            labels: The tokenized text labels.
+            attention_mask: The attention mask for the input.
+            position_ids: The position ids for the input.
+            past_key_values: The past key value cache for the language model attention layers.
+            **kwargs: Additional keyword arguments. Passed directly to the language model.
+        """
+        if inputs_embeds is None:
+            # B x T  ->  B x T x D
+            inputs_embeds = self.get_input_embeddings().forward(input_ids)
+
+        if audio_values is not None and len(audio_values) > 0:
+            inputs_embeds, _, _ = (
+                self._prepare_audio_embeds(
+                    inputs_embeds=inputs_embeds,
+                    audio_values=audio_values,
+                    audio_token_start_idx=audio_token_start_idx,
+                    audio_lens=audio_lens,
+                    audio_token_len=audio_token_len,
+                    audio_batch_size=audio_batch_size,
+                )
+            )
+    
+        lm_output = self.language_model.forward(
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            output_hidden_states=False,  # avoid keeping all layers
+            **kwargs,
+        )
+
+        return lm_output
+
+ 
+
+    def _prepare_audio_embeds(
+        self,
+        inputs_embeds: torch.FloatTensor,
+        audio_values: torch.FloatTensor,
+        audio_token_start_idx: Optional[torch.Tensor] = None,
+        audio_lens: Optional[torch.Tensor] = None,
+        audio_token_len: Optional[torch.Tensor] = None,
+        audio_batch_size: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        assert (
+            audio_token_start_idx is not None
+            and audio_token_len is not None
+            and audio_lens is not None
+            and audio_batch_size is not None
+        ), "audio_token_start_idx/audio_token_len/audio_lens must be provided if audio_values are provided."
+        assert (
+            len(audio_token_start_idx)
+            == len(audio_token_len)
+            == len(audio_lens)
+            == len(audio_values)
+        ), "audio_token_start_idx/audio_token_len/audio_lens/audio_values must have the same batch size."
+        assert len(audio_batch_size) == len(
+            inputs_embeds
+        ), "audio_batch_size and inputs_embeds must have the same batch size."
+
+        # B x A/3200 x (D=max-audio-length-in-batch)
+        audio_tower_output = self.audio_tower.forward(
+            audio_values.to(self.audio_tower.dtype),
+            audio_len=audio_lens,
+        ).last_hidden_state
+        audio_tower_output = audio_tower_output.to(inputs_embeds.dtype)
+        audio_embeds = self.multi_modal_projector.forward(
+            audio_tower_output, audio_token_len=audio_token_len
+        )
+
+        # combine audio and text embeddings
+        merged_audio_start_idx = torch.zeros(
+            (inputs_embeds.shape[0],), device=inputs_embeds.device, dtype=torch.long
+        )
+        merged_audio_token_len = torch.zeros(
+            (inputs_embeds.shape[0],), device=inputs_embeds.device, dtype=torch.long
+        )
+        for i_b, i_a in self._audio_iter(audio_batch_size):
+            start_idx = audio_token_start_idx[i_a]
+            token_len = audio_token_len[i_a]
+            item_embedding = audio_embeds[i_a][:token_len]
+            inputs_embeds[i_b][start_idx : start_idx + token_len] = item_embedding
+            if merged_audio_token_len[i_b] == 0:
+                merged_audio_start_idx[i_b] = start_idx
+            merged_audio_token_len[i_b] += token_len
+
+        return inputs_embeds, merged_audio_start_idx, merged_audio_token_len
+
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        audio_values: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        audio_token_start_idx: Optional[torch.Tensor] = None,
+        audio_lens: Optional[torch.Tensor] = None,
+        audio_token_len: Optional[torch.Tensor] = None,
+        audio_batch_size: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings().forward(input_ids)
+
+        if audio_values is not None:
+            inputs_embeds, _, _ = (
+                self._prepare_audio_embeds(
+                    inputs_embeds=inputs_embeds,
+                    audio_values=audio_values,
+                    audio_token_start_idx=audio_token_start_idx,
+                    audio_lens=audio_lens,
+                    audio_token_len=audio_token_len,
+                    audio_batch_size=audio_batch_size,
+                )
+            )
+
+        return self.language_model.generate(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            **kwargs,
+        )
+
+    @classmethod
+    def _create_multi_modal_projector(
+        cls, config: UltravoxConfig
+    ) -> "UltravoxProjector":
+        if (
+            transformers.modeling_utils._init_weights
+            and config.text_model_id is not None
+        ):
+            projector = UltravoxProjector(config)
+            dtype = config.torch_dtype
+            if isinstance(dtype, str):
+                dtype = getattr(torch, dtype)
+            projector.to(dtype)
+            return projector
+        else:
+            with accelerate.init_empty_weights():
+                projector = UltravoxProjector(config)
+                dtype = config.torch_dtype
+                if isinstance(dtype, str):
+                    dtype = getattr(torch, dtype)
+                projector.to(dtype)
+                return projector
+
+    @classmethod
+    def _create_audio_tower(
+        cls, config: UltravoxConfig
+    ) -> Union[transformers.Wav2Vec2Model, "ModifiedWhisperEncoder"]:
+        # We probably don't want to pass tp_plan or device_map to the audio tower
+        # But potentially other kwargs can be passed in. TODO
+        kwargs = {"torch_dtype": config.torch_dtype}
+        if (
+            transformers.modeling_utils._init_weights
+            and config.audio_model_id is not None
+        ):
+            if "whisper" in config.audio_model_id.lower():
+                audio_tower = ModifiedWhisperEncoder.from_pretrained(
+                    config.audio_model_id, **kwargs
+                )
+                audio_tower.init_latency_mask(
+                    config.audio_latency_block_size, dtype=config.torch_dtype
+                )
+            else:
+                assert config.audio_latency_block_size in (
+                    None,
+                    0,
+                ), "only whisper audio tower supports audio latency masking, got non-zero value for 'audio_latency_block_size'"
+                audio_tower = transformers.AutoModel.from_pretrained(
+                    config.audio_model_id, **kwargs
+                )
+        else:
+            with accelerate.init_empty_weights():
+                if "whisper" in config.audio_config._name_or_path.lower():
+                    audio_tower = ModifiedWhisperEncoder(config.audio_config)
+                    audio_tower.init_latency_mask(
+                        config.audio_latency_block_size,
+                        dtype=config.torch_dtype,
+                    )
+                else:
+                    assert config.audio_latency_block_size in (
+                        None,
+                        0,
+                    ), "only whisper audio tower supports audio latency masking, got non-zero value for 'audio_latency_block_size'"
+                    # we only ever use from_config if the weights are retrained, hence initializing is not
+                    # required. This makes the model quite creation faster since init on CPU is quite slow.
+                    audio_tower = transformers.AutoModel.from_config(
+                        config.audio_config, **kwargs
+                    )
+
+        if isinstance(
+            audio_tower,
+            (transformers.Wav2Vec2BertModel, transformers.WhisperModel),
+        ):
+            # For these models we only need the encoder part
+            # Wav2Vec2BertModel -> Wav2Vec2BertEncoder
+            # WhisperModel -> WhisperEncoder
+            audio_tower = audio_tower.encoder
+
+        audio_tower = apply_lora(audio_tower, config.audio_model_lora_config)
+
+        # Convert all parameters to target dtype
+        audio_tower = audio_tower.to(config.torch_dtype)
+
+        return audio_tower
+
+    @classmethod
+    def _create_language_model(
+        cls, config: UltravoxConfig
+    ) -> transformers.LlamaForCausalLM:
+        if (
+            transformers.modeling_utils._init_weights
+            and config.text_model_id is not None
+        ):
+            language_model = transformers.AutoModelForCausalLM.from_pretrained(
+                config.text_model_id,
+                **{
+                    "attn_implementation": config.text_config._attn_implementation,
+                    "torch_dtype": config.torch_dtype,
+                    **FROM_PRETRAINED_KWARGS,
+                },
+            )
+        else:
+            with accelerate.init_empty_weights():
+                # we only ever use from_config if the weights are retrained, hence initializing is not
+                # required. This makes the model quite creation faster since init on CPU is quite slow.
+                language_model = transformers.AutoModelForCausalLM.from_config(
+                    config.text_config,
+                    attn_implementation=config.text_config._attn_implementation,
+                    torch_dtype=config.torch_dtype,
+                )
+
+        language_model = apply_lora(language_model, config.text_model_lora_config)
+        return language_model
+
+    def merge_and_unload(self):
+        if isinstance(self.language_model, peft.PeftModel):
+            self.language_model = self.language_model.merge_and_unload()
+            # no need to download base language model weights anymore, so we can remove the id
+            self.config.text_model_id = None
+            self.keep_params.update(
+                set(
+                    [
+                        f"language_model.{name}"
+                        for name, _ in self.language_model.named_parameters()
+                    ]
+                )
+            )
+
+        if hasattr(self, "audio_tower") and isinstance(
+            self.audio_tower, peft.PeftModel
+        ):
+            self.audio_tower = self.audio_tower.merge_and_unload()
+            # no need to download base audio model weights anymore, so we can remove the id
+            self.config.audio_model_id = None
+
+            self.keep_params.update(
+                set(
+                    [
+                        f"audio_tower.{name}"
+                        for name, _ in self.audio_tower.named_parameters()
+                    ]
+                )
+            )
+
+        for param in ["text_model_lora_config", "audio_model_lora_config"]:
+            if hasattr(self.config, param):
+                delattr(self.config, param)
+
+    def push_to_hub(self, *args, **kwargs):
+        self.merge_and_unload()
+        return super().push_to_hub(*args, **kwargs)
+
+    def diff_state_dict(
+        self, state_dict: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        if state_dict is None:
+            state_dict = super().state_dict()
+
+        trainable_params = {k for k, v in self.named_parameters() if v.requires_grad}
+        # normalize the keys to match the original model
+        # Example: audio_tower.base_model.model.layers.0._fsdp_wrapped_module.self_attn.k_proj.lora_B.default.weight
+        trainable_params = {
+            k.replace("_fsdp_wrapped_module.", "") for k in trainable_params
+        }
+
+        state_dict = {
+            k: v
+            for k, v in state_dict.items()
+            if k in self.keep_params or k in trainable_params
+        }
+
+        return state_dict
+
+    def save_pretrained(
+        self, *args, state_dict: Optional[Dict[str, Any]] = None, **kwargs
+    ):
+        state_dict = self.diff_state_dict(state_dict)
+
+        super().save_pretrained(*args, state_dict=state_dict, **kwargs)
+
+    def _pre_load_state_dict_hook(self, state_dict: Dict[str, Any], *args, **kwargs):
+        self.keep_params.update(set(state_dict.keys()))
+
+    def print_trainable_parameters(self):
+        """
+        Prints the number of trainable parameters in the model (reuses Peft model's method)
+        """
+        count_params = peft.peft_model.PeftModel.get_nb_trainable_parameters
+
+        trainable_params, all_param = count_params(self)
+
+        logging.info(
+            f"trainable params: {trainable_params:,d} || all params: {all_param:,d}"
+            f" || trainable%: {100 * trainable_params / all_param:.1f}%"
+        )
+
+        lm_trainable_params, lm_all_params = count_params(self.language_model)
+        if hasattr(self, "audio_tower") and self.audio_tower is not None:
+            audio_trainable_params, audio_all_params = count_params(self.audio_tower)
+        else:
+            audio_trainable_params, audio_all_params = 0, 0
+
+        projector_trainable_params = (
+            trainable_params - lm_trainable_params - audio_trainable_params
+        )
+        projector_all_params = all_param - lm_all_params - audio_all_params
+
+        # Calculate percentages only if the total parameters are non-zero
+        audio_percent = (
+            0.0
+            if audio_all_params == 0
+            else 100 * audio_trainable_params / audio_all_params
+        )
+        projector_percent = (
+            0.0
+            if projector_all_params == 0
+            else 100 * projector_trainable_params / projector_all_params
+        )
+
+        logging.info(
+            f"Trainable%:   "
+            f" LLM: {100 * lm_trainable_params / lm_all_params:.1f}%"
+            f" || Audio Encoder: {audio_percent:.1f}%"
+            f" || Projector: {projector_percent:.1f}%"
+        )
+
+    def get_fsdp_wrap_module_class(self) -> Optional[type]:
+        """
+        Dynamically finds and returns the FSDP wrap module class from the language model.
+        This is typically the main transformer block class (e.g., LlamaDecoderLayer).
+        It's used by the `transformer_auto_wrap_policy`.
+        """
+        if not hasattr(self.language_model, "_no_split_modules"):
+            return None
+
+        # _no_split_modules is a list of class names, e.g., ["LlamaDecoderLayer"]
+        # We take the first one as it's the main one.
+        module_name_to_find = self.language_model._no_split_modules[0]
+
+        # Recursively search for a module with the desired class name.
+        for module in self.language_model.modules():
+            if module.__class__.__name__ == module_name_to_find:
+                return module.__class__
+
+        logging.warning(
+            f"Could not find FSDP wrap module class '{module_name_to_find}' in the language model."
+        )
+        return None
+
+
+def get_checkpoint_files(
+    model_id: str,
+    local_files_only: bool = False,
+) -> tuple[list[str], dict | None, list[str]]:
+    resolved_archive_file = transformers.utils.cached_file(
+        model_id,
+        transformers.utils.SAFE_WEIGHTS_NAME,
+        _raise_exceptions_for_missing_entries=False,
+        local_files_only=local_files_only,
+    )
+
+    if resolved_archive_file is not None:
+        # not sharded
+        sharded_metadata = None
+        state_dict = transformers.modeling_utils.load_state_dict(resolved_archive_file)
+        loaded_state_dict_keys = list(state_dict.keys())
+    else:
+        # sharded
+        resolved_archive_file = transformers.utils.cached_file(
+            model_id,
+            transformers.utils.SAFE_WEIGHTS_INDEX_NAME,
+            local_files_only=local_files_only,
+        )
+        resolved_archive_file, sharded_metadata = (
+            transformers.modeling_utils.get_checkpoint_shard_files(
+                model_id,
+                resolved_archive_file,
+                local_files_only=local_files_only,
+            )
+        )
+        loaded_state_dict_keys = sharded_metadata["all_checkpoint_keys"]
+
+    if isinstance(resolved_archive_file, str):
+        resolved_archive_file = [resolved_archive_file]
+
+    return resolved_archive_file, sharded_metadata, loaded_state_dict_keys
+
+
+# TODO: refactor common parts to a shared module
+def is_cache_empty(
+    past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]],
+) -> bool:
+    """
+    Check if the cache is empty.
+    """
+    if past_key_values is None:
+        return True
+    if isinstance(past_key_values, tuple):
+        return all(len(c) == 0 for c in past_key_values)
+    return past_key_values.get_seq_length() == 0
+
+
+T = TypeVar("T", bound=torch.nn.Module)
+
+
+def apply_lora(model: T, lora_config: dict) -> T:
+    """
+    Applies LoRA finetuning to the model. If the `r` parameter is set to 0, the model is frozen instead.
+    """
+    unfreeze_layers = lora_config.pop("unfreeze_layers", None)
+    lora_config = peft.LoraConfig(**lora_config or {})
+
+    if lora_config.r == 0:
+        # freeze the model entirely, except for the specified layers
+        for name, param in model.named_parameters():
+            if not unfreeze_layers or not any(
+                re.match(layer, name) for layer in unfreeze_layers
+            ):
+                param.requires_grad = False
+            else:
+                logging.info(f"Unfreezing layer: {name} with #{param.numel()} params")
+    else:
+        model = peft.get_peft_model(model, lora_config)
+
+    return model
+
+
+class StackAudioFrames(nn.Module):
+    """
+    Stack the audio embedding frames to reduce the sequence length by a factor
+    of `stack_factor`.
+    """
+
+    def __init__(self, stack_factor: int = 8):
+        super().__init__()
+        self.stack_factor = stack_factor
+
+    def forward(self, audio_embeds: torch.Tensor) -> torch.Tensor:
+        B, T, C = audio_embeds.shape
+        T_pad = (T + self.stack_factor - 1) // self.stack_factor * self.stack_factor
+        audio_embeds = F.pad(audio_embeds, (0, 0, 0, T_pad - T))
+        B, T, C = audio_embeds.shape
+        audio_embeds = audio_embeds.view(
+            B, T // self.stack_factor, C * self.stack_factor
+        )
+        return audio_embeds
+
+
+class RMSNorm(transformers.models.llama.modeling_llama.LlamaRMSNorm):
+    def __init__(self, hidden_size: int, init: float = 1, eps: float = 1e-6):
+        super().__init__(hidden_size=hidden_size, eps=eps)
+        self.weight.data.fill_(init)
+
+
+class SwiGLU(nn.Module):
+    def forward(self, x):
+        x, gate = x.chunk(2, dim=-1)
+        return F.silu(gate) * x
+
+
+class UltravoxProjector(nn.Module, transformers.modeling_utils.ModuleUtilsMixin):
+    def __init__(self, config: UltravoxConfig):
+        super().__init__()
+        from types import SimpleNamespace
+
+        self.config = SimpleNamespace(is_decoder=False)
+
+        self._pad_and_stack = StackAudioFrames(config.stack_factor)
+        dim_in = config.audio_config.hidden_size * config.stack_factor
+
+        # Create a copy to avoid modifying the original config object in place.
+        projector_audio_config = copy.deepcopy(config.audio_config)
+        # Ensure the projector transformer layers use the desired attention backend
+        # Some backends expect `attn_implementation` (HF-style), others use `_attn_implementation` (internal/Whisper-style).
+        # We set both for compatibility across training and inference.
+        try:
+            attn_impl_value = config.attn_implementation if config.attn_implementation else "eager"
+            setattr(projector_audio_config, "attn_implementation", attn_impl_value)
+            try:
+                setattr(projector_audio_config, "_attn_implementation", attn_impl_value)
+            except Exception:
+                pass
+        except Exception:
+            pass
+
+        self.ln_pre = RMSNorm(dim_in, init=config.norm_init)
+        self.linear_in = nn.Linear(dim_in, projector_audio_config.d_model)
+
+        self.embed_positions = nn.Embedding(
+            projector_audio_config.max_source_positions, projector_audio_config.d_model
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                whisper.WhisperEncoderLayer(projector_audio_config)
+                for _ in range(config.num_projector_layers)
+            ]
+        )
+
+        self.ln_post = RMSNorm(projector_audio_config.d_model, init=config.norm_init)
+        dim_out = config.text_config.hidden_size
+        self.linear_out = nn.Linear(projector_audio_config.d_model, dim_out)
+
+    def forward(
+        self, audio_features: torch.Tensor, audio_token_len: torch.Tensor
+    ) -> torch.Tensor:
+        audio_features = self._pad_and_stack(audio_features)
+
+        max_len_stacked = audio_features.shape[1]
+        attention_mask = torch.arange(max_len_stacked, device=audio_features.device)[
+            None, :
+        ].lt(audio_token_len[:, None])
+        extended_attention_mask = self.get_extended_attention_mask(
+            attention_mask, attention_mask.shape, audio_features.dtype
+        )
+
+        hidden_states = self.ln_pre(audio_features)
+
+        hidden_states = self.linear_in(hidden_states)
+
+        positions = self.embed_positions(
+            torch.arange(hidden_states.size(1), device=hidden_states.device)
+        )
+        hidden_states = hidden_states + positions
+
+        for layer in self.layers:
+            layer_outputs = layer(
+                hidden_states,
+                attention_mask=extended_attention_mask,
+                layer_head_mask=None,
+            )
+            hidden_states = layer_outputs[0]
+
+        hidden_states = self.ln_post(hidden_states)
+        hidden_states = self.linear_out(hidden_states)
+        return hidden_states
+
+class ModifiedWhisperEncoder(
+    whisper.WhisperEncoder, transformers.modeling_utils.ModuleUtilsMixin
+):
+    """
+    Encoder portion of OpenAI's Whisper model.
+
+    This implementation is a slightly modified version of HF Transformers' Whisper Encoder, with only a few fixes:
+    1. base_model_prefix updated to allow for doing `.from_pretrained` directly on the encoder
+    2. allow less than 30 second of audio padding to be passed in:
+        - relaxed ValueError check for `input_features` length to be less than or equal to `expected_seq_length` instead of strictly equal
+        - embed_pos is now sliced to match the length of `inputs_embeds`
+
+    Original: https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/modeling_whisper.py
+    """
+
+    base_model_prefix = "model.encoder"
+    _no_split_modules = ["WhisperEncoderLayer"]
+    _keys_to_ignore_on_load_unexpected = ["model.decoder.*"]
+
+    def __init__(self, config: transformers.WhisperConfig):
+        super().__init__(config)
+        self.config.is_decoder = False
+
+    @property
+    def max_context_length(self):
+        return (
+            self.config.max_source_positions
+            * self.conv1.stride[0]
+            * self.conv2.stride[0]
+        )
+
+    def init_latency_mask(
+        self, audio_latency_block_size: int | None, dtype: torch.dtype
+    ):
+        if audio_latency_block_size is None:
+            self.audio_streaming_mask = None
+            return
+
+        # Use max_context_length directly in the calculation
+        max_seqlen = self.max_context_length
+        assert (
+            max_seqlen > 0
+        ), f"maximum sequence length must be positive, got {max_seqlen}"
+        assert (
+            max_seqlen % audio_latency_block_size == 0
+        ), f"audio_latency_block_size {audio_latency_block_size} must divide {max_seqlen} evenly."
+        # Given the block size, we calculate number of blocks.
+        audio_latency_nblocks = max_seqlen // audio_latency_block_size
+        audio_streaming_mask = (
+            torch.tril(
+                torch.ones(audio_latency_nblocks, audio_latency_nblocks),
+                diagonal=0,
+            )
+            .repeat_interleave(audio_latency_block_size, dim=0)
+            .repeat_interleave(audio_latency_block_size, dim=1)
+        )
+        audio_streaming_mask = (1.0 - audio_streaming_mask) * torch.finfo(dtype).min
+        audio_streaming_mask = audio_streaming_mask[None, None, :, :]
+        self.register_buffer(
+            "audio_streaming_mask", audio_streaming_mask, persistent=False
+        )
+
+    def forward(
+        self,
+        input_features,
+        audio_len=None,
+        head_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        expected_seq_length = self.max_context_length
+        if input_features.shape[-1] > expected_seq_length:
+            raise ValueError(
+                f"Whisper expects the mel input features to be of length {expected_seq_length} or less, but found {input_features.shape[-1]}. Make sure to pad the input mel features to {expected_seq_length}."
+            )
+
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        inputs_embeds = nn.functional.gelu(self.conv1(input_features))
+        inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))
+
+        inputs_embeds = inputs_embeds.permute(0, 2, 1)
+        embed_pos = self.embed_positions.weight[: inputs_embeds.size(-2)]
+
+        hidden_states = inputs_embeds + embed_pos
+        hidden_states = nn.functional.dropout(
+            hidden_states, p=self.dropout, training=self.training
+        )
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        # Create attention mask based on audio lengths to mask out padding tokens
+        # For each sample in batch:
+        # - Convert raw audio length to feature length after convolutions
+        # - Create boolean mask that is True for valid positions and False for padding
+        # - Convert to extended attention mask format expected by transformer layers
+        #   (1.0 for positions to attend to, large negative for positions to ignore)
+        # This masking ensures consistent behavior between training and inference
+        # by preventing the model from attending to padding tokens in both cases
+        attention_mask = None
+        if audio_len is not None:
+            audio_feature_len = self._get_feat_extract_output_lengths(audio_len)
+            max_seq_len = hidden_states.shape[1]
+            attention_mask = torch.arange(max_seq_len, device=hidden_states.device)[
+                None, :
+            ].lt(audio_feature_len.view(-1, 1))
+            attention_mask = self.get_extended_attention_mask(
+                attention_mask,
+                None,
+                dtype=hidden_states.dtype,
+            )
+
+        if self.audio_streaming_mask is not None:
+            seqlen = hidden_states.size(-2)
+            if attention_mask is not None:
+                attention_mask = torch.minimum(
+                    self.audio_streaming_mask[:, :, :seqlen, :seqlen], attention_mask
+                )  # merge
+            else:
+                attention_mask = self.audio_streaming_mask[:, :, :seqlen, :seqlen]
+            attention_mask = attention_mask.to(hidden_states.dtype)
+
+        # check if head_mask has a correct number of layers specified if desired
+        if head_mask is not None:
+            assert head_mask.size()[0] == (
+                len(self.layers)
+            ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
+
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
+            to_drop = False
+            if self.training:
+                dropout_probability = torch.rand([])
+                if dropout_probability < self.layerdrop:  # skip the layer
+                    to_drop = True
+
+            if to_drop:
+                layer_outputs = (None, None)
+            else:
+                if self.gradient_checkpointing and self.training:
+                    layer_outputs = self._gradient_checkpointing_func(
+                        encoder_layer.__call__,
+                        hidden_states,
+                        attention_mask,
+                        (head_mask[idx] if head_mask is not None else None),
+                        output_attentions,
+                    )
+                else:
+                    layer_outputs = encoder_layer(
+                        hidden_states,
+                        attention_mask,
+                        layer_head_mask=(
+                            head_mask[idx] if head_mask is not None else None
+                        ),
+                        output_attentions=output_attentions,
+                    )
+
+                hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        hidden_states = self.layer_norm(hidden_states)
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_states, encoder_states, all_attentions]
+                if v is not None
+            )
+        return transformers.modeling_outputs.BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_states,
+            attentions=all_attentions,
+        )
+
+
+UltravoxConfig.register_for_auto_class()
+UltravoxModel.register_for_auto_class()
+UltravoxProcessor.register_for_auto_class()
+
+transformers.AutoConfig.register("ultravox", UltravoxConfig)
+transformers.AutoModel.register(UltravoxConfig, UltravoxModel)
+transformers.AutoProcessor.register(
+    UltravoxConfig, UltravoxProcessor, UltravoxProcessor
+)
+
+transformers.activations.ACT2FN["swiglu"] = SwiGLU

ultravox_pipeline.py

@@ -0,0 +1,136 @@
+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
+# Even "from . import X" pattern doesn't work (undocumented and unclear why)
+from .ultravox_model import UltravoxModel
+from .ultravox_processing import UltravoxProcessor
+from .ultravox_tokenizer import from_pretrained_text_tokenizer
+from .ultravox_tokenizer import get_audio_token_id
+
+
+class UltravoxPipeline(transformers.Pipeline):
+    def __init__(
+        self,
+        model: UltravoxModel,
+        tokenizer: Optional[transformers.PreTrainedTokenizerBase] = None,
+        audio_processor: Optional[transformers.ProcessorMixin] = None,
+        chat_template: Optional[str] = None,
+        **kwargs,
+    ):
+        if tokenizer is None:
+            try:
+                tokenizer = from_pretrained_text_tokenizer(model.config._name_or_path)
+            except:  # noqa: E722
+                tokenizer = from_pretrained_text_tokenizer(
+                    model.config.text_model_id or model.config.text_config._name_or_path
+                )
+        if chat_template:
+            tokenizer.chat_template = chat_template
+
+        if chat_template:
+            tokenizer.chat_template = chat_template
+
+        model.config.audio_token_index = get_audio_token_id(tokenizer)
+
+        if audio_processor is None:
+            audio_processor = transformers.AutoProcessor.from_pretrained(
+                model.config.audio_model_id or model.config.audio_config._name_or_path
+            )
+
+        super().__init__(model=model, tokenizer=tokenizer, **kwargs)
+
+        self.processor = UltravoxProcessor(
+            audio_processor=audio_processor,
+            tokenizer=tokenizer,
+            stack_factor=model.config.stack_factor,
+            audio_context_size=model.audio_tower_context_length,
+        )
+
+    def _sanitize_parameters(self, **kwargs):
+        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]):
+        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})
+
+        text = self.processor.tokenizer.apply_chat_template(
+            turns, add_generation_prompt=True, tokenize=False
+        )
+
+        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=audio,
+            sampling_rate=inputs.get("sampling_rate", 16000),
+        )
+        if "audio_values" in output:
+            output["audio_values"] = output["audio_values"].to(self.model.dtype)
+
+        return output
+
+    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.pipelines.PIPELINE_REGISTRY.register_pipeline(
+    "ultravox-pipeline",
+    pipeline_class=UltravoxPipeline,
+    pt_model=transformers.AutoModel,
+    type="multimodal",
+)

ultravox_processing.py

@@ -0,0 +1,387 @@
+import dataclasses
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import transformers
+
+from .ultravox_config import UltravoxConfig
+
+
+@dataclasses.dataclass
+class DataCollatorForSeq2SeqWithAudio(transformers.DataCollatorForSeq2Seq):
+    # when enabled, the alt_input_ids, alt_attention_mask, and alt_labels fields are used for computing the KL loss in UltravoxModel
+    include_alt_fields: bool = False
+
+    def __call__(self, features, *args, **kwargs):
+        audio_values = [x for f in features for x in f.pop("audio_values", [])]
+        audio_lens = [x for f in features for x in f.pop("audio_lens", [])]
+        audio_token_len = [x for f in features for x in f.pop("audio_token_len", [])]
+        audio_token_start_idx = [
+            x for f in features for x in f.pop("audio_token_start_idx", [])
+        ]
+
+        if self.include_alt_fields:
+            # these fields are hard-coded in the transformer data collator, so they need special handling before calling the super method
+            alt_features = [
+                {
+                    "input_ids": f.pop("alt_input_ids"),
+                    "attention_mask": f.pop("alt_attention_mask"),
+                    "labels": f.pop("alt_labels"),
+                }
+                for f in features
+            ]
+
+        batch = super().__call__(features, *args, **kwargs)
+        if self.include_alt_fields:
+            alt_batch = super().__call__(alt_features, *args, **kwargs)
+            batch["alt_input_ids"] = alt_batch["input_ids"]
+            batch["alt_attention_mask"] = alt_batch["attention_mask"]
+            batch["alt_labels"] = alt_batch["labels"]
+
+        # Only process audio fields if we have non-empty audio values
+        if audio_values and len(audio_values) > 0 and len(audio_values[0]) > 0:
+            batch["audio_token_start_idx"] = torch.stack(audio_token_start_idx)
+            batch["audio_lens"] = torch.stack(audio_lens)
+            batch["audio_token_len"] = torch.stack(audio_token_len)
+            # Pad the last dimension of all audio_values to the same length, with 0s on the right.
+            max_len = max([x.shape[-1] for x in audio_values])
+            batch["audio_values"] = torch.stack(
+                [F.pad(x, (0, max_len - x.shape[-1])) for x in audio_values]
+            )
+            if self.tokenizer.padding_side == "left":
+                input_ids_lens = torch.LongTensor(
+                    [f["input_ids"].shape[-1] for f in features]
+                )
+                displacement = batch["input_ids"].shape[-1] - input_ids_lens
+                displacement = displacement.repeat_interleave(
+                    batch["audio_batch_size"].squeeze(-1)
+                )
+                batch["audio_token_start_idx"] += displacement.to(
+                    batch["audio_token_start_idx"].device
+                )
+        return batch
+
+
+class UltravoxProcessor(transformers.ProcessorMixin):
+    """
+    Constructs an Ultravox processor which wraps an audio processor and a tokenizer into a single processor.
+
+    Args:
+        audio_processor: The audio processor for the audio encoder.
+        tokenizer: The tokenizer for the language model.
+    """
+
+    attributes = ["audio_processor", "tokenizer"]
+    audio_processor_class = ("WhisperProcessor",)
+    tokenizer_class = (
+        "PreTrainedTokenizer",
+        "PreTrainedTokenizerFast",
+    )
+
+    tokenizer: transformers.PreTrainedTokenizerBase
+    audio_processor: transformers.ProcessorMixin
+
+    def __init__(
+        self,
+        audio_processor=None,
+        tokenizer=None,
+        audio_padding: str = "longest",
+        encoder_ds_factor: int = 2,
+        stack_factor: int = 8,
+        audio_placeholder: str = "<|audio|>",
+        # Defaults to whisper encoder context size
+        audio_context_size: Optional[int] = 3000,
+    ):
+        """
+        Args:
+            audio_processor: The audio processor for the audio encoder.
+            tokenizer: The tokenizer for the language model.
+            audio_padding: The padding strategy for the audio encoder.
+            stack_factor: The factor by which the audio encoder output is stacked in the multimodal projector.
+            encoder_ds_factor: The downsampling factor of the audio encoder.
+            audio_placeholder: The placeholder for the audio in the text.
+            audio_context_size: The maximum number of frames that the audio encoder can handle.
+        """
+        self.audio_padding = audio_padding
+        self.encoder_ds_factor = encoder_ds_factor
+        self.stack_factor = stack_factor
+        self.audio_placeholder = audio_placeholder
+        self.audio_context_size = audio_context_size
+        assert (
+            tokenizer.eos_token is not None
+        ), "The tokenizer has no EOS token. Cannot recover."
+        self.vocab = tokenizer.get_vocab()
+        # VLLM currently relies on updating audio_token_replacement, hence to be safe
+        # we should not update it. This dependency should be removed in the future.
+        self.audio_token_replacement = tokenizer.eos_token
+        if tokenizer.pad_token_id is None:
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+
+        # Use a dummy audio processor to satisfy the base class for text-only training
+        if audio_processor is None:
+            audio_processor = transformers.AutoProcessor.from_pretrained(
+                "openai/whisper-tiny"
+            )
+
+        super().__init__(audio_processor=audio_processor, tokenizer=tokenizer)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        config: UltravoxConfig = transformers.AutoConfig.from_pretrained(
+            pretrained_model_name_or_path, **kwargs
+        )
+        audio_processor = transformers.AutoProcessor.from_pretrained(
+            config.audio_model_id
+            or config.audio_config._name_or_path
+            or "openai/whisper-tiny"
+        )
+
+        tokenizer = transformers.AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path, **kwargs
+        )
+        tokenizer.padding_side = "left"
+        tokenizer.pad_token = tokenizer.eos_token
+
+        return cls(
+            audio_processor=audio_processor,
+            tokenizer=tokenizer,
+            stack_factor=config.stack_factor,
+        )
+
+    def _chunk_and_pad_audio(
+        self,
+        audio_values: torch.Tensor,
+        audio_lens: torch.Tensor,
+        include_audio_num_chunks: bool = False,
+    ) -> Dict[str, Any]:
+        """
+        Processes the audio batch by chunking any items in the batch according to the audio_context_size,
+        padding the last chunk if needed, and returns a dictionary with updated audio data.
+
+        Args:
+            audio_values (torch.Tensor): A tensor of audio values (e.g., in B, D, T format).
+            audio_lens (torch.Tensor): A tensor of audio lengths.
+
+        Returns:
+            Dict[str, Any]: Dictionary with the following keys:
+                - "audio_values": The concatenated audio tensor after chunking and padding.
+                - "audio_lens": Tensor of lengths for each chunk.
+                - "audio_is_continuation": Tensor of booleans indicating if the chunk is a continuation of the previous chunk.
+                - "audio_batch_size": A Tensor with one integer representing the number of chunks.
+
+        """
+        chunked_audio_values: List[torch.Tensor] = []
+        chunked_audio_lens: List[int] = []
+        is_continuation_list: List[bool] = []
+        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
+            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]
+                if is_continuation and chunk.shape[-1] < context_size:
+                    # N.B. We only need to pad continuation chunks. If none of the samples require chunking, the
+                    # batch might not (need to) be padded all the way to the audio_context_size, in which case
+                    # we've already included the padding above. On the other hand, if we have any continuation
+                    # chunks we know that the batch needs to be padded to audio_context_size because that's what
+                    # we're slicing to.
+                    chunk = F.pad(chunk, (0, context_size - chunk.shape[-1]))
+                chunked_audio_values.append(chunk)
+                chunked_audio_lens.append(
+                    min(int(audio_lens[i].item()) - offset, context_size)
+                )
+                is_continuation_list.append(is_continuation)
+
+        data = {
+            "audio_values": torch.stack(chunked_audio_values, dim=0),
+            "audio_lens": torch.tensor(
+                chunked_audio_lens, dtype=torch.int64, device=audio_values.device
+            ),
+            "audio_is_continuation": torch.tensor(
+                is_continuation_list, dtype=torch.bool, device=audio_values.device
+            ),
+            "audio_batch_size": torch.tensor(
+                [len(chunked_audio_values)], device=audio_values.device
+            ),
+        }
+        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,
+        text: Optional[str] = None,
+        audio: Optional[Union[np.ndarray, torch.Tensor]] = None,
+        audios: Optional[
+            Union[
+                List[Union[np.ndarray, torch.Tensor]], Union[np.ndarray, torch.Tensor]
+            ]
+        ] = None,
+        sampling_rate: Optional[int] = None,
+        return_tensors: Optional[
+            Union[str, transformers.TensorType]
+        ] = transformers.TensorType.PYTORCH,
+        include_audio_num_chunks: bool = False,
+        **kwargs,
+    ) -> transformers.BatchFeature:
+        """
+        Main method to prepare for the model one text sequence and audio. This method forwards the `text`
+        and `kwargs` arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizerFast.__call__`] if `text` is not `None` to encode
+        the text. To prepare the audio(s), this method forwards the `audio`, `sampling_rate` and `kwargs` arguments to
+        audio processor's [`~WhisperProcessor.__call__`] if `audio` is not `None`. Please refer to the docstring
+        of the above two methods for more information.
+
+        Args:
+            text (`str`, `List[str]`):
+                The sequence to be encoded. Sequence can be a string or (pretokenized string).
+            audio (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`):
+                The audio to be prepared. Audio can be a single-channel (1-dimensional) NumPy array or PyTorch tensor.
+            audios (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`):
+                A list or two dimensional array of audio to be prepared.
+            sampling_rate (`int`, *optional*, defaults to 16000):
+                Sampling rate of the input audio. We expect 16kHz audio. Don't change this value unless you know what
+                you are doing.
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework. Acceptable values are:
+
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+                - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+        Returns:
+            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+              `None`).
+            - **audio_values** -- Processed audio values to be fed to a model. Returned when `audio` is not `None`.
+            - **audio_token_len** -- Predicted number of audio frames: this value is guaranteed to be a close upper bound.
+              Returned when `audio` is not `None`.
+            - **audio_token_start_idx** -- The index in the tokenized text where the audio starts. Returned when `audio` is not `None`.
+        """
+        # TODO: Add support for multiple text inputs.
+        if audio is not None and audios is not None:
+            raise ValueError("Only one of `audio` or `audios` should be provided.")
+        elif audio is not None:
+            audios = audio if isinstance(audio, list) or audio.ndim == 2 else [audio]
+        elif audios is None:
+            audios = []
+
+        data = {}
+        audio_is_continuation = []
+        if len(audios) > 0:
+            audios = [x.numpy() if isinstance(x, torch.Tensor) else x for x in audios]
+
+            # Pad out each audio to at least 2 hops (the minimum required by the processor).
+            hop_length = self.audio_processor.feature_extractor.hop_length
+            audios = [
+                (
+                    np.pad(x, (0, 2 * hop_length - len(x)), mode="constant")
+                    if len(x) < 2 * hop_length
+                    else x
+                )
+                for x in audios
+            ]
+
+            # Main audio processing. The processor is model-specific.
+            x: transformers.BatchFeature = self.audio_processor(
+                audios,
+                sampling_rate=sampling_rate,
+                padding="longest",
+                pad_to_multiple_of=hop_length,  # The attention mask effectively gets padded to the hop length, so pad the audio to be consistent.
+                truncation=False,
+                return_attention_mask=True,
+                **kwargs,
+            )
+
+            data.update(
+                self._chunk_and_pad_audio(
+                    audio_values=torch.as_tensor(
+                        x.input_features if "input_features" in x else x.input_values
+                    ),
+                    audio_lens=torch.as_tensor(x.attention_mask).sum(-1),
+                    include_audio_num_chunks=include_audio_num_chunks,
+                )
+            )
+
+            audio_is_continuation = data.pop("audio_is_continuation")
+            data["audio_token_len"] = torch.ceil(
+                data["audio_lens"] / (self.encoder_ds_factor * self.stack_factor)
+            ).to(dtype=torch.int)
+
+        if text is not None:
+            if not isinstance(text, str):
+                raise ValueError("Text must be a string. Batch mode not supported yet.")
+
+            # Special tokens like BOS should already have been added by the caller.
+            tokenized_parts = self.tokenizer(
+                text.split(
+                    "<|audio|>"  # The placeholder isn't part of the vocabulary, so split the text around it.
+                ),
+                add_special_tokens=False,
+                **kwargs,
+            )
+
+            audio_token_start_idx = []
+            placeholder_index = -1
+            split_input_ids = tokenized_parts["input_ids"]
+            input_ids: List[int] = []
+
+            audio_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]:
+                    placeholder_index += 1
+                    if placeholder_index >= len(split_input_ids):
+                        raise ValueError(
+                            f"Text contains too few audio placeholders. (Expected {len(audios)} placeholders)"
+                        )
+
+                    input_ids.extend(split_input_ids[placeholder_index])
+
+                audio_token_start_idx.append(len(input_ids))
+
+                input_ids.extend([audio_replacement_token_id] * token_len)
+
+            # Include any tokens after the last audio.
+            placeholder_index += 1
+            if placeholder_index != len(split_input_ids) - 1:
+                raise ValueError(
+                    f"Text contains too many audio placeholders. (Expected {len(audios)} placeholders)"
+                )
+            input_ids.extend(split_input_ids[placeholder_index])
+
+            if "audio_token_len" in data:
+                data["audio_token_start_idx"] = torch.as_tensor(audio_token_start_idx)
+
+            data["input_ids"] = [input_ids]
+            data["attention_mask"] = [[1] * len(input_ids)]
+
+            # Ensure that there are no audio placeholders after the last audio.
+
+        return transformers.BatchFeature(data=data, tensor_type=return_tensors)
+
+    def batch_decode(self, *args, **kwargs):
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        audio_processor_input_names = self.audio_processor.model_input_names
+        return list(set(tokenizer_input_names + audio_processor_input_names))
+
+
+UltravoxProcessor.register_for_auto_class()
+
+transformers.AutoProcessor.register(UltravoxConfig, UltravoxProcessor)

ultravox_tokenizer.py

@@ -0,0 +1,25 @@
+import logging
+
+import transformers
+
+AUDIO_TOKEN = "<|audio|>"
+
+
+def from_pretrained_text_tokenizer(
+    *args, **kwargs
+) -> transformers.PreTrainedTokenizerBase:
+    """
+    Create a tokenizer with the additional special token for audio.
+    This is mainly used for VLLM to work properly. This repo does not currently require it.
+    """
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(*args, **kwargs)
+    tokenizer.add_special_tokens({"additional_special_tokens": [AUDIO_TOKEN]})
+    logging.info(f"Audio token id: {get_audio_token_id(tokenizer)}")
+    return tokenizer
+
+
+def get_audio_token_id(tokenizer: transformers.PreTrainedTokenizerBase) -> int:
+    audio_token_id = tokenizer.encode(AUDIO_TOKEN, add_special_tokens=False)
+    assert len(audio_token_id) == 1, "Audio token should be a single token"
+    return audio_token_id[0]