

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		:   List[Any]										=       logging.get_logger(__name__)

# General docstring
_UpperCAmelCase		:   Optional[int]										=       """PoolFormerConfig"""

# Base docstring
_UpperCAmelCase		:   Dict										=       """sail/poolformer_s12"""
_UpperCAmelCase		:   str										=       [1, 512, 7, 7]

# Image classification docstring
_UpperCAmelCase		:   Optional[int]										=       """sail/poolformer_s12"""
_UpperCAmelCase		:   Any										=       """tabby, tabby cat"""

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





def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__ = 0.0   ,							UpperCamelCase__ = False      ):







							'''simple docstring'''






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







class  lowercase				(       nn.Module			):




				def __init__(					self      ,   snake_case = None       ):
											super().__init__()
											snake_case_							=    drop_prob




				def     a							(					self      ,   snake_case       ):
											return drop_path(snake_case      ,   self.drop_prob      ,   self.training       )




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







class  lowercase				(       nn.Module			):




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

											snake_case_							=    nn.Convad(snake_case      ,   snake_case      ,   kernel_size=snake_case      ,   stride=snake_case      ,   padding=snake_case       )
											snake_case_							=    norm_layer(snake_case       ) if norm_layer else nn.Identity()




				def     a							(					self      ,   snake_case       ):
											snake_case_							=    self.projection(snake_case       )
											snake_case_							=    self.norm(snake_case       )
											return embeddings







class  lowercase				(       nn.GroupNorm			):




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







class  lowercase				(       nn.Module			):




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




				def     a							(					self      ,   snake_case       ):
											return self.pool(snake_case       ) - hidden_states







class  lowercase				(       nn.Module			):




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




				def     a							(					self      ,   snake_case       ):
											snake_case_							=    self.conva(snake_case       )
											snake_case_							=    self.act_fn(snake_case       )
											snake_case_							=    self.drop(snake_case       )
											snake_case_							=    self.conva(snake_case       )
											snake_case_							=    self.drop(snake_case       )

											return hidden_states







class  lowercase				(       nn.Module			):




				def __init__(					self      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case      ,   snake_case       ):
											super().__init__()
											snake_case_							=    PoolFormerPooling(snake_case       )
											snake_case_							=    PoolFormerOutput(snake_case      ,   snake_case      ,   snake_case      ,   snake_case       )
											snake_case_							=    PoolFormerGroupNorm(snake_case       )
											snake_case_							=    PoolFormerGroupNorm(snake_case       )

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




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

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

																		snake_case_							=    (output,) + outputs
																		return outputs

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

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

																		snake_case_							=    (output,) + outputs
																		return outputs







class  lowercase				(       nn.Module			):




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

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

											# Transformer blocks
											snake_case_							=    []
											snake_case_							=    0
											for i in range(config.num_encoder_blocks       ):
																		# each block consists of layers
																		snake_case_							=    []
																		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       )       )

											snake_case_							=    nn.ModuleList(snake_case       )




				def     a							(					self      ,   snake_case      ,   snake_case=False      ,   snake_case=True       ):
											snake_case_							=    () if output_hidden_states else None

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

																		if output_hidden_states:
																									snake_case_							=    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  lowercase				(       lowercase_			):
				__SCREAMING_SNAKE_CASE					: Tuple       			=			PoolFormerConfig
				__SCREAMING_SNAKE_CASE					: List[Any]       			=			'''poolformer'''
				__SCREAMING_SNAKE_CASE					: Optional[int]       			=			'''pixel_values'''
				__SCREAMING_SNAKE_CASE					: Any       			=			True




				def     a							(					self      ,   snake_case       ):
											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     a							(					self      ,   snake_case      ,   snake_case=False       ):
											if isinstance(snake_case      ,   snake_case       ):
																		snake_case_							=    value



