File size: 4,670 Bytes
ce6d303 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | from typing import Optional, List
from dataclasses import dataclass, field
from transformers import PretrainedConfig, Qwen3Config
@dataclass
class MossAudioEncoderConfig:
d_model: int = 1280
output_dim: int = 1280
num_mel_bins: int = 128
encoder_layers: int = 32
encoder_attention_heads: int = 20
encoder_ffn_dim: int = 5120
downsample_rate: int = 8
downsample_hidden_size: int = 480
encoder_attention_window_size: int = 100
max_source_positions: int = 1500
dropout: float = 0.1
attention_dropout: float = 0.1
activation_dropout: float = 0.0
activation_function: str = "gelu"
layer_norm_eps: float = 1e-5
_attn_implementation: str = "eager"
pretrained_path: str = ""
n_window: int = 200
conv_chunksize: int = 64
deepstack_encoder_layer_indexes: List[int] = field(default_factory=lambda: [8, 16, 24])
@classmethod
def from_dict(cls, config_dict):
if config_dict is None:
return cls()
allowed_keys = set(cls.__dataclass_fields__.keys())
filtered = {k: v for k, v in config_dict.items() if k in allowed_keys}
return cls(**filtered)
def to_dict(self):
return {
"d_model": self.d_model,
"output_dim": self.output_dim,
"num_mel_bins": self.num_mel_bins,
"encoder_layers": self.encoder_layers,
"encoder_attention_heads": self.encoder_attention_heads,
"encoder_ffn_dim": self.encoder_ffn_dim,
"downsample_rate": self.downsample_rate,
"downsample_hidden_size": self.downsample_hidden_size,
"encoder_attention_window_size": self.encoder_attention_window_size,
"max_source_positions": self.max_source_positions,
"dropout": self.dropout,
"attention_dropout": self.attention_dropout,
"activation_dropout": self.activation_dropout,
"activation_function": self.activation_function,
"layer_norm_eps": self.layer_norm_eps,
"_attn_implementation": self._attn_implementation,
"pretrained_path": self.pretrained_path,
"n_window": self.n_window,
"conv_chunksize": self.conv_chunksize,
"deepstack_encoder_layer_indexes": list(self.deepstack_encoder_layer_indexes or []),
}
class MossAudioConfig(PretrainedConfig):
model_type = "moss_audio"
is_composition = True
def __init__(
self,
audio_config=None,
language_config=None,
adapter_hidden_size=8192,
ignore_index=-100,
deepstack_num_inject_layers: Optional[int] = None,
**kwargs,
):
if isinstance(audio_config, dict):
audio_config = MossAudioEncoderConfig.from_dict(audio_config)
elif audio_config is None:
audio_config = MossAudioEncoderConfig()
if isinstance(language_config, dict):
language_config = Qwen3Config(**language_config)
elif language_config is None:
language_config = Qwen3Config()
self.audio_config = audio_config
self.language_config = language_config
self.adapter_hidden_size = adapter_hidden_size
self.ignore_index = ignore_index
self.deepstack_num_inject_layers = deepstack_num_inject_layers
_propagate_keys = {
"num_hidden_layers", "eos_token_id", "bos_token_id", "vocab_size",
"tie_word_embeddings",
}
for key in ("num_hidden_layers", "eos_token_id", "bos_token_id", "vocab_size"):
kwargs.setdefault(key, getattr(language_config, key, None))
kwargs.setdefault("tie_word_embeddings", False)
if hasattr(language_config, "to_dict"):
_lang_keys = set(language_config.to_dict().keys())
for key in list(kwargs.keys()):
if key in _lang_keys and key not in _propagate_keys:
kwargs.pop(key)
super().__init__(**kwargs)
def to_dict(self):
output = super().to_dict()
output["audio_config"] = (
self.audio_config.to_dict() if hasattr(self.audio_config, "to_dict") else self.audio_config
)
output["language_config"] = (
self.language_config.to_dict()
if hasattr(self.language_config, "to_dict")
else self.language_config
)
output["adapter_hidden_size"] = self.adapter_hidden_size
output["ignore_index"] = self.ignore_index
output["deepstack_num_inject_layers"] = self.deepstack_num_inject_layers
return output
__all__ = ["MossAudioEncoderConfig", "MossAudioConfig"]
|