main
py 120 lines 4.56 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from __future__ import annotations
5
6 from transformers import CONFIG_MAPPING, PreTrainedConfig
7
8 # NOTE: Temporary shim for FunAudioChat checkpoints.
9 # These checkpoints use `model_type="funaudiochat"`, which is not currently
10 # recognized by released Transformers, and the public checkpoint does not
11 # provide an `auto_map` to enable `trust_remote_code=True`.
12 # Remove this file once Transformers adds native support (or the checkpoint
13 # provides an `auto_map`) and vLLM can rely on `AutoConfig.from_pretrained()`.
14
15
16 class FunAudioChatAudioEncoderConfig(PreTrainedConfig):
17 model_type = "funaudiochat_audio_encoder"
18
19 def __init__(
20 self,
21 _attn_implementation: str | None = None,
22 num_mel_bins: int = 128,
23 encoder_layers: int = 32,
24 encoder_attention_heads: int = 20,
25 encoder_ffn_dim: int = 5120,
26 d_model: int = 1280,
27 dropout: float = 0.0,
28 attention_dropout: float = 0.0,
29 activation_function: str = "gelu",
30 activation_dropout: float = 0.0,
31 scale_embedding: bool = False,
32 initializer_range: float = 0.02,
33 max_source_positions: int = 1500,
34 n_window: int = 100,
35 output_dim: int = 3584,
36 bos_token_id: int | None = None,
37 codebook_size: int | None = None,
38 continuous_features_mode: str = "replace",
39 crq_transformer_config: dict | None = None,
40 eos_token_id: int | None = None,
41 group_size: int = 5,
42 enable_audio_invert_tower: bool = True,
43 pad_token_id: int | None = None,
44 **kwargs,
45 ) -> None:
46 attn_impl = kwargs.pop("_attn_implementation", None) or _attn_implementation
47 super().__init__(**kwargs)
48 # Match HF default for attention implementation selection.
49 self._attn_implementation = attn_impl or "sdpa"
50
51 self.num_mel_bins = num_mel_bins
52 self.d_model = d_model
53 self.encoder_layers = encoder_layers
54 self.encoder_attention_heads = encoder_attention_heads
55 self.encoder_ffn_dim = encoder_ffn_dim
56 self.dropout = dropout
57 self.attention_dropout = attention_dropout
58 self.activation_function = activation_function
59 self.activation_dropout = activation_dropout
60 self.num_hidden_layers = encoder_layers
61 self.initializer_range = initializer_range
62 self.scale_embedding = scale_embedding
63 self.max_source_positions = max_source_positions
64 self.n_window = n_window
65 self.output_dim = output_dim
66
67 self.bos_token_id = bos_token_id
68 self.codebook_size = codebook_size
69 self.continuous_features_mode = continuous_features_mode
70 self.crq_transformer_config = crq_transformer_config
71 self.eos_token_id = eos_token_id
72 self.group_size = group_size
73 self.enable_audio_invert_tower = enable_audio_invert_tower
74 self.pad_token_id = pad_token_id
75
76
77 class FunAudioChatConfig(PreTrainedConfig):
78 model_type = "funaudiochat"
79 attribute_map = {
80 "audio_token_id": "audio_token_index",
81 }
82
83 def __init__(
84 self,
85 audio_config: PreTrainedConfig | dict | None = None,
86 text_config: PreTrainedConfig | dict | None = None,
87 audio_token_index: int = 151646,
88 ignore_index: int = -100,
89 hidden_size: int | None = None,
90 **kwargs,
91 ) -> None:
92 self.audio_token_index = audio_token_index
93 self.ignore_index = ignore_index
94
95 if audio_config is None:
96 self.audio_config = FunAudioChatAudioEncoderConfig()
97 elif isinstance(audio_config, dict):
98 default_model_type = FunAudioChatAudioEncoderConfig.model_type
99 audio_config.setdefault("model_type", default_model_type)
100 self.audio_config = FunAudioChatAudioEncoderConfig(**audio_config)
101 else:
102 self.audio_config = audio_config
103
104 if text_config is None:
105 self.text_config = CONFIG_MAPPING["qwen2"]()
106 elif isinstance(text_config, dict):
107 # Default to qwen2 for backwards compatibility; FunAudioChat uses
108 # qwen3 in practice for recent checkpoints.
109 text_config.setdefault("model_type", "qwen2")
110 self.text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config)
111 else:
112 self.text_config = text_config
113
114 self.hidden_size = (
115 int(self.text_config.hidden_size)
116 if hidden_size is None
117 else int(hidden_size)
118 )
119
120 super().__init__(**kwargs)