_UpperCAmelCase		:   Optional[Any]										=       R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

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







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




				def __init__(					self      ,   snake_case       ):
											super().__init__(snake_case       )
											snake_case_							=    config

											snake_case_							=    PoolFormerEncoder(snake_case       )

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




				def     a							(					self       ):
											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     a							(					self      ,   snake_case = None      ,   snake_case = None      ,   snake_case = None      ,   ):
											snake_case_							=    (
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											snake_case_							=    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'       )

											snake_case_							=    self.encoder(
											    snake_case      ,   output_hidden_states=snake_case      ,   return_dict=snake_case      ,   )
											snake_case_							=    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  lowercase				(       nn.Module			):




				def __init__(					self      ,   snake_case       ):
											super().__init__()
											snake_case_							=    nn.Linear(config.hidden_size      ,   config.hidden_size       )




				def     a							(					self      ,   snake_case       ):
											snake_case_							=    self.dense(snake_case       )
											return output







@add_start_docstrings(
    '''
    PoolFormer Model transformer with an image classification head on top
    '''	,  lowercase_	,  )
class  lowercase				(       lowercase_			):




				def __init__(					self      ,   snake_case       ):
											super().__init__(snake_case       )
											snake_case_							=    config.num_labels
											snake_case_							=    PoolFormerModel(snake_case       )

											# Final norm
											snake_case_							=    PoolFormerGroupNorm(config.hidden_sizes[-1]       )
											# Classifier head
											snake_case_							=    (
											    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     a							(					self      ,   snake_case = None      ,   snake_case = None      ,   snake_case = None      ,   snake_case = None      ,   ):
											snake_case_							=    return_dict if return_dict is not None else self.config.use_return_dict

											snake_case_							=    self.poolformer(
											    snake_case      ,   output_hidden_states=snake_case      ,   return_dict=snake_case      ,   )

											snake_case_							=    outputs[0]

											snake_case_							=    self.classifier(self.norm(snake_case       ).mean([-2, -1]       )       )

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

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

											if not return_dict:
																		snake_case_							=    (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       )







import uuid
from typing import Any, Dict, List, Optional, Union

from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
					import tensorflow as tf

if is_torch_available():
					import torch


_UpperCAmelCase		:   Dict										=       logging.get_logger(__name__)







class  lowercase				:




				def __init__(					self      ,   snake_case = None      ,   snake_case = None      ,   snake_case=None      ,   snake_case=None       ):
											if not conversation_id:
																		snake_case_							=    uuid.uuida()
											if past_user_inputs is None:
																		snake_case_							=    []
											if generated_responses is None:
																		snake_case_							=    []

											snake_case_							=    conversation_id
											snake_case_							=    past_user_inputs
											snake_case_							=    generated_responses
											snake_case_							=    text




				def __eq__(					self      ,   snake_case       ):
											if not isinstance(snake_case      ,   snake_case       ):
																		return False
											if self.uuid == other.uuid:
																		return True
											return (
											    self.new_user_input == other.new_user_input
											    and self.past_user_inputs == other.past_user_inputs
											    and self.generated_responses == other.generated_responses
											)




				def     a							(					self      ,   snake_case      ,   snake_case = False       ):
											if self.new_user_input:
																		if overwrite:
																									logger.warning(
																									    F'''User input added while unprocessed input was existing: "{self.new_user_input}" was overwritten '''
																									    F'''with: "{text}".'''       )
																									snake_case_							=    text
																		else:
																									logger.warning(
																									    F'''User input added while unprocessed input was existing: "{self.new_user_input}" new input '''
																									    F'''ignored: "{text}". Set `overwrite` to True to overwrite unprocessed user input'''       )
											else:
																		snake_case_							=    text




				def     a							(					self       ):
											if self.new_user_input:
																		self.past_user_inputs.append(self.new_user_input       )
											snake_case_							=    None




				def     a							(					self      ,   snake_case       ):
											self.generated_responses.append(snake_case       )




				def     a							(					self       ):
											for user_input, generated_response in zip(self.past_user_inputs      ,   self.generated_responses       ):
																		yield True, user_input
																		yield False, generated_response
											if self.new_user_input:
																		yield True, self.new_user_input




				def __repr__(					self       ):
											snake_case_							=    F'''Conversation id: {self.uuid} \n'''
											for is_user, text in self.iter_texts():
																		snake_case_							=    'user' if is_user else 'bot'
																		output += F'''{name} >> {text} \n'''
											return output







@add_end_docstrings(
    lowercase_	,  R'''
        min_length_for_response (`int`, *optional*, defaults to 32):
            The minimum length (in number of tokens) for a response.
        minimum_tokens (`int`, *optional*, defaults to 10):
            The minimum length of tokens to leave for a response.
    '''	,  )
class  lowercase				(       lowercase_			):




				def __init__(					self      ,   *snake_case      ,   **snake_case       ):
											super().__init__(*snake_case      ,   **snake_case       )
											if self.tokenizer.pad_token_id is None:
																		snake_case_							=    self.tokenizer.eos_token




				def     a							(					self      ,   snake_case=None      ,   snake_case=None      ,   snake_case=None      ,   **snake_case       ):
											snake_case_							=    {}
											snake_case_							=    {}
											snake_case_							=    {}

											if min_length_for_response is not None:
																		snake_case_							=    min_length_for_response
											if minimum_tokens is not None:
																		snake_case_							=    minimum_tokens

											if "max_length" in generate_kwargs:
																		snake_case_							=    generate_kwargs['max_length']
																		# self.max_length = generate_kwargs.get("max_length", self.model.config.max_length)
											if clean_up_tokenization_spaces is not None:
																		snake_case_							=    clean_up_tokenization_spaces

											if generate_kwargs:
																		forward_params.update(snake_case       )
											return preprocess_params, forward_params, postprocess_params




				def __call__(					self      ,   snake_case      ,   snake_case=0      ,   **snake_case       ):
											snake_case_							=    super().__call__(snake_case      ,   num_workers=snake_case      ,   **snake_case       )
											if isinstance(snake_case      ,   snake_case       ) and len(snake_case       ) == 1:
																		return outputs[0]
											return outputs




				def     a							(					self      ,   snake_case      ,   snake_case=32       ):
											if not isinstance(snake_case      ,   snake_case       ):
																		raise ValueError('ConversationalPipeline, expects Conversation as inputs'       )
											if conversation.new_user_input is None:
																		raise ValueError(
																		    F'''Conversation with UUID {type(conversation.uuid       )} does not contain new user input to process. '''
																		    'Add user inputs with the conversation\'s `add_user_input` method'       )
											if hasattr(self.tokenizer      ,   '_build_conversation_input_ids'       ):
																		snake_case_							=    self.tokenizer._build_conversation_input_ids(snake_case       )
											else:
																		# If the tokenizer cannot handle conversations, we default to only the old version
																		snake_case_							=    self._legacy_parse_and_tokenize(snake_case       )

											if self.framework == "pt":
																		snake_case_							=    torch.LongTensor([input_ids]       )
											elif self.framework == "tf":
																		snake_case_							=    tf.constant([input_ids]       )
											return {"input_ids": input_ids, "conversation": conversation}




				def     a							(					self      ,   snake_case      ,   snake_case=10      ,   **snake_case       ):
											snake_case_							=    generate_kwargs.get('max_length'      ,   self.model.config.max_length       )

											snake_case_							=    model_inputs['input_ids'].shape[1]
											if max_length - minimum_tokens < n:
																		logger.warning(F'''Conversation input is to long ({n}), trimming it to ({max_length} - {minimum_tokens})'''       )
																		snake_case_							=    max_length - minimum_tokens
																		snake_case_							=    model_inputs['input_ids'][:, -trim:]
																		if "attention_mask" in model_inputs:
																									snake_case_							=    model_inputs['attention_mask'][:, -trim:]
											snake_case_							=    model_inputs.pop('conversation'       )
											snake_case_							=    max_length
											snake_case_							=    self.model.generate(**snake_case      ,   **snake_case       )
											if self.model.config.is_encoder_decoder:
																		snake_case_							=    1
											else:
																		snake_case_							=    n
											return {"output_ids": output_ids[:, start_position:], "conversation": conversation}




				def     a							(					self      ,   snake_case      ,   snake_case=True       ):
											snake_case_							=    model_outputs['output_ids']
											snake_case_							=    self.tokenizer.decode(
											    output_ids[0]      ,   skip_special_tokens=snake_case      ,   clean_up_tokenization_spaces=snake_case      ,   )
											snake_case_							=    model_outputs['conversation']
											conversation.mark_processed()
											conversation.append_response(snake_case       )
											return conversation




				def     a							(					self      ,   snake_case       ):
											snake_case_							=    self.tokenizer.eos_token_id
											snake_case_							=    []
											for is_user, text in conversation.iter_texts():
																		if eos_token_id is not None:
																									input_ids.extend(self.tokenizer.encode(snake_case      ,   add_special_tokens=snake_case       ) + [eos_token_id]       )
																		else:
																									input_ids.extend(self.tokenizer.encode(snake_case      ,   add_special_tokens=snake_case       )       )

											if len(snake_case       ) > self.tokenizer.model_max_length:
																		snake_case_							=    input_ids[-self.tokenizer.model_max_length :]
											return input_ids




