
'''simple docstring'''
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, List, Mapping, Optional

from packaging import version


if TYPE_CHECKING:
		from ... import PreTrainedTokenizer, TensorType

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import is_torch_available, logging


lowercase__											=logging.get_logger(__name__)

lowercase__											={
    'bigscience/bloom': 'https://huggingface.co/bigscience/bloom/resolve/main/config.json',
    'bigscience/bloom-560m': 'https://huggingface.co/bigscience/bloom-560m/blob/main/config.json',
    'bigscience/bloom-1b1': 'https://huggingface.co/bigscience/bloom-1b1/blob/main/config.json',
    'bigscience/bloom-1b7': 'https://huggingface.co/bigscience/bloom-1b7/blob/main/config.json',
    'bigscience/bloom-3b': 'https://huggingface.co/bigscience/bloom-3b/blob/main/config.json',
    'bigscience/bloom-7b1': 'https://huggingface.co/bigscience/bloom-7b1/blob/main/config.json',
}







class   a_ (      UpperCamelCase__					):
			lowerCamelCase__	:			List[str]       =    'bloom'
			lowerCamelCase__	:			Optional[int]       =    ['past_key_values']
			lowerCamelCase__	:			Dict       =    {
			    'num_hidden_layers': 'n_layer',
			    'num_attention_heads': 'n_head',
			}




			def __init__(       self					,   UpperCAmelCase=25_08_80					,   UpperCAmelCase=64					,   UpperCAmelCase=2					,   UpperCAmelCase=8					,   UpperCAmelCase=1e-5					,   UpperCAmelCase=0.02					,   UpperCAmelCase=True					,   UpperCAmelCase=1					,   UpperCAmelCase=2					,   UpperCAmelCase=False					,   UpperCAmelCase=0.0					,   UpperCAmelCase=0.0					,   UpperCAmelCase=1					,   UpperCAmelCase=False					,   **UpperCAmelCase					,   ):
				a_      				=						vocab_size
				# Backward compatibility with n_embed kwarg
				a_      				=						kwargs.pop("""n_embed"""					,   UpperCAmelCase	)
				a_      				=						hidden_size if n_embed is None else n_embed
				a_      				=						n_layer
				a_      				=						n_head
				a_      				=						layer_norm_epsilon
				a_      				=						initializer_range
				a_      				=						use_cache
				a_      				=						pretraining_tp
				a_      				=						apply_residual_connection_post_layernorm
				a_      				=						hidden_dropout
				a_      				=						attention_dropout

				a_      				=						bos_token_id
				a_      				=						eos_token_id
				a_      				=						slow_but_exact

				super().__init__(bos_token_id=UpperCAmelCase					,   eos_token_id=UpperCAmelCase					,   **UpperCAmelCase	)







class   a_ (      UpperCamelCase__					):
			lowerCamelCase__	:			Optional[int]       =    version.parse('1.12'					)




			def __init__(       self					,   UpperCAmelCase					,   UpperCAmelCase = "default"					,   UpperCAmelCase = None					,   UpperCAmelCase = False					,   ):
				super().__init__(UpperCAmelCase					,   task=UpperCAmelCase					,   patching_specs=UpperCAmelCase					,   use_past=UpperCAmelCase	)
				if not getattr(self._config					,   """pad_token_id"""					,   UpperCAmelCase	):
					# TODO: how to do that better?
					a_      				=						0




			@property
			def 					lowerCAmelCase__					(       self	):
				a_      				=						OrderedDict({"""input_ids""": {0: """batch""", 1: """sequence"""}}	)
				if self.use_past:
					# BLOOM stores values on dynamic axis 2. For more details see: https://github.com/huggingface/transformers/pull/18344
					self.fill_with_past_key_values_(UpperCAmelCase					,   direction="""inputs"""					,   inverted_values_shape=UpperCAmelCase	)
					a_      				=						{0: """batch""", 1: """past_sequence + sequence"""}
				else:
					a_      				=						{0: """batch""", 1: """sequence"""}

				return common_inputs




			@property
			def 					lowerCAmelCase__					(       self	):
				return self._config.n_layer




			@property
			def 					lowerCAmelCase__					(       self	):
				return self._config.n_head




			@property
			def 					lowerCAmelCase__					(       self	):
				return 1e-3




			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase					,   UpperCAmelCase = -1					,   UpperCAmelCase = -1					,   UpperCAmelCase = False					,   UpperCAmelCase = None					,   ):
				a_      				=						super(UpperCAmelCase					,   self	).generate_dummy_inputs(
				    UpperCAmelCase					,   batch_size=UpperCAmelCase					,   seq_length=UpperCAmelCase					,   is_pair=UpperCAmelCase					,   framework=UpperCAmelCase	)

				# We need to order the input in the way they appears in the forward()
				a_      				=						OrderedDict({"""input_ids""": common_inputs["""input_ids"""]}	)

				# Need to add the past_keys
				if self.use_past:
					if not is_torch_available():
						raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""	)
					else:
						import torch

						a_				,  a_      				=						common_inputs["""input_ids"""].shape
						# Not using the same length for past_key_values
						a_      				=						seqlen + 2
						a_      				=						self._config.hidden_size // self.num_attention_heads
						a_      				=						(
						    batch * self.num_attention_heads,
						    head_dim,
						    past_key_values_length,
						)
						a_      				=						(
						    batch * self.num_attention_heads,
						    past_key_values_length,
						    head_dim,
						)
						a_      				=						[
						    (torch.zeros(UpperCAmelCase	), torch.zeros(UpperCAmelCase	)) for _ in range(self.num_layers	)
						]

				a_      				=						common_inputs["""attention_mask"""]
				if self.use_past:
					a_      				=						ordered_inputs["""attention_mask"""].dtype
					a_      				=						torch.cat(
					    [ordered_inputs["""attention_mask"""], torch.ones(UpperCAmelCase					,   UpperCAmelCase					,   dtype=UpperCAmelCase	)]					,   dim=1	)

				return ordered_inputs






			@property
			def 					lowerCAmelCase__					(       self	):
				return 13


