Spaces:
Build error
Build error
| # Copyright 2024 The YourMT3 Authors. | |
| # | |
| # 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 | |
| # | |
| # Please see the details in the LICENSE file. | |
| import math | |
| from typing import Optional, Union | |
| from torch import nn | |
| from transformers.configuration_utils import PretrainedConfig | |
| from transformers.modeling_utils import PreTrainedModel | |
| class ConformerYMT3Config(PretrainedConfig): | |
| r""" | |
| This is the configuration class to store the configuration of a [`ConformerYMT3Encoder`]. It is used to | |
| instantiate an ConformerYMT3Encoder according to the specified arguments, defining the model architecture. | |
| Instantiating a configuration with the defaults will yield a similar configuration to that of the Wav2Vec2Conformer | |
| [facebook/wav2vec2-conformer-rel-pos-large](https://huggingface.co/facebook/wav2vec2-conformer-rel-pos-large) | |
| architecture. | |
| Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the | |
| documentation from [`PretrainedConfig`] for more information. | |
| Args: | |
| d_model (`int`, *optional*, defaults to 512): | |
| Dimensionality of the encoder layers and the pooler layer. | |
| num_layers (`int`, *optional*, defaults to 12): | |
| Number of hidden layers in the Transformer encoder. | |
| num_heads (`int`, *optional*, defaults to 12): | |
| Number of attention heads for each attention layer in the Transformer encoder. | |
| intermediate_size (`int`, *optional*, defaults to 2048): | |
| Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. | |
| hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`): | |
| The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`, | |
| `"relu"`, `"selu"` and `"gelu_new"` are supported. | |
| dropout_rate (`float`, *optional*, defaults to 0.05): | |
| The dropout probability for all fully connected layers in the embeddings, encoder, and pooler. | |
| layerdrop (`float`, *optional*, defaults to 0.1): | |
| The LayerDrop probability. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556) for more | |
| details. | |
| initializer_range (`float`, *optional*, defaults to 0.02): | |
| The standard deviation of the truncated_normal_initializer for initializing all weight matrices. | |
| layer_norm_eps (`float`, *optional*, defaults to 1e-12): | |
| The epsilon used by the layer normalization layers. | |
| conv_dim (`Tuple[int]` or `List[int]`, *optional*, defaults to `(512, 512, 512, 512, 512, 512, 512)`): | |
| A tuple of integers defining the number of input and output channels of each 1D convolutional layer in the | |
| feature encoder. The length of *conv_dim* defines the number of 1D convolutional layers. | |
| conv_stride (`Tuple[int]` or `List[int]`, *optional*, defaults to `(5, 2, 2, 2, 2, 2, 2)`): | |
| A tuple of integers defining the stride of each 1D convolutional layer in the feature encoder. The length | |
| of *conv_stride* defines the number of convolutional layers and has to match the length of *conv_dim*. | |
| conv_kernel (`Tuple[int]` or `List[int]`, *optional*, defaults to `(10, 3, 3, 3, 3, 3, 3)`): | |
| A tuple of integers defining the kernel size of each 1D convolutional layer in the feature encoder. The | |
| length of *conv_kernel* defines the number of convolutional layers and has to match the length of | |
| *conv_dim*. | |
| conv_bias (`bool`, *optional*, defaults to `False`): | |
| Whether the 1D convolutional layers have a bias. | |
| output_hidden_size (`int`, *optional*): | |
| Dimensionality of the encoder output layer. If not defined, this defaults to *hidden-size*. Only relevant | |
| if `add_adapter is True`. | |
| position_encoding_type (`str`, *optional*, defaults to `"relative"`): | |
| Can be specified to `relative` or `rotary` for relative or rotary position embeddings respectively. If left | |
| `None` no relative position embedding is applied. | |
| rotary_embedding_base (`int`, *optional*, defaults to 10000): | |
| If `"rotary"` position embeddings are used, defines the size of the embedding base. | |
| num_max_positions (`int`, *optional*, defaults to 5000): | |
| if `"relative"` position embeddings are used, defines the maximum source input positions. | |
| conv_depthwise_kernel_size (`int`, defaults to 31): | |
| Kernel size of convolutional depthwise 1D layer in Conformer blocks. | |
| Example: | |
| ```python | |
| >>> from transformers import ConformerYMT3Config, ConformerYMT3Encoder | |
| >>> # Initializing a ConformerYMT3Encoder configuration | |
| >>> configuration = ConformerYMT3Config() | |
| >>> # Initializing a model (with random weights) from the facebook/wav2vec2-conformer-rel-pos-large style configuration | |
| >>> model = ConformerYMT3Encoder(configuration) | |
| >>> # Accessing the model configuration | |
| >>> configuration = model.config | |
| ```""" | |
| model_type = "conformer-ymt3" | |
| def __init__( | |
| self, | |
| d_model=512, # 768 | |
| num_layers=8, # ConformerYMT3Encoder | |
| num_heads=8, # ConformerYMT3SelfAttention | |
| intermediate_size=2048, # 3072,# used in intermediate_dense of ConformerYMT3FeedForward | |
| hidden_act="gelu", # used in intermediate_act_fn of ConformerYMT3FeedForward | |
| dropout_rate=0.1, | |
| layerdrop=0.1, | |
| initializer_range=0.02, | |
| layer_norm_eps=1e-5, | |
| conv_dim=(512, 512, 512, 512, 512, 512, 512), | |
| conv_stride=(5, 2, 2, 2, 2, 2, 2), | |
| conv_kernel=(10, 3, 3, 3, 3, 3, 3), | |
| conv_bias=False, | |
| position_encoding_type="rotary", | |
| rotary_embedding_base=10000, | |
| num_max_positions=1024, | |
| conv_depthwise_kernel_size=31, | |
| **kwargs, | |
| ): | |
| super().__init__(**kwargs) | |
| self.d_model = d_model | |
| self.conv_dim = list(conv_dim) | |
| self.conv_stride = list(conv_stride) | |
| self.conv_kernel = list(conv_kernel) | |
| self.conv_bias = conv_bias | |
| self.num_layers = num_layers | |
| self.intermediate_size = intermediate_size | |
| self.hidden_act = hidden_act | |
| self.num_heads = num_heads | |
| self.dropout_rate = dropout_rate | |
| self.layerdrop = layerdrop | |
| self.layer_norm_eps = layer_norm_eps | |
| self.initializer_range = initializer_range | |
| self.num_max_positions = num_max_positions | |
| self.position_encoding_type = position_encoding_type | |
| self.rotary_embedding_base = rotary_embedding_base | |
| # Conformer-block related | |
| self.conv_depthwise_kernel_size = conv_depthwise_kernel_size | |
| class ConformerYMT3PreTrainedModel(PreTrainedModel): | |
| """ | |
| An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained | |
| models. | |
| """ | |
| config_class = ConformerYMT3Config | |
| base_model_prefix = "wav2vec2_conformer" | |
| main_input_name = "input_values" | |
| supports_gradient_checkpointing = True | |
| def _init_weights(self, module): | |
| """Initialize the weights""" | |
| if module.__class__.__name__ == "ConformerYMT3SelfAttention": | |
| if hasattr(module, "pos_bias_u"): | |
| nn.init.xavier_uniform_(module.pos_bias_u) | |
| if hasattr(module, "pos_bias_v"): | |
| nn.init.xavier_uniform_(module.pos_bias_v) | |
| elif isinstance(module, nn.Linear): | |
| module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) | |
| if module.bias is not None: | |
| module.bias.data.zero_() | |
| elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): | |
| module.bias.data.zero_() | |
| module.weight.data.fill_(1.0) | |
| elif isinstance(module, nn.Conv1d): | |
| nn.init.kaiming_normal_(module.weight) | |
| if module.bias is not None: | |
| k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0])) | |
| nn.init.uniform_(module.bias, a=-k, b=k) | |
| def _set_gradient_checkpointing(self, module, value=False): | |
| if module.__class__.__name__ == "ConformerYMT3Encoder": | |
| module.gradient_checkpointing = value | |