OpenMOSS-Team/MOSS-Audio-Tokenizer

1# coding=utf-82# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""MossAudioTokenizer model configuration"""16 17from typing import Any18 19from transformers.configuration_utils import PreTrainedConfig20from transformers.utils import logging21 22 23logger = logging.get_logger(__name__)24 25 26class MossAudioTokenizerConfig(PreTrainedConfig):27    r"""28    This is the configuration class to store the configuration of a [`MossAudioTokenizerModel`]. It is used to instantiate a29    MossAudioTokenizer model according to the specified arguments, defining the model architecture.30 31    Instantiating a configuration with the defaults will yield a similar configuration to that of the32    [VoiceAgentGroup/moss_audio_tokenizer](https://huggingface.co/VoiceAgentGroup/moss_audio_tokenizer) architecture.33 34    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the35    documentation from [`PreTrainedConfig`] for more information.36 37    Args:38        sampling_rate (`int`, *optional*, defaults to 24000):39            The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz).40        downsample_rate (`int`, *optional*, defaults to 1920):41            Total downsampling rate from waveform to tokens.42        causal_transformer_context_duration (`float`, *optional*, defaults to 10.0):43            Context duration in seconds for causal transformer.44        encoder_kwargs (`list[dict]`, *optional*):45            List of encoder module configurations. Each dict specifies a module type and its parameters.46        decoder_kwargs (`list[dict]`, *optional*):47            List of decoder module configurations in execution order.48        quantizer_type (`str`, *optional*, defaults to `"rvq"`):49            Quantizer type. Options include `"rvq"`, `"spec_rvq"`, `"rlfq"`, `"random_prefix_rlfq"`.50        quantizer_kwargs (`dict`, *optional*):51            Configuration for the quantizer including `input_dim`, `rvq_dim`, `output_dim`, `num_quantizers`,52            `codebook_size`, and `codebook_dim`.53 54    Example:55 56    ```python57    >>> from transformers import MossAudioTokenizerModel, MossAudioTokenizerConfig58 59    >>> # Initializing a MossAudioTokenizer style configuration60    >>> configuration = MossAudioTokenizerConfig()61 62    >>> # Initializing a model (with random weights) from the configuration63    >>> model = MossAudioTokenizerModel(configuration)64 65    >>> # Accessing the model configuration66    >>> configuration = model.config67    ```68    """69 70    model_type = "moss-audio-tokenizer"71 72    # Backward-compatible alias used by some checkpoints.73    attribute_map = {"sample_rate": "sampling_rate"}74 75    sampling_rate: int76    downsample_rate: int77    causal_transformer_context_duration: float78    encoder_kwargs: list[dict[str, Any]]79    decoder_kwargs: list[dict[str, Any]]80    quantizer_type: str81    quantizer_kwargs: dict[str, Any]82 83    def __init__(84        self,85        version: str | None = None,86        sampling_rate: int = 24000,87        downsample_rate: int = 1920,88        causal_transformer_context_duration: float = 10.0,89        encoder_kwargs: list[dict[str, Any]] | None = None,90        decoder_kwargs: list[dict[str, Any]] | None = None,91        quantizer_type: str = "rlfq",92        quantizer_kwargs: dict[str, Any] | None = None,93        **kwargs,94    ):95        # Some checkpoints might include an incorrect/legacy `model_type` (e.g. "speech_tokenizer").96        # We drop it to avoid overriding the class-level `model_type`.97        kwargs.pop("model_type", None)98 99        # `version` is accepted for compatibility but not used in modeling.100        self.version = version101        self.sampling_rate = sampling_rate102        self.downsample_rate = downsample_rate103        self.causal_transformer_context_duration = causal_transformer_context_duration104        # Default encoder configuration105        if encoder_kwargs is None:106            encoder_kwargs = [107                {108                    "module_type": "PatchedPretransform",109                    "patch_size": 240,110                },111                {112                    "module_type": "Transformer",113                    "input_dimension": 240,114                    "output_dimension": 384,115                    "d_model": 768,116                    "num_heads": 12,117                    "num_layers": 12,118                    "dim_feedforward": 3072,119                    "causal": True,120                    "norm": "layer_norm",121                    "positional_embedding": "rope",122                    "max_period": 10000,123                    "gating": "none",124                    "layer_scale": 0.01,125                    "conv_layout": True,126                },127                {128                    "module_type": "PatchedPretransform",129                    "patch_size": 2,130                },131                {132                    "module_type": "Transformer",133                    "input_dimension": 768,134                    "output_dimension": 384,135                    "d_model": 768,136                    "num_heads": 12,137                    "num_layers": 12,138                    "dim_feedforward": 3072,139                    "causal": True,140                    "norm": "layer_norm",141                    "positional_embedding": "rope",142                    "max_period": 10000,143                    "gating": "none",144                    "layer_scale": 0.01,145                    "conv_layout": True,146                },147                {148                    "module_type": "PatchedPretransform",149                    "patch_size": 2,150                },151                {152                    "module_type": "Transformer",153                    "input_dimension": 768,154                    "output_dimension": 640,155                    "d_model": 768,156                    "num_heads": 12,157                    "num_layers": 12,158                    "dim_feedforward": 3072,159                    "causal": True,160                    "norm": "layer_norm",161                    "positional_embedding": "rope",162                    "max_period": 10000,163                    "gating": "none",164                    "layer_scale": 0.01,165                    "conv_layout": True,166                },167                {168                    "module_type": "PatchedPretransform",169                    "patch_size": 2,170                },171                {172                    "module_type": "Transformer",173                    "input_dimension": 1280,174                    "output_dimension": 768,175                    "d_model": 1280,176                    "num_heads": 20,177                    "num_layers": 32,178                    "dim_feedforward": 5120,179                    "causal": True,180                    "norm": "layer_norm",181                    "positional_embedding": "rope",182                    "max_period": 10000,183                    "gating": "none",184                    "layer_scale": 0.01,185                    "conv_layout": True,186                },187            ]188        self.encoder_kwargs = encoder_kwargs189 190        # Default decoder configuration (execution order)191        if decoder_kwargs is None:192            decoder_kwargs = [193                {194                    "module_type": "Transformer",195                    "input_dimension": 768,196                    "output_dimension": 1280,197                    "d_model": 1280,198                    "num_heads": 20,199                    "num_layers": 32,200                    "dim_feedforward": 5120,201                    "causal": True,202                    "norm": "layer_norm",203                    "positional_embedding": "rope",204                    "max_period": 10000,205                    "gating": "none",206                    "layer_scale": 0.01,207                    "conv_layout": True,208                },209                {210                    "module_type": "PatchedPretransform",211                    "patch_size": 2,212                },213                {214                    "module_type": "Transformer",215                    "input_dimension": 640,216                    "output_dimension": 768,217                    "d_model": 768,218                    "num_heads": 12,219                    "num_layers": 12,220                    "dim_feedforward": 3072,221                    "causal": True,222                    "norm": "layer_norm",223                    "positional_embedding": "rope",224                    "max_period": 10000,225                    "gating": "none",226                    "layer_scale": 0.01,227                    "conv_layout": True,228                },229                {230                    "module_type": "PatchedPretransform",231                    "patch_size": 2,232                },233                {234                    "module_type": "Transformer",235                    "input_dimension": 384,236                    "output_dimension": 768,237                    "d_model": 768,238                    "num_heads": 12,239                    "num_layers": 12,240                    "dim_feedforward": 3072,241                    "causal": True,242                    "norm": "layer_norm",243                    "positional_embedding": "rope",244                    "max_period": 10000,245                    "gating": "none",246                    "layer_scale": 0.01,247                    "conv_layout": True,248                },249                {250                    "module_type": "PatchedPretransform",251                    "patch_size": 2,252                },253                {254                    "module_type": "Transformer",255                    "input_dimension": 384,256                    "output_dimension": 768,257                    "d_model": 768,258                    "num_heads": 12,259                    "num_layers": 12,260                    "dim_feedforward": 3072,261                    "causal": True,262                    "norm": "layer_norm",263                    "positional_embedding": "rope",264                    "max_period": 10000,265                    "gating": "none",266                    "layer_scale": 0.01,267                    "conv_layout": True,268                },269                {270                    "module_type": "PatchedPretransform",271                    "patch_size": 2,272                },273                {274                    "module_type": "Transformer",275                    "input_dimension": 384,276                    "output_dimension": 240,277                    "d_model": 768,278                    "num_heads": 12,279                    "num_layers": 12,280                    "dim_feedforward": 3072,281                    "causal": True,282                    "norm": "layer_norm",283                    "positional_embedding": "rope",284                    "max_period": 10000,285                    "gating": "none",286                    "layer_scale": 0.01,287                    "conv_layout": True,288                },289                {290                    "module_type": "PatchedPretransform",291                    "patch_size": 240,292                },293            ]294        self.decoder_kwargs = decoder_kwargs295 296        # Default quantizer configuration297        if quantizer_kwargs is None:298            quantizer_kwargs = {299                "input_dim": 768,300                "rvq_dim": 512,301                "output_dim": 768,302                "num_quantizers": 32,303                "codebook_size": 1024,304                "codebook_dim": 8,305                "quantizer_type": "rlfq",306            }307 308        # Handle quantizer_type from kwargs or config309        kw_qtype = quantizer_kwargs.get("quantizer_type", None)310        if kw_qtype is not None:311            self.quantizer_type = kw_qtype312        else:313            self.quantizer_type = quantizer_type314            quantizer_kwargs["quantizer_type"] = quantizer_type315 316        self.quantizer_kwargs = quantizer_kwargs317 318        super().__init__(**kwargs)319 320    @property321    def num_quantizers(self) -> int:322        """Return the number of quantizers from quantizer_kwargs."""323        return self.quantizer_kwargs.get("num_quantizers", 32)324 325    @property326    def codebook_size(self) -> int:327        """Return the codebook size from quantizer_kwargs."""328        return self.quantizer_kwargs.get("codebook_size", 4096)329 330    @property331    def frame_rate(self) -> float:332        """Return the frame rate (tokens per second)."""333        return self.sampling_rate / self.downsample_rate334 335 336__all__ = ["MossAudioTokenizerConfig"]337