'''simple docstring'''
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


lowercase__											=logging.get_logger(__name__)

# General docstring
lowercase__											='PoolFormerConfig'

# Base docstring
lowercase__											='sail/poolformer_s12'
lowercase__											=[1, 5_12, 7, 7]

# Image classification docstring
lowercase__											='sail/poolformer_s12'
lowercase__											='tabby, tabby cat'

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






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







class   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase = None	):
				super().__init__()
				a_      				=						drop_prob




			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				return drop_path(UpperCAmelCase					,   self.drop_prob					,   self.training	)






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







class   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase=None	):
				super().__init__()
				a_      				=						patch_size if isinstance(UpperCAmelCase					,   collections.abc.Iterable	) else (patch_size, patch_size)
				a_      				=						stride if isinstance(UpperCAmelCase					,   collections.abc.Iterable	) else (stride, stride)
				a_      				=						padding if isinstance(UpperCAmelCase					,   collections.abc.Iterable	) else (padding, padding)

				a_      				=						nn.Convad(UpperCAmelCase					,   UpperCAmelCase					,   kernel_size=UpperCAmelCase					,   stride=UpperCAmelCase					,   padding=UpperCAmelCase	)
				a_      				=						norm_layer(UpperCAmelCase	) if norm_layer else nn.Identity()






			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				a_      				=						self.projection(UpperCAmelCase	)
				a_      				=						self.norm(UpperCAmelCase	)
				return embeddings







class   a_ (      nn.GroupNorm					):




			def __init__(       self					,   UpperCAmelCase					,   **UpperCAmelCase	):
				super().__init__(1					,   UpperCAmelCase					,   **UpperCAmelCase	)







class   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase	):
				super().__init__()
				a_      				=						nn.AvgPoolad(UpperCAmelCase					,   stride=1					,   padding=pool_size // 2					,   count_include_pad=UpperCAmelCase	)






			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				return self.pool(UpperCAmelCase	) - hidden_states







class   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase	):
				super().__init__()
				a_      				=						nn.Convad(UpperCAmelCase					,   UpperCAmelCase					,   1	)
				a_      				=						nn.Convad(UpperCAmelCase					,   UpperCAmelCase					,   1	)
				a_      				=						PoolFormerDropPath(UpperCAmelCase	)
				if isinstance(config.hidden_act					,   UpperCAmelCase	):
					a_      				=						ACTaFN[config.hidden_act]
				else:
					a_      				=						config.hidden_act






			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				a_      				=						self.conva(UpperCAmelCase	)
				a_      				=						self.act_fn(UpperCAmelCase	)
				a_      				=						self.drop(UpperCAmelCase	)
				a_      				=						self.conva(UpperCAmelCase	)
				a_      				=						self.drop(UpperCAmelCase	)

				return hidden_states







