mispeech/ced-mini

1# coding=utf-82# Copyright 2023 Xiaomi Corporation 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""" PyTorch CED (Ced) model."""16 17import collections18import math19from functools import partial20from typing import Any, Callable, Optional, Tuple, Union21 22import torch23import torch.utils.checkpoint24from torch import nn25 26from transformers.modeling_outputs import SequenceClassifierOutput27from transformers.modeling_utils import PreTrainedModel28from transformers.utils import (29    add_code_sample_docstrings,30    add_start_docstrings,31    add_start_docstrings_to_model_forward,32    logging,33)34from .configuration_ced import CedConfig35 36 37logger = logging.get_logger(__name__)38 39_CONFIG_FOR_DOC = "CedConfig"40_SEQ_CLASS_EXPECTED_OUTPUT = "'Speech synthesizer'"41_SEQ_CLASS_EXPECTED_LOSS = 0.6942 43# Audio classification docstring44_SEQ_CLASS_CHECKPOINT = "mispeech/ced-tiny"45 46 47CED_PRETRAINED_MODEL_ARCHIVE_LIST = [48    "mispeech/ced-tiny",49    "mispeech/ced-mini",50    "mispeech/ced-small",51    "mispeech/ced-base",52    # See all CED models at https://huggingface.co/models?search=mispeech%2Fced53]54 55 56class CedPreTrainedModel(PreTrainedModel):57    """58    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained59    models.60    """61 62    config_class = CedConfig63    base_model_prefix = "ced"64    main_input_name = "input_values"65    supports_gradient_checkpointing = True66 67    def _init_weights(self, module):68        """Initialize the weights"""69        if isinstance(module, nn.Linear):70            trunc_normal_(module.weight, std=0.02)71            if module.bias is not None:72                nn.init.zeros_(module.bias)73        elif isinstance(module, nn.LayerNorm):74            nn.init.constant_(module.bias, 0)75            nn.init.constant_(module.weight, 1.0)76 77 78Conv_Kernel = Union[int, Tuple[int, int]]79 80 81def to_2tuple(x: Any) -> Tuple[Any, Any]:82    if isinstance(x, collections.abc.Iterable):83        return x84    return (x, x)85 86 87class CedAudioPatchEmbed(nn.Module):88    def __init__(89        self,90        input_size: Conv_Kernel = 224,91        patch_size: Conv_Kernel = 16,92        patch_stride: Conv_Kernel = 16,93        in_chans: int = 1,94        embed_dim: int = 768,95        norm_layer: Optional[Callable] = None,96        flatten: bool = False,97    ):98        super().__init__()99        self.input_size = to_2tuple(input_size)100        self.patch_size = to_2tuple(patch_size)101        self.patch_stride = to_2tuple(patch_stride)102        self.grid_size = (103            self.input_size[0] // self.patch_stride[0],104            self.input_size[1] // self.patch_stride[1],105        )106        self.num_patches = self.grid_size[0] * self.grid_size[1]107        self.flatten = flatten108 109        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_stride)110        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()111 112    def forward(self, x):113        x = self.proj(x)114        if self.flatten:115            x = torch.permute(torch.flatten(x, 2, 3), (0, 2, 1))116        x = self.norm(x)117        return x118 119 120class CedAttention(nn.Module):121    def __init__(122        self,123        dim,124        num_heads=8,125        qkv_bias=False,126        attn_drop=0.0,127        proj_drop=0.0,128        causal: bool = False,129    ):130        super().__init__()131        assert dim % num_heads == 0, "dim should be divisible by num_heads"132        self.num_heads = num_heads133        head_dim = dim // num_heads134        self.scale = head_dim**-0.5135 136        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)137        self.attn_drop = nn.Dropout(attn_drop)138        self.proj = nn.Linear(dim, dim)139        self.proj_drop = nn.Dropout(proj_drop)140        self.causal = causal141 142    def forward(self, x):143        B, N, C = x.shape144        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)145        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)146 147        attn = (q @ k.transpose(-2, -1)) * self.scale148        # if mask is not None:149        # # Mask is a tensor of shape [B, T, T]150        # # Different from self.causal == True, the mask might be something like:151        # # [False, False, True]152        # # [False, False, True]153        # # [True, True, True]154        # # We use -inf to pad here, since if we would pad by any number, the entries at rows only containing155        # # [True, True, True] would lead to weights such as: [0.33,0.33,0.33], which is not correct156        # mask_value = torch.as_tensor(-float('inf'))157        # print(mask.shape, attn.shape)158        # attn = attn.masked_fill(mask, mask_value)159        if self.causal:160            mask_value = -torch.finfo(attn.dtype).max161            i, j = attn.shape[-2:]162            mask = torch.ones(i, j, device=q.device, dtype=torch.bool).triu(j - i + 1)163            attn = attn.masked_fill(mask, mask_value)164        attn = attn.softmax(dim=-1)165        # Only for the case that a mask with all True entries on a row is passed.166        # attn = torch.nan_to_num(attn)167        attn = self.attn_drop(attn)168 169        x = (attn @ v).transpose(1, 2).reshape(B, N, C)170        x = self.proj(x)171        x = self.proj_drop(x)172        return x173 174 175class CedMlp(nn.Module):176    def __init__(177        self,178        in_features: int,179        hidden_features: Optional[int] = None,180        out_features: Optional[int] = None,181        act_layer: Callable = nn.GELU,182        drop: float = 0.0,183    ):184        super().__init__()185        out_features = out_features or in_features186        hidden_features = hidden_features or in_features187        self.fc1 = nn.Linear(in_features, hidden_features)188        self.act = act_layer()189        self.fc2 = nn.Linear(hidden_features, out_features)190        self.drop = nn.Dropout(drop)191 192    def forward(self, x):193        x = self.fc1(x)194        x = self.act(x)195        x = self.drop(x)196        x = self.fc2(x)197        x = self.drop(x)198        return x199 200 201# Drop path is taken from Timm202# https://github.com/huggingface/pytorch-image-models/blob/7c67d6aca992f039eece0af5f7c29a43d48c00e4/timm/models/layers/drop.py#L155203class DropPath(nn.Module):204    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""205 206    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):207        super(DropPath, self).__init__()208        self.drop_prob = drop_prob209        self.scale_by_keep = scale_by_keep210 211    def forward(self, x):212        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)213 214    def extra_repr(self):215        return f"drop_prob={round(self.drop_prob,3):0.3f}"216 217 218def drop_path(x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True):219    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).220 221    This is the same as the DropConnect impl I (https://github.com/rwightman) created for EfficientNet, etc networks,222    however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...223    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the224    layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the225    argument.226 227    """228    if drop_prob == 0.0 or not training:229        return x230    keep_prob = 1 - drop_prob231    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets232    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)233    if keep_prob > 0.0 and scale_by_keep:234        random_tensor.div_(keep_prob)235    return x * random_tensor236 237 238class CedBlock(nn.Module):239    def __init__(240        self,241        dim,242        num_heads,243        mlp_ratio=4.0,244        qkv_bias=False,245        drop=0.0,246        attn_drop=0.0,247        drop_path=0.0,248        act_layer: Callable = nn.GELU,249        norm_layer: Callable = nn.LayerNorm,250        attention_type: Callable = CedAttention,251        attention_kwargs={},252        **kwargs,253    ):254        super().__init__()255        self.norm1 = norm_layer(dim)256        self.attn = attention_type(257            dim,258            num_heads=num_heads,259            qkv_bias=qkv_bias,260            attn_drop=attn_drop,261            proj_drop=drop,262            **attention_kwargs,263        )264        self.ls1 = nn.Identity()265        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()266 267        self.norm2 = norm_layer(dim)268        self.mlp = CedMlp(269            in_features=dim,270            hidden_features=int(dim * mlp_ratio),271            act_layer=act_layer,272            drop=drop,273        )274        self.ls2 = nn.Identity()275        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()276 277    def forward(self, x):278        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))279        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))280        return x281 282 283# Taken from timm284def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):285    return _no_grad_trunc_normal_(tensor, mean, std, a, b)286 287 288def _no_grad_trunc_normal_(tensor, mean, std, a, b):289    # Cut & paste from PyTorch official master until it's in a few official releases - RW290    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf291    def norm_cdf(x):292        # Computes standard normal cumulative distribution function293        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0294 295    with torch.no_grad():296        # Values are generated by using a truncated uniform distribution and297        # then using the inverse CDF for the normal distribution.298        # Get upper and lower cdf values299        l = norm_cdf((a - mean) / std)300        u = norm_cdf((b - mean) / std)301 302        # Uniformly fill tensor with values from [l, u], then translate to303        # [2l-1, 2u-1].304        tensor.uniform_(2 * l - 1, 2 * u - 1)305 306        # Use inverse cdf transform for normal distribution to get truncated307        # standard normal308        tensor.erfinv_()309 310        # Transform to proper mean, std311        tensor.mul_(std * math.sqrt(2.0))312        tensor.add_(mean)313 314        # Clamp to ensure it's in the proper range315        tensor.clamp_(min=a, max=b)316        return tensor317 318 319CED_START_DOCSTRING = r"""320 321    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the322    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads323    etc.)324 325    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.326    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage327    and behavior.328 329    Parameters:330        config ([`CedConfig`]): Model configuration class with all the parameters of the model.331            Initializing with a config file does not load the weights associated with the model, only the332            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.333"""334 335CED_INPUTS_DOCSTRING = r"""336    Args:337        input_values (`torch.FloatTensor` of shape `(batch_size, n_mels, sequence_length)`):338            The sequence of audio features extracted from the audio signal. Can be obtained from a raw audio waveform339            using `~transformers.CedFeatureExtractor.__call__`.340"""341 342 343@add_start_docstrings(344    "The bare Ced Model transformer outputting raw hidden-states without any specific head on top.",345    CED_START_DOCSTRING,346)347class CedModel(CedPreTrainedModel):348    def __init__(self, config: CedConfig) -> None:349        super().__init__(config)350        self.config = config351        self.name = config.name352 353        # Allowed length in number of frames, otherwise the positional embedding will throw an error354        self.maximal_allowed_length = self.config.target_length355 356        self.init_bn = torch.nn.BatchNorm2d(config.n_mels, momentum=0.01)357 358        self.patch_embed = CedAudioPatchEmbed(359            input_size=(config.n_mels, config.target_length),360            embed_dim=config.embed_dim,361            patch_size=config.patch_size,362            flatten=False,363            patch_stride=config.patch_stride,364        )365 366        self.time_pos_embed = nn.Parameter(torch.randn(1, config.embed_dim, 1, self.patch_embed.grid_size[1]) * 0.02)367        self.freq_pos_embed = nn.Parameter(torch.randn(1, config.embed_dim, self.patch_embed.grid_size[0], 1) * 0.02)368        norm_layer = partial(nn.LayerNorm, eps=1e-6)369        act_layer = nn.GELU370        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.depth, device="cpu")]  # stochastic depth decay rule371        self.pos_drop = nn.Dropout(p=config.drop_rate)372        self.blocks = nn.Sequential(373            *[374                CedBlock(375                    dim=config.embed_dim,376                    num_heads=config.num_heads,377                    mlp_ratio=config.mlp_ratio,378                    qkv_bias=config.qkv_bias,379                    drop=config.drop_rate,380                    attn_drop=config.attn_drop_rate,381                    drop_path=dpr[i],382                    norm_layer=norm_layer,383                    act_layer=act_layer,384                    attention_type=CedAttention,385                )386                for i in range(config.depth)387            ]388        )389        self.norm = norm_layer(config.embed_dim)390 391        # Initialize weights and apply final processing392        self.post_init()393 394    def _freeze_parameters(self):395        for param in self.parameters():396            param.requires_grad = False397        self._requires_grad = False398 399    def forward_features(self, x: torch.Tensor) -> torch.Tensor:400        x = self.patch_embed(x)401        _, _, _, t = x.shape402        x = x + self.time_pos_embed[:, :, :, :t]403        x = x + self.freq_pos_embed[:, :, :, :]  # Just to support __getitem__ in posembed404 405        # x = rearrange(x, 'b c f t -> b (f t) c')406        x = torch.permute(torch.flatten(x, 2, 3), (0, 2, 1))407 408        if self.config.pooling == "token":409            cls_token = self.cls_token.expand(x.shape[0], -1, -1)410            cls_token = cls_token + self.token_pos_embed411            x = torch.cat((cls_token, x), dim=1)412        x = self.pos_drop(x)413        x = self.blocks(x)414        x = self.norm(x)415        return x416 417    def forward(self, input_values: torch.Tensor):418        r"""419        Runs a forward pass of the CED model as an audio encoder.420        """421        x = torch.unsqueeze(input_values, 1)422 423        x = torch.permute(x, (0, 2, 1, 3))424        x = self.init_bn(x)425        x = torch.permute(x, (0, 2, 1, 3))426 427        if x.shape[-1] > self.maximal_allowed_length:428            splits = x.split(self.maximal_allowed_length, -1)429 430            if splits[-1].shape[-1] < self.maximal_allowed_length:431                if self.config.pad_last:432                    pad = torch.zeros(*x.shape[:-1], self.maximal_allowed_length, device=x.device)433                    pad[..., : splits[-1].shape[-1]] = splits[-1]434                    splits = torch.stack((*splits[:-1], pad), dim=0)435                else:436                    splits = torch.stack(splits[:-1], dim=0)437            else:438                splits = torch.stack(splits[:-1], dim=0)439            n_splits = len(splits)440            x = torch.flatten(splits, 0, 1)  # spl b c f t-> (spl b) c f t441        else:442            n_splits = 1443 444        x = self.forward_features(x)445        x = torch.reshape(x, (x.shape[0] // n_splits, -1, x.shape[-1]))446 447        return SequenceClassifierOutput(logits=x)448 449 450@add_start_docstrings(451    """452    Ced model with an audio classification head on top (a linear layer on top of the pooled output).453    """,454    CED_START_DOCSTRING,455)456class CedForAudioClassification(CedPreTrainedModel):457    def __init__(self, config: CedConfig) -> None:458        super().__init__(config)459        self.config = config460 461        self.encoder = CedModel(config)462 463        # Classifier head464        self.outputlayer = nn.Sequential(465            nn.LayerNorm(config.embed_dim),466            nn.Linear(config.embed_dim, config.outputdim),467        )468 469        # Initialize weights and apply final processing470        self.post_init()471 472    def forward_head(self, x: torch.Tensor) -> torch.Tensor:473        if self.config.pooling == "token":474            x = x[:, 0]475            return self.outputlayer(x).sigmoid()476        elif self.config.pooling == "mean":477            x = x.mean(1)478            return self.outputlayer(x).sigmoid()479        elif self.config.pooling == "logit":480            x = x.mean(1)481            return self.outputlayer(x)482        elif self.config.pooling == "dm":483            # Unpack using the frequency dimension, which is constant484            # 'b (f t) d -> b f t d', f=self.patch_embed.grid_size[0])485            x = torch.reshape(x, (x.shape[0], self.patch_embed.grid_size[0], -1, x.shape[3]))486 487            # First poolin frequency, then sigmoid the (B T D) output488            x = self.outputlayer(x.mean(1)).sigmoid()489            return x.mean(1)490        else:491            return x.mean(1)492 493    def freeze_encoder(self):494        self.encoder._freeze_parameters()495 496    @add_start_docstrings_to_model_forward(CED_INPUTS_DOCSTRING.format("batch_size, sequence_length"))497    @add_code_sample_docstrings(498        checkpoint=_SEQ_CLASS_CHECKPOINT,499        output_type=SequenceClassifierOutput,500        config_class=_CONFIG_FOR_DOC,501        modality="audio",502        model_cls="CedForAudioClassification",503        expected_output=_SEQ_CLASS_EXPECTED_OUTPUT,504        expected_loss=_SEQ_CLASS_EXPECTED_LOSS,505    )506    def forward(self, input_values: torch.Tensor, labels: Optional[torch.Tensor] = None):507        """508        Runs a forward pass of the CED model for audio classification task.509 510        Examples:511 512        ```python513        >>> from transformers import AutoFeatureExtractor, AutoModelForAudioClassification514        >>> from datasets import load_dataset515        >>> import torch516 517        >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")518        >>> dataset = dataset.sort("id")519        >>> sampling_rate = dataset.features["audio"].sampling_rate520 521        >>> feature_extractor = AutoFeatureExtractor.from_pretrained("mispeech/ced-tiny")522        >>> model = AutoModelForAudioClassification.from_pretrained("mispeech/ced-tiny")523 524        >>> inputs = feature_extractor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt")525 526        >>> with torch.no_grad():527        ...     logits = model(**inputs).logits528 529        >>> predicted_class_ids = torch.argmax(logits, dim=-1).item()530        >>> predicted_label = model.config.id2label[predicted_class_ids]531        >>> predicted_label532        'Speech synthesizer'533        ```534        """535        last_hidden_states = self.encoder(input_values).logits536        logits = self.forward_head(last_hidden_states)537 538        if labels is not None:539            try:540                loss_fct = getattr(nn.modules.loss, self.config.loss)()541            except AttributeError:542                raise NotImplementedError(f"Loss {self.config.loss} not implemented.")543 544            labels = nn.functional.one_hot(labels, num_classes=self.config.outputdim).float()545            loss = loss_fct(logits, labels)546        else:547            loss = None548 549        return SequenceClassifierOutput(logits=logits, loss=loss, hidden_states=last_hidden_states)550