
import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


UpperCamelCase_   	=	logging.get_logger(__name__)

# General docstring
UpperCamelCase_   	=	"PoolFormerConfig"

# Base docstring
UpperCamelCase_   	=	"sail/poolformer_s12"
UpperCamelCase_   	=	[1, 5_1_2, 7, 7]

# Image classification docstring
UpperCamelCase_   	=	"sail/poolformer_s12"
UpperCamelCase_   	=	"tabby, tabby cat"

UpperCamelCase_   	=	[
    "sail/poolformer_s12",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]






def 		_UpperCAmelCase   (	UpperCamelCase:     List[str]			,							UpperCamelCase:     float = 0.0			,							UpperCamelCase:     bool = False  ):






	"""simple docstring"""


	if drop_prob == 0.0 or not training:
		return input
	__lowerCAmelCase         =							1 - drop_prob
	__lowerCAmelCase         =							(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
	__lowerCAmelCase         =							keep_prob + torch.rand(UpperCamelCase			,							dtype=input.dtype			,							device=input.device  )
	random_tensor.floor_()  # binarize
	__lowerCAmelCase         =							input.div(UpperCamelCase  ) * random_tensor
	return output

class       a		(   nn.Module					):







							def __init__(						self  :  Union[str, Any]      ,				snake_case__  :  Optional[float] = None ):
								"""simple docstring"""







								super().__init__()
								__lowerCAmelCase         =							drop_prob







							def       UpperCAmelCase__					(						self  :  Optional[Any]      ,				snake_case__  :  torch.Tensor ):
								"""simple docstring"""







								return drop_path(snake_case__      ,				self.drop_prob      ,				self.training )







							def       UpperCAmelCase__					(						self  :  List[Any] ):
								"""simple docstring"""







								return "p={}".format(self.drop_prob )

class       a		(   nn.Module					):







							def __init__(						self  :  Tuple      ,				snake_case__  :  Tuple      ,				snake_case__  :  Optional[int]      ,				snake_case__  :  List[str]      ,				snake_case__  :  Any      ,				snake_case__  :  Optional[Any]      ,				snake_case__  :  List[str]=None ):
								"""simple docstring"""







								super().__init__()
								__lowerCAmelCase         =							patch_size if isinstance(snake_case__      ,				collections.abc.Iterable ) else (patch_size, patch_size)
								__lowerCAmelCase         =							stride if isinstance(snake_case__      ,				collections.abc.Iterable ) else (stride, stride)
								__lowerCAmelCase         =							padding if isinstance(snake_case__      ,				collections.abc.Iterable ) else (padding, padding)

								__lowerCAmelCase         =							nn.Convad(snake_case__      ,				snake_case__      ,				kernel_size=snake_case__      ,				stride=snake_case__      ,				padding=snake_case__ )
								__lowerCAmelCase         =							norm_layer(snake_case__ ) if norm_layer else nn.Identity()







							def       UpperCAmelCase__					(						self  :  Optional[int]      ,				snake_case__  :  Union[str, Any] ):
								"""simple docstring"""







								__lowerCAmelCase         =							self.projection(snake_case__ )
								__lowerCAmelCase         =							self.norm(snake_case__ )
								return embeddings

class       a		(   nn.GroupNorm					):







							def __init__(						self  :  str      ,				snake_case__  :  str      ,				**snake_case__  :  List[str] ):
								"""simple docstring"""







								super().__init__(1      ,				snake_case__      ,				**snake_case__ )

class       a		(   nn.Module					):







							def __init__(						self  :  Any      ,				snake_case__  :  Union[str, Any] ):
								"""simple docstring"""







								super().__init__()
								__lowerCAmelCase         =							nn.AvgPoolad(snake_case__      ,				stride=1      ,				padding=pool_size // 2      ,				count_include_pad=snake_case__ )







							def       UpperCAmelCase__					(						self  :  str      ,				snake_case__  :  List[Any] ):
								"""simple docstring"""







								return self.pool(snake_case__ ) - hidden_states

class       a		(   nn.Module					):







							def __init__(						self  :  Optional[int]      ,				snake_case__  :  Union[str, Any]      ,				snake_case__  :  List[Any]      ,				snake_case__  :  Tuple      ,				snake_case__  :  Any ):
								"""simple docstring"""







								super().__init__()
								__lowerCAmelCase         =							nn.Convad(snake_case__      ,				snake_case__      ,				1 )
								__lowerCAmelCase         =							nn.Convad(snake_case__      ,				snake_case__      ,				1 )
								__lowerCAmelCase         =							PoolFormerDropPath(snake_case__ )
								if isinstance(config.hidden_act      ,				snake_case__ ):
									__lowerCAmelCase         =							ACTaFN[config.hidden_act]
								else:
									__lowerCAmelCase         =							config.hidden_act







							def       UpperCAmelCase__					(						self  :  Optional[Any]      ,				snake_case__  :  str ):
								"""simple docstring"""







								__lowerCAmelCase         =							self.conva(snake_case__ )
								__lowerCAmelCase         =							self.act_fn(snake_case__ )
								__lowerCAmelCase         =							self.drop(snake_case__ )
								__lowerCAmelCase         =							self.conva(snake_case__ )
								__lowerCAmelCase         =							self.drop(snake_case__ )

								return hidden_states

class       a		(   nn.Module					):







							def __init__(						self  :  Optional[int]      ,				snake_case__  :  Any      ,				snake_case__  :  Tuple      ,				snake_case__  :  Optional[int]      ,				snake_case__  :  Optional[int]      ,				snake_case__  :  str      ,				snake_case__  :  Optional[int] ):
								"""simple docstring"""







								super().__init__()
								__lowerCAmelCase         =							PoolFormerPooling(snake_case__ )
								__lowerCAmelCase         =							PoolFormerOutput(snake_case__      ,				snake_case__      ,				snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							PoolFormerGroupNorm(snake_case__ )
								__lowerCAmelCase         =							PoolFormerGroupNorm(snake_case__ )

								# Useful for training neural nets
								__lowerCAmelCase         =							PoolFormerDropPath(snake_case__ ) if drop_path > 0.0 else nn.Identity()
								__lowerCAmelCase         =							config.use_layer_scale
								if config.use_layer_scale:
									__lowerCAmelCase         =							nn.Parameter(
									    config.layer_scale_init_value * torch.ones((snake_case__) )      ,				requires_grad=snake_case__ )
									__lowerCAmelCase         =							nn.Parameter(
									    config.layer_scale_init_value * torch.ones((snake_case__) )      ,				requires_grad=snake_case__ )







							def       UpperCAmelCase__					(						self  :  Dict      ,				snake_case__  :  Optional[int] ):
								"""simple docstring"""







								if self.use_layer_scale:
									__lowerCAmelCase         =							self.pooling(self.before_norm(snake_case__ ) )
									__lowerCAmelCase         =							self.layer_scale_a.unsqueeze(-1 ).unsqueeze(-1 ) * pooling_output
									# First residual connection
									__lowerCAmelCase         =							hidden_states + self.drop_path(snake_case__ )
									__lowerCAmelCase         =							()

									__lowerCAmelCase         =							self.output(self.after_norm(snake_case__ ) )
									__lowerCAmelCase         =							self.layer_scale_a.unsqueeze(-1 ).unsqueeze(-1 ) * layer_output
									# Second residual connection
									__lowerCAmelCase         =							hidden_states + self.drop_path(snake_case__ )

									__lowerCAmelCase         =							(output,) + outputs
									return outputs

								else:
									__lowerCAmelCase         =							self.drop_path(self.pooling(self.before_norm(snake_case__ ) ) )
									# First residual connection
									__lowerCAmelCase         =							pooling_output + hidden_states
									__lowerCAmelCase         =							()

									# Second residual connection inside the PoolFormerOutput block
									__lowerCAmelCase         =							self.drop_path(self.output(self.after_norm(snake_case__ ) ) )
									__lowerCAmelCase         =							hidden_states + layer_output

									__lowerCAmelCase         =							(output,) + outputs
									return outputs

class       a		(   nn.Module					):







							def __init__(						self  :  Any      ,				snake_case__  :  int ):
								"""simple docstring"""







								super().__init__()
								__lowerCAmelCase         =							config
								# stochastic depth decay rule
								__lowerCAmelCase         =							[x.item() for x in torch.linspace(0      ,				config.drop_path_rate      ,				sum(config.depths ) )]

								# patch embeddings
								__lowerCAmelCase         =							[]
								for i in range(config.num_encoder_blocks ):
									embeddings.append(
									    PoolFormerEmbeddings(
									        patch_size=config.patch_sizes[i]      ,				stride=config.strides[i]      ,				padding=config.padding[i]      ,				num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]      ,				hidden_size=config.hidden_sizes[i]      ,				) )
								__lowerCAmelCase         =							nn.ModuleList(snake_case__ )

								# Transformer blocks
								__lowerCAmelCase         =							[]
								__lowerCAmelCase         =							0
								for i in range(config.num_encoder_blocks ):
									# each block consists of layers
									__lowerCAmelCase         =							[]
									if i != 0:
										cur += config.depths[i - 1]
									for j in range(config.depths[i] ):
										layers.append(
										    PoolFormerLayer(
										        snake_case__      ,				num_channels=config.hidden_sizes[i]      ,				pool_size=config.pool_size      ,				hidden_size=config.hidden_sizes[i]      ,				intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio )      ,				drop_path=dpr[cur + j]      ,				) )
									blocks.append(nn.ModuleList(snake_case__ ) )

								__lowerCAmelCase         =							nn.ModuleList(snake_case__ )







							def       UpperCAmelCase__					(						self  :  Tuple      ,				snake_case__  :  Dict      ,				snake_case__  :  List[Any]=False      ,				snake_case__  :  Dict=True ):
								"""simple docstring"""







								__lowerCAmelCase         =							() if output_hidden_states else None

								__lowerCAmelCase         =							pixel_values
								for idx, layers in enumerate(zip(self.patch_embeddings      ,				self.block ) ):
									__lowerCAmelCase   ,						__lowerCAmelCase         =							layers
									# Get patch embeddings from hidden_states
									__lowerCAmelCase         =							embedding_layer(snake_case__ )
									# Send the embeddings through the blocks
									for _, blk in enumerate(snake_case__ ):
										__lowerCAmelCase         =							blk(snake_case__ )
										__lowerCAmelCase         =							layer_outputs[0]

									if output_hidden_states:
										__lowerCAmelCase         =							all_hidden_states + (hidden_states,)

								if not return_dict:
									return tuple(v for v in [hidden_states, all_hidden_states] if v is not None )

								return BaseModelOutputWithNoAttention(last_hidden_state=snake_case__      ,				hidden_states=snake_case__ )

class       a		(   __UpperCAmelCase					):
							lowercase_		:						Union[str, Any]   							=  PoolFormerConfig
							lowercase_		:						int   							=  'poolformer'
							lowercase_		:						Optional[Any]   							=  'pixel_values'
							lowercase_		:						Dict   							=  True







							def       UpperCAmelCase__					(						self  :  Any      ,				snake_case__  :  Any ):
								"""simple docstring"""







								if isinstance(snake_case__      ,				(nn.Linear, nn.Convad) ):
									module.weight.data.normal_(mean=0.0      ,				std=self.config.initializer_range )
									if module.bias is not None:
										module.bias.data.zero_()
								elif isinstance(snake_case__      ,				nn.LayerNorm ):
									module.bias.data.zero_()
									module.weight.data.fill_(1.0 )







							def       UpperCAmelCase__					(						self  :  str      ,				snake_case__  :  List[str]      ,				snake_case__  :  str=False ):
								"""simple docstring"""







								if isinstance(snake_case__      ,				snake_case__ ):
									__lowerCAmelCase         =							value
UpperCamelCase_   	=	R"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

UpperCamelCase_   	=	R"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`PoolFormerImageProcessor.__call__`] for details.\n"

@add_start_docstrings(
    'The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top.'      ,	__UpperCAmelCase      ,	)
class       a		(   __UpperCAmelCase					):







							def __init__(						self  :  str      ,				snake_case__  :  Tuple ):
								"""simple docstring"""







								super().__init__(snake_case__ )
								__lowerCAmelCase         =							config

								__lowerCAmelCase         =							PoolFormerEncoder(snake_case__ )

								# Initialize weights and apply final processing
								self.post_init()







							def       UpperCAmelCase__					(						self  :  Union[str, Any] ):
								"""simple docstring"""







								return self.embeddings.patch_embeddings







							@add_start_docstrings_to_model_forward(snake_case__ )
							@add_code_sample_docstrings(
							    checkpoint=_CHECKPOINT_FOR_DOC      ,				output_type=snake_case__      ,				config_class=_CONFIG_FOR_DOC      ,				modality="vision"      ,				expected_output=_EXPECTED_OUTPUT_SHAPE      ,				)
							def       UpperCAmelCase__					(						self  :  int      ,				snake_case__  :  Optional[torch.FloatTensor] = None      ,				snake_case__  :  Optional[bool] = None      ,				snake_case__  :  Optional[bool] = None      ,				):
								"""simple docstring"""







								__lowerCAmelCase         =							(
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								__lowerCAmelCase         =							return_dict if return_dict is not None else self.config.use_return_dict

								if pixel_values is None:
									raise ValueError("You have to specify pixel_values" )

								__lowerCAmelCase         =							self.encoder(
								    snake_case__      ,				output_hidden_states=snake_case__      ,				return_dict=snake_case__      ,				)
								__lowerCAmelCase         =							encoder_outputs[0]

								if not return_dict:
									return (sequence_output, None) + encoder_outputs[1:]

								return BaseModelOutputWithNoAttention(
								    last_hidden_state=snake_case__      ,				hidden_states=encoder_outputs.hidden_states      ,				)

class       a		(   nn.Module					):







							def __init__(						self  :  List[Any]      ,				snake_case__  :  Tuple ):
								"""simple docstring"""







								super().__init__()
								__lowerCAmelCase         =							nn.Linear(config.hidden_size      ,				config.hidden_size )







							def       UpperCAmelCase__					(						self  :  Tuple      ,				snake_case__  :  List[Any] ):
								"""simple docstring"""







								__lowerCAmelCase         =							self.dense(snake_case__ )
								return output

@add_start_docstrings(
    '\n    PoolFormer Model transformer with an image classification head on top\n    '      ,	__UpperCAmelCase      ,	)
class       a		(   __UpperCAmelCase					):







							def __init__(						self  :  Union[str, Any]      ,				snake_case__  :  Optional[Any] ):
								"""simple docstring"""







								super().__init__(snake_case__ )
								__lowerCAmelCase         =							config.num_labels
								__lowerCAmelCase         =							PoolFormerModel(snake_case__ )

								# Final norm
								__lowerCAmelCase         =							PoolFormerGroupNorm(config.hidden_sizes[-1] )
								# Classifier head
								__lowerCAmelCase         =							(
								    nn.Linear(config.hidden_sizes[-1]      ,				config.num_labels ) if config.num_labels > 0 else nn.Identity()
								)

								# Initialize weights and apply final processing
								self.post_init()







							@add_start_docstrings_to_model_forward(snake_case__ )
							@add_code_sample_docstrings(
							    checkpoint=_IMAGE_CLASS_CHECKPOINT      ,				output_type=snake_case__      ,				config_class=_CONFIG_FOR_DOC      ,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,				)
							def       UpperCAmelCase__					(						self  :  Tuple      ,				snake_case__  :  Optional[torch.FloatTensor] = None      ,				snake_case__  :  Optional[torch.LongTensor] = None      ,				snake_case__  :  Optional[bool] = None      ,				snake_case__  :  Optional[bool] = None      ,				):
								"""simple docstring"""







								__lowerCAmelCase         =							return_dict if return_dict is not None else self.config.use_return_dict

								__lowerCAmelCase         =							self.poolformer(
								    snake_case__      ,				output_hidden_states=snake_case__      ,				return_dict=snake_case__      ,				)

								__lowerCAmelCase         =							outputs[0]

								__lowerCAmelCase         =							self.classifier(self.norm(snake_case__ ).mean([-2, -1] ) )

								__lowerCAmelCase         =							None
								if labels is not None:
									if self.config.problem_type is None:
										if self.num_labels == 1:
											__lowerCAmelCase         =							"regression"
										elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
											__lowerCAmelCase         =							"single_label_classification"
										else:
											__lowerCAmelCase         =							"multi_label_classification"

									if self.config.problem_type == "regression":
										__lowerCAmelCase         =							MSELoss()
										if self.num_labels == 1:
											__lowerCAmelCase         =							loss_fct(logits.squeeze()      ,				labels.squeeze() )
										else:
											__lowerCAmelCase         =							loss_fct(snake_case__      ,				snake_case__ )
									elif self.config.problem_type == "single_label_classification":
										__lowerCAmelCase         =							CrossEntropyLoss()
										__lowerCAmelCase         =							loss_fct(logits.view(-1      ,				self.num_labels )      ,				labels.view(-1 ) )
									elif self.config.problem_type == "multi_label_classification":
										__lowerCAmelCase         =							BCEWithLogitsLoss()
										__lowerCAmelCase         =							loss_fct(snake_case__      ,				snake_case__ )

								if not return_dict:
									__lowerCAmelCase         =							(logits,) + outputs[2:]
									return ((loss,) + output) if loss is not None else output

								return ImageClassifierOutputWithNoAttention(loss=snake_case__      ,				logits=snake_case__      ,				hidden_states=outputs.hidden_states )








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


UpperCamelCase_   	=	logging.get_logger(__name__)

UpperCamelCase_   	=	{
    "speechbrain/m-ctc-t-large": "https://huggingface.co/speechbrain/m-ctc-t-large/resolve/main/config.json",
    # See all M-CTC-T models at https://huggingface.co/models?filter=mctct
}

class       a		(   __UpperCAmelCase					):
							lowercase_		:						Union[str, Any]   							=  'mctct'







							def __init__(						self  :  Dict      ,				snake_case__  :  Any=8_065      ,				snake_case__  :  Tuple=1_536      ,				snake_case__  :  Union[str, Any]=36      ,				snake_case__  :  List[Any]=6_144      ,				snake_case__  :  Tuple=4      ,				snake_case__  :  List[str]=384      ,				snake_case__  :  List[Any]=920      ,				snake_case__  :  Any=1E-5      ,				snake_case__  :  List[Any]=0.3      ,				snake_case__  :  List[str]="relu"      ,				snake_case__  :  List[Any]=0.0_2      ,				snake_case__  :  Optional[Any]=0.3      ,				snake_case__  :  List[str]=0.3      ,				snake_case__  :  List[Any]=1      ,				snake_case__  :  Optional[int]=0      ,				snake_case__  :  Any=2      ,				snake_case__  :  Any=1      ,				snake_case__  :  List[str]=0.3      ,				snake_case__  :  Tuple=1      ,				snake_case__  :  Union[str, Any]=(7,)      ,				snake_case__  :  List[Any]=(3,)      ,				snake_case__  :  List[Any]=80      ,				snake_case__  :  Optional[int]=1      ,				snake_case__  :  List[str]=None      ,				snake_case__  :  int="sum"      ,				snake_case__  :  Dict=False      ,				**snake_case__  :  List[str]      ,				):
								"""simple docstring"""







								super().__init__(**snake_case__      ,				pad_token_id=snake_case__      ,				bos_token_id=snake_case__      ,				eos_token_id=snake_case__ )
								__lowerCAmelCase         =							vocab_size
								__lowerCAmelCase         =							hidden_size
								__lowerCAmelCase         =							num_hidden_layers
								__lowerCAmelCase         =							intermediate_size
								__lowerCAmelCase         =							num_attention_heads
								__lowerCAmelCase         =							attention_head_dim
								__lowerCAmelCase         =							max_position_embeddings
								__lowerCAmelCase         =							layer_norm_eps
								__lowerCAmelCase         =							layerdrop
								__lowerCAmelCase         =							hidden_act
								__lowerCAmelCase         =							initializer_range
								__lowerCAmelCase         =							hidden_dropout_prob
								__lowerCAmelCase         =							attention_probs_dropout_prob
								__lowerCAmelCase         =							pad_token_id
								__lowerCAmelCase         =							bos_token_id
								__lowerCAmelCase         =							eos_token_id
								__lowerCAmelCase         =							conv_glu_dim
								__lowerCAmelCase         =							conv_dropout
								__lowerCAmelCase         =							num_conv_layers
								__lowerCAmelCase         =							input_feat_per_channel
								__lowerCAmelCase         =							input_channels
								__lowerCAmelCase         =							conv_channels
								__lowerCAmelCase         =							ctc_loss_reduction
								__lowerCAmelCase         =							ctc_zero_infinity

								# prevents config testing fail with exporting to json
								__lowerCAmelCase         =							list(snake_case__ )
								__lowerCAmelCase         =							list(snake_case__ )

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






