






from ...configuration_utils import PretrainedConfig
from ...utils import logging


__A      =				logging.get_logger(__name__)

__A      =				{
    'facebook/s2t-small-librispeech-asr': (
        'https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/config.json'
    ),
    # See all Speech2Text models at https://huggingface.co/models?filter=speech_to_text
}





class 						SCREAMING_SNAKE_CASE				( __lowerCAmelCase					):





							"""simple docstring"""

							A_						  =					'''speech_to_text'''
							A_						  =					['''past_key_values''']
							A_						  =					{'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}

							def __init__(  self:   int		,     __A:   Optional[Any]=1_00_00		,     __A:   Dict=12		,     __A:   Any=20_48		,     __A:   Tuple=4		,     __A:   List[str]=6		,     __A:   str=20_48		,     __A:   List[Any]=4		,     __A:   Optional[int]=0.0		,     __A:   List[Any]=0.0		,     __A:   Union[str, Any]=True		,     __A:   Tuple=True		,     __A:   int="relu"		,     __A:   Optional[Any]=2_56		,     __A:   Union[str, Any]=0.1		,     __A:   int=0.0		,     __A:   Union[str, Any]=0.0		,     __A:   Dict=0.02		,     __A:   Dict=2		,     __A:   int=True		,     __A:   List[Any]=1		,     __A:   List[Any]=0		,     __A:   Union[str, Any]=2		,     __A:   Optional[Any]=60_00		,     __A:   List[Any]=10_24		,     __A:   Tuple=2		,     __A:   List[str]=(5, 5)		,     __A:   int=10_24		,     __A:   Optional[int]=80		,     __A:   Dict=1		,     **__A:   Union[str, Any]		,     )    ->				Tuple:
												_A								= vocab_size
												_A								= d_model
												_A								= encoder_ffn_dim
												_A								= encoder_layers
												_A								= encoder_attention_heads
												_A								= decoder_ffn_dim
												_A								= decoder_layers
												_A								= decoder_attention_heads
												_A								= dropout
												_A								= attention_dropout
												_A								= activation_dropout
												_A								= activation_function
												_A								= init_std
												_A								= encoder_layerdrop
												_A								= decoder_layerdrop
												_A								= use_cache
												_A								= encoder_layers
												_A								= scale_embedding  # scale factor will be sqrt(d_model) if True
												_A								= max_source_positions
												_A								= max_target_positions
												_A								= num_conv_layers
												_A								= list(lowerCAmelCase_   )
												_A								= conv_channels
												_A								= input_feat_per_channel
												_A								= input_channels

												if len(self.conv_kernel_sizes   ) != self.num_conv_layers:
																	raise ValueError(
																	    '''Configuration for convolutional module is incorrect. '''
																	    '''It is required that `len(config.conv_kernel_sizes)` == `config.num_conv_layers` '''
																	    f"""but is `len(config.conv_kernel_sizes) = {len(self.conv_kernel_sizes   )}`, """
																	    f"""`config.num_conv_layers = {self.num_conv_layers}`."""   )

												super().__init__(
												    pad_token_id=lowerCAmelCase_		,     bos_token_id=lowerCAmelCase_		,     eos_token_id=lowerCAmelCase_		,     is_encoder_decoder=lowerCAmelCase_		,     decoder_start_token_id=lowerCAmelCase_		,     **lowerCAmelCase_		,     )














import flax.linen as nn
import jax
import jax.numpy as jnp





class 						SCREAMING_SNAKE_CASE				( nn.Module					):





							"""simple docstring"""

							A_						  =					42
							A_						  =					jnp.floataa

							def       __A (  self:   Tuple   )    ->				Tuple:
												_A								= nn.Conv(
												    self.out_channels		,     kernel_size=(3, 3)		,     strides=(1, 1)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )



							def __call__(  self:   Dict		,     __A:   Dict   )    ->				Tuple:
												_A							,_A							,_A							,_A								= hidden_states.shape
												_A								= jax.image.resize(
												    __A		,     shape=(batch, height * 2, width * 2, channels)		,     method='''nearest'''		,     )
												_A								= self.conv(__A   )
												return hidden_states





class 						SCREAMING_SNAKE_CASE				( nn.Module					):





							"""simple docstring"""

							A_						  =					42
							A_						  =					jnp.floataa

							def       __A (  self:   List[str]   )    ->				Tuple:
												_A								= nn.Conv(
												    self.out_channels		,     kernel_size=(3, 3)		,     strides=(2, 2)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )



							def __call__(  self:   Union[str, Any]		,     __A:   List[Any]   )    ->				Union[str, Any]:
												# pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
												# hidden_states = jnp.pad(hidden_states, pad_width=pad)
												_A								= self.conv(__A   )
												return hidden_states





class 						SCREAMING_SNAKE_CASE				( nn.Module					):





							"""simple docstring"""

							A_						  =					42
							A_						  =					None
							A_						  =					0.0
							A_						  =					None
							A_						  =					jnp.floataa

							def       __A (  self:   Dict   )    ->				Dict:
												_A								= self.in_channels if self.out_channels is None else self.out_channels

												_A								= nn.GroupNorm(num_groups=32		,     epsilon=1e-5   )
												_A								= nn.Conv(
												    __A		,     kernel_size=(3, 3)		,     strides=(1, 1)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )

												_A								= nn.Dense(__A		,     dtype=self.dtype   )

												_A								= nn.GroupNorm(num_groups=32		,     epsilon=1e-5   )
												_A								= nn.Dropout(self.dropout_prob   )
												_A								= nn.Conv(
												    __A		,     kernel_size=(3, 3)		,     strides=(1, 1)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )

												_A								= self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

												_A								= None
												if use_nin_shortcut:
																	_A								= nn.Conv(
																	    __A		,     kernel_size=(1, 1)		,     strides=(1, 1)		,     padding='''VALID'''		,     dtype=self.dtype		,     )



							def __call__(  self:   Dict		,     __A:   List[Any]		,     __A:   List[Any]		,     __A:   Any=True   )    ->				List[Any]:
												_A								= hidden_states
												_A								= self.norma(__A   )
												_A								= nn.swish(__A   )
												_A								= self.conva(__A   )

												_A								= self.time_emb_proj(nn.swish(__A   )   )
												_A								= jnp.expand_dims(jnp.expand_dims(__A		,     1   )		,     1   )
												_A								= hidden_states + temb

												_A								= self.norma(__A   )
												_A								= nn.swish(__A   )
												_A								= self.dropout(__A		,     __A   )
												_A								= self.conva(__A   )

												if self.conv_shortcut is not None:
																	_A								= self.conv_shortcut(__A   )

												return hidden_states + residual



