Upload 3 files

configuration_ced.py

@@ -0,0 +1,137 @@
+# coding=utf-8
+# Copyright 2023 Xiaomi Corporation and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" CED model configuration"""
+
+
+from transformers import PretrainedConfig
+from transformers.utils import logging
+from transformers.utils.hub import cached_file
+
+logger = logging.get_logger(__name__)
+
+CED_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "mispeech/ced-tiny": "https://huggingface.co/mispeech/ced-tiny/resolve/main/config.json",
+}
+
+
+class CedConfig(PretrainedConfig):
+    model_type = "ced"
+
+    r"""
+    Configuration class for the CED model.
+
+    Args:
+        name (str, optional, *optional*):
+            Name of the pre-defined configuration. Can be "ced-tiny", "ced-mini", "ced-small" or "ced-base".
+        attn_drop_rate (float, *optional*, defaults to 0.0):
+            Dropout probability for attention weights. Default to 0.0.
+        depth (int, *optional*, defaults to 12): Number of transformer layers. Default to 12.
+        drop_path_rate (float, *optional*, defaults to 0.0): Drop path is taken from timm. Default to 0.0.
+        drop_rate (float, *optional*, defaults to 0.0):
+            Dropout probability for input embeddings. Default to 0.0.
+        embed_dim (int, *optional*, defaults to 768):
+            Dimensionality of the audio patch embeddings. Default to 768.
+        eval_avg (str, *optional*, defaults to `"mean"`):
+            Type of pooling to use for evaluation. Can be "mean", "token", "dm" or "logit". Default to "mean".
+        mlp_ratio (float, *optional*, defaults to 4.0):
+            Ratio of hidden size in the feedforward layer to the embedding size. Default to 4.0.
+        num_heads (int, *optional*, defaults to 12): Number of attention heads. Default to 12.
+        outputdim (int, *optional*, defaults to 527): Dimensionality of the output. Default to 527.
+        patch_size (int, *optional*, defaults to 16): Size of the patches. Default to 16.
+        patch_stride (int, *optional*, defaults to 16): Stride of the patches. Default to 16.
+        pooling (str, *optional*, defaults to `"mean"`):
+            Type of pooling to use for the output. Can be "mean", "token", "dm" or "logit". Default to "mean".
+        qkv_bias (bool, *optional*, defaults to `True`):
+            Whether to include bias terms in the query, key and value projections. Default to True.
+        target_length (int, *optional*, defaults to 1012): Frames of an audio chunk. Default to 1012.
+    """
+
+    def __init__(
+        self,
+        name=None,
+        attn_drop_rate=0.0,
+        depth=12,
+        drop_path_rate=0.0,
+        drop_rate=0.0,
+        embed_dim=768,
+        eval_avg="mean",
+        mlp_ratio=4.0,
+        num_heads=12,
+        outputdim=527,
+        patch_size=16,
+        patch_stride=16,
+        pooling="mean",
+        qkv_bias=True,
+        target_length=1012,
+        **kwargs,
+    ):
+        r"""
+        TODO: Add docstring
+        """
+
+        super().__init__(**kwargs)
+
+        if name == "ced-tiny":
+            embed_dim = 192
+            num_heads = 3
+        elif name == "ced-mini":
+            embed_dim = 256
+            num_heads = 4
+        elif name == "ced-small":
+            embed_dim = 384
+            num_heads = 6
+        elif name == "ced-base":
+            embed_dim = 768
+            num_heads = 12
+        else:
+            logger.info("No model name specified for CedConfig, use default settings.")
+
+        assert pooling in ("mean", "token", "dm", "logit")
+        self.name = name
+        self.attn_drop_rate = attn_drop_rate
+        self.center = kwargs.get("center", True)
+        self.depth = depth
+        self.drop_path_rate = drop_path_rate
+        self.drop_rate = drop_rate
+        self.embed_dim = embed_dim
+        self.eval_avg = eval_avg
+        self.f_max = kwargs.get("f_max", 8000)
+        self.f_min = kwargs.get("f_min", 0)
+        self.hop_size = kwargs.get("hop_size", 160)
+        self.mlp_ratio = mlp_ratio
+        self.n_fft = kwargs.get("n_fft", 512)
+        self.n_mels = kwargs.get("n_mels", 64)
+        self.n_mels = kwargs.get("n_mels", 64)
+        self.num_heads = num_heads
+        self.outputdim = outputdim
+        self.pad_last = kwargs.get("pad_last", True)
+        self.patch_size = patch_size
+        self.patch_stride = patch_stride
+        self.pooling = pooling
+        self.qkv_bias = qkv_bias
+        self.target_length = target_length
+        self.win_size = kwargs.get("win_size", 512)
+        self.loss = "BCE"
+
+        if self.outputdim == 527:
+            with open(cached_file("topel/ConvNeXt-Tiny-AT", "class_labels_indices.csv"), "r") as f:
+                self.id2label = {
+                    int(line.split(",", maxsplit=3)[0]): line.split(",", maxsplit=3)[2].replace('"', "").strip("\n")
+                    for line in f.readlines()[1:]
+                }
+            self.label2id = {v: k for k, v in self.id2label.items()}
+        else:
+            self.id2label = None
+            self.label2id = None

