Update ultravox_config.py

ultravox_config.py

@@ -19,6 +19,8 @@ class LoraConfigSimplified:
     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 LossFunction(str, Enum):
@@ -28,8 +30,10 @@ class LossFunction(str, Enum):
 
 @dataclasses.dataclass
 class LossConfig:
-    loss_function: LossFunction = LossFunction.KL_Divergence
+    loss_function: LossFunction = LossFunction.CrossEntropy
     kl_temperature: float = 2.0
+    # Number of tokens to ignore from the beginning of the sequence. Only used in LSM
+    initial_tokens_to_ignore: int = 0
 
     @property
     def requires_alt_fields(self):
@@ -45,7 +49,7 @@ class UltravoxConfig(transformers.PretrainedConfig):
     documentation from [`PretrainedConfig`] for more information.
 
     Args:
-        audio_config (`Wav2Vec2Config`,  *optional*):
+        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`.
@@ -63,15 +67,17 @@ class UltravoxConfig(transformers.PretrainedConfig):
             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.
 
 
     Example:
 
     ```python
-    >>> from transformers import UltravoxForConditionalGeneration, Wav2Vec2Config, UltravoxConfig, LlamaConfig
+    >>> from transformers import UltravoxModel, WhisperConfig, UltravoxConfig, LlamaConfig
 
     >>> # Initializing an audio encoder config
-    >>> audio_config = Wav2Vec2Config()
+    >>> audio_config = WhisperConfig()
 
     >>> # Initializing a Llama config
     >>> text_config = LlamaConfig()
@@ -80,13 +86,13 @@ class UltravoxConfig(transformers.PretrainedConfig):
     >>> configuration = UltravoxConfig(audio_config, text_config)
 
     >>> # Initializing a completely untrained model from the configuration
-    >>> model = UltravoxForConditionalGeneration(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="facebook/wav2vec2-base-960h", text_model_id="meta-llama/Llama-2-7b-chat-hf")
+    >>> config = UltravoxConfig(audio_model_id="openai/whisper-tiny", text_model_id="meta-llama/Llama-2-7b-chat-hf")
     ```"""
 
     model_type = "ultravox"
@@ -103,8 +109,10 @@ class UltravoxConfig(transformers.PretrainedConfig):
         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
         text_model_lora_config: Optional[LoraConfigSimplified] = None,
         audio_model_lora_config: Optional[LoraConfigSimplified] = None,
+        audio_latency_block_size: Optional[int] = None,
         **kwargs,
     ):
         self.ignore_index = ignore_index
@@ -116,7 +124,7 @@ class UltravoxConfig(transformers.PretrainedConfig):
         self.stack_factor = stack_factor
         self.norm_init = norm_init
         self.projector_act = projector_act
-
+        self.projector_ln_mid = projector_ln_mid
         if text_model_id is not None:
             self.text_config: transformers.LlamaConfig = (
                 transformers.AutoConfig.from_pretrained(text_model_id)
@@ -134,7 +142,7 @@ class UltravoxConfig(transformers.PretrainedConfig):
         else:
             audio_config = audio_config or {}
             self.audio_config = transformers.CONFIG_MAPPING[
-                audio_config.get("model_type", "wav2vec2")
+                audio_config.get("model_type", "whisper")
             ](**audio_config)
 
         self.text_model_lora_config = (
@@ -147,6 +155,7 @@ class UltravoxConfig(transformers.PretrainedConfig):
             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
 
         self.vocab_size = self.text_config.vocab_size
 
@@ -160,7 +169,12 @@ class UltravoxConfig(transformers.PretrainedConfig):
         # 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