| from transformers import PretrainedConfig |
| from torch import nn |
| from .types import TransformerLayerCFG, TransformerEncoderCFG |
|
|
|
|
|
|
| class WavJEPAConfig(PretrainedConfig): |
| model_type = "wavjepa-base" |
| model_size = "base" |
| in_channels: int = 1 |
|
|
| def __init__( |
| self, |
| extractor_layers_spec: str = "[(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512,2,2)]", |
| extractor_dropout : float = 0.0, |
| extractor_mode : str = "default", |
| extractor_conv_bias : bool = False, |
| extractor_depthwise: bool = False, |
| encoder_d_model: int = 768, |
| encoder_nhead : int = 12, |
| encoder_batch_first = True, |
| encoder_norm_first = False, |
| encoder_bias = True, |
| encoder_mlp_ratio = 4.0, |
| encoder_dropout = 0.0, |
| encoder_num_layers: int = 12, |
| encoder_enable_nested_tensor = False, |
| encoder_mask_check = True, |
| decoder_d_model: int = 384, |
| decoder_nhead : int = 12, |
| decoder_batch_first = True, |
| decoder_norm_first = False, |
| decoder_bias = True, |
| decoder_mlp_ratio = 4.0, |
| decoder_dropout = 0.0, |
| decoder_num_layers: int = 12, |
| decoder_enable_nested_tensor = False, |
| decoder_mask_check = True, |
| **kwargs |
| ): |
| self.encoder_cfg = TransformerEncoderCFG.create( |
| num_layers = encoder_num_layers, |
| enable_nested_tensor = encoder_enable_nested_tensor, |
| mask_check = encoder_mask_check, |
| ) |
| self.decoder_cfg = TransformerEncoderCFG.create( |
| num_layers = decoder_num_layers, |
| enable_nested_tensor = decoder_enable_nested_tensor, |
| mask_check = decoder_mask_check, |
| ) |
| |
| self.encoder_layers_cfg = TransformerLayerCFG.create( |
| d_model = encoder_d_model, |
| nhead = encoder_nhead, |
| batch_first = encoder_batch_first, |
| norm_first = encoder_norm_first, |
| bias = encoder_bias, |
| mlp_ratio = encoder_mlp_ratio, |
| dropout = encoder_dropout, |
| layer_norm_eps = 1e-6 |
| ) |
| self.decoder_layers_cfg = TransformerLayerCFG.create( |
| d_model = decoder_d_model, |
| nhead = decoder_nhead, |
| batch_first = decoder_batch_first, |
| norm_first = decoder_norm_first, |
| bias = decoder_bias, |
| mlp_ratio = decoder_mlp_ratio, |
| dropout = decoder_dropout, |
| layer_norm_eps = 1e-6 |
| ) |
| self.extractor_config = dict( |
| conv_layers_spec = extractor_layers_spec, |
| in_channels = self.in_channels, |
| dropout = extractor_dropout, |
| mode = extractor_mode, |
| conv_bias = extractor_conv_bias, |
| depthwise = extractor_depthwise) |
|
|
| super().__init__(**kwargs) |