feature_extraction_ced.py

@@ -0,0 +1,166 @@
+# coding=utf-8
+# Copyright 2023 Xiaomi Corporation and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Feature extractor class for CED.
+"""
+
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+import torchaudio.transforms as audio_transforms
+
+from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class CedFeatureExtractor(SequenceFeatureExtractor):
+    r"""
+    CedFeatureExtractor extracts Mel spectrogram features from audio signals.
+
+    Args:
+        f_min (int, *optional*, defaults to 0): Minimum frequency for the Mel filterbank.
+        sampling_rate (int, *optional*, defaults to 16000):
+            Sampling rate of the input audio signal.
+        win_size (int, *optional*, defaults to 512): Window size for the STFT.
+        center (bool, *optional*, defaults to `True`):
+            Whether to pad the signal on both sides to center it.
+        n_fft (int, *optional*, defaults to 512): Number of FFT points for the STFT.
+        f_max (int, optional, *optional*): Maximum frequency for the Mel filterbank.
+        hop_size (int, *optional*, defaults to 160): Hop size for the STFT.
+        feature_size (int, *optional*, defaults to 64): Number of Mel bands to generate.
+        padding_value (float, *optional*, defaults to 0.0): Value for padding.
+
+    Returns:
+        BatchFeature: A BatchFeature object containing the extracted features.
+    """
+
+    def __init__(
+        self,
+        f_min: int = 0,
+        sampling_rate: int = 16000,
+        win_size: int = 512,
+        center: bool = True,
+        n_fft: int = 512,
+        f_max: Optional[int] = None,
+        hop_size: int = 160,
+        feature_size: int = 64,
+        padding_value: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(
+            feature_size=feature_size,
+            sampling_rate=sampling_rate,
+            padding_value=padding_value,
+            **kwargs,
+        )
+        self.f_min = f_min
+        self.win_size = win_size
+        self.center = center
+        self.n_fft = n_fft
+        self.f_max = f_max
+        self.hop_size = hop_size
+
+        self.model_input_names = ["input_values"]
+
+    def __call__(
+        self,
+        x: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]],
+        sampling_rate: Optional[int] = None,
+        max_length: Optional[int] = None,
+        truncation: bool = False,
+        return_tensors="pt",
+    ) -> BatchFeature:
+        r"""
+        Extracts Mel spectrogram features from an audio signal tensor.
+
+        Args:
+            x: Input audio signal tensor.
+            sampling_rate (int, *optional*, defaults to `None`):
+                Sampling rate of the input audio signal.
+            max_length (int, *optional*, defaults to None):
+                Maximum length of the input audio signal.
+            truncation (bool, *optional*, defaults to `False`):
+                Whether to truncate the input signal to max_length.
+            return_tensors (str, *optional*, defaults to "pt"):
+                If set to "pt", the return type will be a PyTorch tensor.
+
+        Returns:
+            BatchFeature: A dictionary containing the extracted features.
+        """
+        if sampling_rate is None:
+            sampling_rate = self.sampling_rate
+
+        if return_tensors != "pt":
+            raise NotImplementedError("Only return_tensors='pt' is currently supported.")
+
+        mel_spectrogram = audio_transforms.MelSpectrogram(
+            f_min=self.f_min,
+            sample_rate=sampling_rate,
+            win_length=self.win_size,
+            center=self.center,
+            n_fft=self.n_fft,
+            f_max=self.f_max,
+            hop_length=self.hop_size,
+            n_mels=self.feature_size,
+        )
+        amplitude_to_db = audio_transforms.AmplitudeToDB(top_db=120)
+
+        if isinstance(x, np.ndarray):
+            if x.ndim == 1:
+                x = x[np.newaxis, :]
+            if x.ndim != 2:
+                raise ValueError("np.ndarray input must be a 1D or 2D.")
+            x = torch.from_numpy(x)
+        elif isinstance(x, torch.Tensor):
+            if x.dim() == 1:
+                x = x.unsqueeze(0)
+            if x.dim() != 2:
+                raise ValueError("torch.Tensor input must be a 1D or 2D.")
+        elif isinstance(x, (list, tuple)):
+            longest_length = max(x_.shape[0] for x_ in x)
+            if not truncation and max_length is not None and max_length < longest_length:
+                max_length = longest_length
+            if not truncation and max_length is None:
+                max_length = longest_length
+
+
+            if all(isinstance(x_, np.ndarray) for x_ in x):
+                if not all(x_.ndim == 1 for x_ in x):
+                    raise ValueError("All np.ndarray in a list must be 1D.")
+
+                x_trim = [x_[:max_length] for x_ in x]
+                x_pad = [np.pad(x_, (0, max_length - x_.shape[0]), mode="constant", constant_values=0) for x_ in x_trim]
+                x = torch.stack([torch.from_numpy(x_) for x_ in x_pad])
+            elif all(isinstance(x_, torch.Tensor) for x_ in x):
+                if not all(x_.dim() == 1 for x_ in x):
+                    raise ValueError("All torch.Tensor in a list must be 1D.")
+                x_pad = [torch.nn.functional.pad(x_, (0, max_length - x_.shape[0]), value=0) for x_ in x]
+                x = torch.stack(x_pad)
+            else:
+                raise ValueError("Input list must be numpy arrays or PyTorch tensors.")
+        else:
+            raise ValueError(
+                "Input must be a numpy array, a list of numpy arrays, a PyTorch tensor, or a list of PyTorch tensor."
+            )
+
+        x = x.float()
+        x = mel_spectrogram(x)
+        x = amplitude_to_db(x)
+        return BatchFeature({"input_values": x})

modeling_ced.py

@@ -0,0 +1,549 @@
+# coding=utf-8
+# Copyright 2023 Xiaomi Corporation and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch CED (Ced) model."""
+
+import collections
+import math
+from functools import partial
+from typing import Any, Callable, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+
+from transformers.modeling_outputs import SequenceClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    add_code_sample_docstrings,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+)
+from .configuration_ced import CedConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "CedConfig"
+_SEQ_CLASS_EXPECTED_OUTPUT = "'Speech synthesizer'"
+_SEQ_CLASS_EXPECTED_LOSS = 0.69
+
+# Audio classification docstring
+_SEQ_CLASS_CHECKPOINT = "mispeech/ced-tiny"
+
+
+CED_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "mispeech/ced-tiny",
+    "mispeech/ced-mini",
+    "mispeech/ced-small",
+    "mispeech/ced-base",
+    # See all CED models at https://huggingface.co/models?search=mispeech%2Fced
+]
+
+
+class CedPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = CedConfig
+    base_model_prefix = "ced"
+    main_input_name = "input_values"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.constant_(module.bias, 0)
+            nn.init.constant_(module.weight, 1.0)
+
+
+Conv_Kernel = Union[int, Tuple[int, int]]
+
+
+def to_2tuple(x: Any) -> Tuple[Any, Any]:
+    if isinstance(x, collections.abc.Iterable):
+        return x
+    return (x, x)
+
+
+class CedAudioPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        input_size: Conv_Kernel = 224,
+        patch_size: Conv_Kernel = 16,
+        patch_stride: Conv_Kernel = 16,
+        in_chans: int = 1,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten: bool = False,
+    ):
+        super().__init__()
+        self.input_size = to_2tuple(input_size)
+        self.patch_size = to_2tuple(patch_size)
+        self.patch_stride = to_2tuple(patch_stride)
+        self.grid_size = (
+            self.input_size[0] // self.patch_stride[0],
+            self.input_size[1] // self.patch_stride[1],
+        )
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_stride)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        x = self.proj(x)
+        if self.flatten:
+            x = torch.permute(torch.flatten(x, 2, 3), (0, 2, 1))
+        x = self.norm(x)
+        return x
+
+
+class CedAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        causal: bool = False,
+    ):
+        super().__init__()
+        assert dim % num_heads == 0, "dim should be divisible by num_heads"
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.causal = causal
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        # if mask is not None:
+        # # Mask is a tensor of shape [B, T, T]
+        # # Different from self.causal == True, the mask might be something like:
+        # # [False, False, True]
+        # # [False, False, True]
+        # # [True, True, True]
+        # # We use -inf to pad here, since if we would pad by any number, the entries at rows only containing
+        # # [True, True, True] would lead to weights such as: [0.33,0.33,0.33], which is not correct
+        # mask_value = torch.as_tensor(-float('inf'))
+        # print(mask.shape, attn.shape)
+        # attn = attn.masked_fill(mask, mask_value)
+        if self.causal:
+            mask_value = -torch.finfo(attn.dtype).max
+            i, j = attn.shape[-2:]
+            mask = torch.ones(i, j, device=q.device, dtype=torch.bool).triu(j - i + 1)
+            attn = attn.masked_fill(mask, mask_value)
+        attn = attn.softmax(dim=-1)
+        # Only for the case that a mask with all True entries on a row is passed.
+        # attn = torch.nan_to_num(attn)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class CedMlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable = nn.GELU,
+        drop: float = 0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+# Drop path is taken from Timm
+# https://github.com/huggingface/pytorch-image-models/blob/7c67d6aca992f039eece0af5f7c29a43d48c00e4/timm/models/layers/drop.py#L155
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+    def extra_repr(self):
+        return f"drop_prob={round(self.drop_prob,3):0.3f}"
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I (https://github.com/rwightman) created for EfficientNet, etc networks,
+    however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
+    layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
+    argument.
+
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class CedBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer: Callable = nn.GELU,
+        norm_layer: Callable = nn.LayerNorm,
+        attention_type: Callable = CedAttention,
+        attention_kwargs={},
+        **kwargs,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = attention_type(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+            **attention_kwargs,
+        )
+        self.ls1 = nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = CedMlp(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            drop=drop,
+        )
+        self.ls2 = nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x):
+        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))
+        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
+        return x
+
+
+# Taken from timm
+def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.0))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+CED_START_DOCSTRING = r"""
+
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`CedConfig`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+CED_INPUTS_DOCSTRING = r"""
+    Args:
+        input_values (`torch.FloatTensor` of shape `(batch_size, n_mels, sequence_length)`):
+            The sequence of audio features extracted from the audio signal. Can be obtained from a raw audio waveform
+            using `~transformers.CedFeatureExtractor.__call__`.
+"""
+
+
+@add_start_docstrings(
+    "The bare Ced Model transformer outputting raw hidden-states without any specific head on top.",
+    CED_START_DOCSTRING,
+)
+class CedModel(CedPreTrainedModel):
+    def __init__(self, config: CedConfig) -> None:
+        super().__init__(config)
+        self.config = config
+        self.name = config.name
+
+        # Allowed length in number of frames, otherwise the positional embedding will throw an error
+        self.maximal_allowed_length = self.config.target_length
+
+        self.init_bn = torch.nn.BatchNorm2d(config.n_mels, momentum=0.01)
+
+        self.patch_embed = CedAudioPatchEmbed(
+            input_size=(config.n_mels, config.target_length),
+            embed_dim=config.embed_dim,
+            patch_size=config.patch_size,
+            flatten=False,
+            patch_stride=config.patch_stride,
+        )
+
+        self.time_pos_embed = nn.Parameter(torch.randn(1, config.embed_dim, 1, self.patch_embed.grid_size[1]) * 0.02)
+        self.freq_pos_embed = nn.Parameter(torch.randn(1, config.embed_dim, self.patch_embed.grid_size[0], 1) * 0.02)
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+        act_layer = nn.GELU
+        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.depth)]  # stochastic depth decay rule
+        self.pos_drop = nn.Dropout(p=config.drop_rate)
+        self.blocks = nn.Sequential(
+            *[
+                CedBlock(
+                    dim=config.embed_dim,
+                    num_heads=config.num_heads,
+                    mlp_ratio=config.mlp_ratio,
+                    qkv_bias=config.qkv_bias,
+                    drop=config.drop_rate,
+                    attn_drop=config.attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    act_layer=act_layer,
+                    attention_type=CedAttention,
+                )
+                for i in range(config.depth)
+            ]
+        )
+        self.norm = norm_layer(config.embed_dim)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def _freeze_parameters(self):
+        for param in self.parameters():
+            param.requires_grad = False
+        self._requires_grad = False
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.patch_embed(x)
+        _, _, _, t = x.shape
+        x = x + self.time_pos_embed[:, :, :, :t]
+        x = x + self.freq_pos_embed[:, :, :, :]  # Just to support __getitem__ in posembed
+
+        # x = rearrange(x, 'b c f t -> b (f t) c')
+        x = torch.permute(torch.flatten(x, 2, 3), (0, 2, 1))
+
+        if self.config.pooling == "token":
+            cls_token = self.cls_token.expand(x.shape[0], -1, -1)
+            cls_token = cls_token + self.token_pos_embed
+            x = torch.cat((cls_token, x), dim=1)
+        x = self.pos_drop(x)
+        x = self.blocks(x)
+        x = self.norm(x)
+        return x
+
+    def forward(self, input_values: torch.Tensor):
+        r"""
+        Runs a forward pass of the CED model as an audio encoder.
+        """
+        x = torch.unsqueeze(input_values, 1)
+
+        x = torch.permute(x, (0, 2, 1, 3))
+        x = self.init_bn(x)
+        x = torch.permute(x, (0, 2, 1, 3))
+
+        if x.shape[-1] > self.maximal_allowed_length:
+            splits = x.split(self.maximal_allowed_length, -1)
+
+            if splits[-1].shape[-1] < self.maximal_allowed_length:
+                if self.config.pad_last:
+                    pad = torch.zeros(*x.shape[:-1], self.maximal_allowed_length, device=x.device)
+                    pad[..., : splits[-1].shape[-1]] = splits[-1]
+                    splits = torch.stack((*splits[:-1], pad), dim=0)
+                else:
+                    splits = torch.stack(splits[:-1], dim=0)
+            else:
+                splits = torch.stack(splits[:-1], dim=0)
+            n_splits = len(splits)
+            x = torch.flatten(splits, 0, 1)  # spl b c f t-> (spl b) c f t
+        else:
+            n_splits = 1
+
+        x = self.forward_features(x)
+        x = torch.reshape(x, (x.shape[0] // n_splits, -1, x.shape[-1]))
+
+        return SequenceClassifierOutput(logits=x)
+
+
+@add_start_docstrings(
+    """
+    Ced model with an audio classification head on top (a linear layer on top of the pooled output).
+    """,
+    CED_START_DOCSTRING,
+)
+class CedForAudioClassification(CedPreTrainedModel):
+    def __init__(self, config: CedConfig) -> None:
+        super().__init__(config)
+        self.config = config
+
+        self.encoder = CedModel(config)
+
+        # Classifier head
+        self.outputlayer = nn.Sequential(
+            nn.LayerNorm(config.embed_dim),
+            nn.Linear(config.embed_dim, config.outputdim),
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward_head(self, x: torch.Tensor) -> torch.Tensor:
+        if self.config.pooling == "token":
+            x = x[:, 0]
+            return self.outputlayer(x).sigmoid()
+        elif self.config.pooling == "mean":
+            x = x.mean(1)
+            return self.outputlayer(x).sigmoid()
+        elif self.config.pooling == "logit":
+            x = x.mean(1)
+            return self.outputlayer(x)
+        elif self.config.pooling == "dm":
+            # Unpack using the frequency dimension, which is constant
+            # 'b (f t) d -> b f t d', f=self.patch_embed.grid_size[0])
+            x = torch.reshape(x, (x.shape[0], self.patch_embed.grid_size[0], -1, x.shape[3]))
+
+            # First poolin frequency, then sigmoid the (B T D) output
+            x = self.outputlayer(x.mean(1)).sigmoid()
+            return x.mean(1)
+        else:
+            return x.mean(1)
+
+    def freeze_encoder(self):
+        self.encoder._freeze_parameters()
+
+    @add_start_docstrings_to_model_forward(CED_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_code_sample_docstrings(
+        checkpoint=_SEQ_CLASS_CHECKPOINT,
+        output_type=SequenceClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+        modality="audio",
+        model_cls="CedForAudioClassification",
+        expected_output=_SEQ_CLASS_EXPECTED_OUTPUT,
+        expected_loss=_SEQ_CLASS_EXPECTED_LOSS,
+    )
+    def forward(self, input_values: torch.Tensor, labels: Optional[torch.Tensor] = None):
+        """
+        Runs a forward pass of the CED model for audio classification task.
+
+        Examples:
+
+        ```python
+        >>> from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
+        >>> from datasets import load_dataset
+        >>> import torch
+
+        >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
+        >>> dataset = dataset.sort("id")
+        >>> sampling_rate = dataset.features["audio"].sampling_rate
+
+        >>> feature_extractor = AutoFeatureExtractor.from_pretrained("mispeech/ced-tiny")
+        >>> model = AutoModelForAudioClassification.from_pretrained("mispeech/ced-tiny")
+
+        >>> inputs = feature_extractor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt")
+
+        >>> with torch.no_grad():
+        ...     logits = model(**inputs).logits
+
+        >>> predicted_class_ids = torch.argmax(logits, dim=-1).item()
+        >>> predicted_label = model.config.id2label[predicted_class_ids]
+        >>> predicted_label
+        'Speech synthesizer'
+        ```
+        """
+        last_hidden_states = self.encoder(input_values).logits
+        logits = self.forward_head(last_hidden_states)
+
+        if labels is not None:
+            try:
+                loss_fct = getattr(nn.modules.loss, self.config.loss)()
+            except AttributeError:
+                raise NotImplementedError(f"Loss {self.config.loss} not implemented.")
+
+            labels = nn.functional.one_hot(labels, num_classes=self.config.outputdim).float()
+            loss = loss_fct(logits, labels)
+        else:
+            loss = None
+
+        return SequenceClassifierOutput(logits=logits, loss=loss, hidden_states=last_hidden_states)