class   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase	):
				super().__init__()
				a_      				=						PoolFormerPooling(UpperCAmelCase	)
				a_      				=						PoolFormerOutput(UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase	)
				a_      				=						PoolFormerGroupNorm(UpperCAmelCase	)
				a_      				=						PoolFormerGroupNorm(UpperCAmelCase	)

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






			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				if self.use_layer_scale:
					a_      				=						self.pooling(self.before_norm(UpperCAmelCase	)	)
					a_      				=						self.layer_scale_a.unsqueeze(-1	).unsqueeze(-1	) * pooling_output
					# First residual connection
					a_      				=						hidden_states + self.drop_path(UpperCAmelCase	)
					a_      				=						()

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

					a_      				=						(output,) + outputs
					return outputs

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

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

					a_      				=						(output,) + outputs
					return outputs







class   a_ (      nn.Module					):




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

				# patch embeddings
				a_      				=						[]
				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]					,   )	)
				a_      				=						nn.ModuleList(UpperCAmelCase	)

				# Transformer blocks
				a_      				=						[]
				a_      				=						0
				for i in range(config.num_encoder_blocks	):
					# each block consists of layers
					a_      				=						[]
					if i != 0:
						cur += config.depths[i - 1]
					for j in range(config.depths[i]	):
						layers.append(
						    PoolFormerLayer(
						        UpperCAmelCase					,   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(UpperCAmelCase	)	)

				a_      				=						nn.ModuleList(UpperCAmelCase	)






			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase					,   UpperCAmelCase=False					,   UpperCAmelCase=True	):
				a_      				=						() if output_hidden_states else None

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

					if output_hidden_states:
						a_      				=						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=UpperCAmelCase					,   hidden_states=UpperCAmelCase	)







class   a_ (      UpperCamelCase__					):
			lowerCamelCase__	:			Union[str, Any]       =    PoolFormerConfig
			lowerCamelCase__	:			Optional[Any]       =    'poolformer'
			lowerCamelCase__	:			List[Any]       =    'pixel_values'
			lowerCamelCase__	:			int       =    True




			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				if isinstance(UpperCAmelCase					,   (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(UpperCAmelCase					,   nn.LayerNorm	):
					module.bias.data.zero_()
					module.weight.data.fill_(1.0	)






			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase					,   UpperCAmelCase=False	):
				if isinstance(UpperCAmelCase					,   UpperCAmelCase	):
					a_      				=						value







lowercase__											=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'

lowercase__											=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					,   UpperCAmelCase	):
				super().__init__(UpperCAmelCase	)
				a_      				=						config

				a_      				=						PoolFormerEncoder(UpperCAmelCase	)

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




			def 					lowerCAmelCase__					(       self	):
				return self.embeddings.patch_embeddings






			@add_start_docstrings_to_model_forward(UpperCAmelCase	)
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC					,   output_type=UpperCAmelCase					,   config_class=_CONFIG_FOR_DOC					,   modality="""vision"""					,   expected_output=_EXPECTED_OUTPUT_SHAPE					,   )
			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase = None					,   UpperCAmelCase = None					,   UpperCAmelCase = None					,   ):
				a_      				=						(
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				a_      				=						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"""	)

				a_      				=						self.encoder(
				    UpperCAmelCase					,   output_hidden_states=UpperCAmelCase					,   return_dict=UpperCAmelCase					,   )
				a_      				=						encoder_outputs[0]

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

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







class   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase	):
				super().__init__()
				a_      				=						nn.Linear(config.hidden_size					,   config.hidden_size	)






			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				a_      				=						self.dense(UpperCAmelCase	)
				return output







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




			def __init__(       self					,   UpperCAmelCase	):
				super().__init__(UpperCAmelCase	)
				a_      				=						config.num_labels
				a_      				=						PoolFormerModel(UpperCAmelCase	)

				# Final norm
				a_      				=						PoolFormerGroupNorm(config.hidden_sizes[-1]	)
				# Classifier head
				a_      				=						(
				    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(UpperCAmelCase	)
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT					,   output_type=UpperCAmelCase					,   config_class=_CONFIG_FOR_DOC					,   expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT					,   )
			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase = None					,   UpperCAmelCase = None					,   UpperCAmelCase = None					,   UpperCAmelCase = None					,   ):
				a_      				=						return_dict if return_dict is not None else self.config.use_return_dict

				a_      				=						self.poolformer(
				    UpperCAmelCase					,   output_hidden_states=UpperCAmelCase					,   return_dict=UpperCAmelCase					,   )

				a_      				=						outputs[0]

				a_      				=						self.classifier(self.norm(UpperCAmelCase	).mean([-2, -1]	)	)

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

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

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

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