
"""simple docstring"""




import unittest

from transformers import FunnelConfig, FunnelTokenizer, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    MODEL_FOR_PRETRAINING_MAPPING,
	    FunnelBaseModel,
	    FunnelForMaskedLM,
	    FunnelForMultipleChoice,
	    FunnelForPreTraining,
	    FunnelForQuestionAnswering,
	    FunnelForSequenceClassification,
	    FunnelForTokenClassification,
	    FunnelModel,
	)




class 					UpperCamelCase     :

				"""simple docstring"""



				def __init__(    self		,     lowercase__		,     lowercase__=13		,     lowercase__=7		,     lowercase__=True		,     lowercase__=True		,     lowercase__=True		,     lowercase__=True		,     lowercase__=99		,     lowercase__=[1, 1, 2]		,     lowercase__=1		,     lowercase__=32		,     lowercase__=4		,     lowercase__=8		,     lowercase__=37		,     lowercase__="gelu_new"		,     lowercase__=0.1		,     lowercase__=0.1		,     lowercase__=0.0		,     lowercase__=512		,     lowercase__=3		,     lowercase__=0.02		,     lowercase__=3		,     lowercase__=4		,     lowercase__=None		,     lowercase__=False		,     )   ->   int:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     parent
									SCREAMING_SNAKE_CASE           =     batch_size
									SCREAMING_SNAKE_CASE           =     seq_length
									SCREAMING_SNAKE_CASE           =     is_training
									SCREAMING_SNAKE_CASE           =     use_input_mask
									SCREAMING_SNAKE_CASE           =     use_token_type_ids
									SCREAMING_SNAKE_CASE           =     use_labels
									SCREAMING_SNAKE_CASE           =     vocab_size
									SCREAMING_SNAKE_CASE           =     block_sizes
									SCREAMING_SNAKE_CASE           =     num_decoder_layers
									SCREAMING_SNAKE_CASE           =     d_model
									SCREAMING_SNAKE_CASE           =     n_head
									SCREAMING_SNAKE_CASE           =     d_head
									SCREAMING_SNAKE_CASE           =     d_inner
									SCREAMING_SNAKE_CASE           =     hidden_act
									SCREAMING_SNAKE_CASE           =     hidden_dropout
									SCREAMING_SNAKE_CASE           =     attention_dropout
									SCREAMING_SNAKE_CASE           =     activation_dropout
									SCREAMING_SNAKE_CASE           =     max_position_embeddings
									SCREAMING_SNAKE_CASE           =     type_vocab_size
									SCREAMING_SNAKE_CASE           =     2
									SCREAMING_SNAKE_CASE           =     num_labels
									SCREAMING_SNAKE_CASE           =     num_choices
									SCREAMING_SNAKE_CASE           =     scope
									SCREAMING_SNAKE_CASE           =     initializer_std

									# Used in the tests to check the size of the first attention layer
									SCREAMING_SNAKE_CASE           =     n_head
									# Used in the tests to check the size of the first hidden state
									SCREAMING_SNAKE_CASE           =     self.d_model
									# Used in the tests to check the number of output hidden states/attentions
									SCREAMING_SNAKE_CASE           =     sum(self.block_sizes      ) + (0 if base else self.num_decoder_layers)
									# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
									# the last hidden state of the first block (which is the first hidden state of the decoder).
									if not base:
														SCREAMING_SNAKE_CASE           =     self.num_hidden_layers + 2


				def 							A (    self      )   ->   Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     ids_tensor([self.batch_size, self.seq_length]		,     self.vocab_size      )

									SCREAMING_SNAKE_CASE           =     None
									if self.use_input_mask:
														SCREAMING_SNAKE_CASE           =     ids_tensor([self.batch_size, self.seq_length]		,     vocab_size=2      )

									SCREAMING_SNAKE_CASE           =     None
									if self.use_token_type_ids:
														SCREAMING_SNAKE_CASE           =     ids_tensor([self.batch_size, self.seq_length]		,     self.type_vocab_size      )

									SCREAMING_SNAKE_CASE           =     None
									SCREAMING_SNAKE_CASE           =     None
									SCREAMING_SNAKE_CASE           =     None
									if self.use_labels:
														SCREAMING_SNAKE_CASE           =     ids_tensor([self.batch_size]		,     self.type_sequence_label_size      )
														SCREAMING_SNAKE_CASE           =     ids_tensor([self.batch_size, self.seq_length]		,     self.num_labels      )
														SCREAMING_SNAKE_CASE           =     ids_tensor([self.batch_size]		,     self.num_choices      )
														SCREAMING_SNAKE_CASE           =     ids_tensor([self.batch_size, self.seq_length]		,     1      )

									SCREAMING_SNAKE_CASE           =     self.get_config()

									return (
									    config,
									    input_ids,
									    token_type_ids,
									    input_mask,
									    sequence_labels,
									    token_labels,
									    choice_labels,
									    fake_token_labels,
									)


				def 							A (    self      )   ->   List[Any]:
									"""simple docstring"""







									return FunnelConfig(
									    vocab_size=self.vocab_size		,     block_sizes=self.block_sizes		,     num_decoder_layers=self.num_decoder_layers		,     d_model=self.d_model		,     n_head=self.n_head		,     d_head=self.d_head		,     d_inner=self.d_inner		,     hidden_act=self.hidden_act		,     hidden_dropout=self.hidden_dropout		,     attention_dropout=self.attention_dropout		,     activation_dropout=self.activation_dropout		,     max_position_embeddings=self.max_position_embeddings		,     type_vocab_size=self.type_vocab_size		,     initializer_std=self.initializer_std		,     )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     )   ->   int:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     FunnelModel(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     attention_mask=lowercase__		,     token_type_ids=lowercase__      )
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     token_type_ids=lowercase__      )
									SCREAMING_SNAKE_CASE           =     model(lowercase__      )
									self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, self.seq_length, self.d_model)      )

									SCREAMING_SNAKE_CASE           =     False
									SCREAMING_SNAKE_CASE           =     model(lowercase__      )
									self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, self.seq_length, self.d_model)      )

									SCREAMING_SNAKE_CASE           =     False
									SCREAMING_SNAKE_CASE           =     model(lowercase__      )
									self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, self.seq_length, self.d_model)      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     )   ->   int:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     FunnelBaseModel(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     attention_mask=lowercase__		,     token_type_ids=lowercase__      )
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     token_type_ids=lowercase__      )
									SCREAMING_SNAKE_CASE           =     model(lowercase__      )
									self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, 2, self.d_model)      )

									SCREAMING_SNAKE_CASE           =     False
									SCREAMING_SNAKE_CASE           =     model(lowercase__      )
									self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, 3, self.d_model)      )

									SCREAMING_SNAKE_CASE           =     False
									SCREAMING_SNAKE_CASE           =     model(lowercase__      )
									self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, 2, self.d_model)      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     )   ->   List[str]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.num_labels
									SCREAMING_SNAKE_CASE           =     FunnelForPreTraining(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     attention_mask=lowercase__		,     token_type_ids=lowercase__		,     labels=lowercase__      )
									self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.seq_length)      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     )   ->   Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     FunnelForMaskedLM(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     attention_mask=lowercase__		,     token_type_ids=lowercase__		,     labels=lowercase__      )
									self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.seq_length, self.vocab_size)      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     )   ->   Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.num_labels
									SCREAMING_SNAKE_CASE           =     FunnelForSequenceClassification(lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     attention_mask=lowercase__		,     token_type_ids=lowercase__		,     labels=lowercase__      )
									self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.num_labels)      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     )   ->   Dict:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.num_choices
									SCREAMING_SNAKE_CASE           =     FunnelForMultipleChoice(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     input_ids.unsqueeze(1      ).expand(-1		,     self.num_choices		,     -1      ).contiguous()
									SCREAMING_SNAKE_CASE           =     token_type_ids.unsqueeze(1      ).expand(-1		,     self.num_choices		,     -1      ).contiguous()
									SCREAMING_SNAKE_CASE           =     input_mask.unsqueeze(1      ).expand(-1		,     self.num_choices		,     -1      ).contiguous()
									SCREAMING_SNAKE_CASE           =     model(
									    lowercase__		,     attention_mask=lowercase__		,     token_type_ids=lowercase__		,     labels=lowercase__		,     )
									self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.num_choices)      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     )   ->   List[str]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.num_labels
									SCREAMING_SNAKE_CASE           =     FunnelForTokenClassification(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     attention_mask=lowercase__		,     token_type_ids=lowercase__		,     labels=lowercase__      )
									self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.seq_length, self.num_labels)      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     )   ->   List[str]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     FunnelForQuestionAnswering(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     model(
									    lowercase__		,     attention_mask=lowercase__		,     token_type_ids=lowercase__		,     start_positions=lowercase__		,     end_positions=lowercase__		,     )
									self.parent.assertEqual(result.start_logits.shape		,     (self.batch_size, self.seq_length)      )
									self.parent.assertEqual(result.end_logits.shape		,     (self.batch_size, self.seq_length)      )




				def 							A (    self      )   ->   Dict:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.prepare_config_and_inputs()
									(
									    (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   
									)           =     config_and_inputs
									SCREAMING_SNAKE_CASE           =     {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
									return config, inputs_dict






@require_torch
class 					UpperCamelCase     (	__magic_name__  ,				__magic_name__  ,				unittest.TestCase				):

				"""simple docstring"""



				UpperCAmelCase_    :				List[str]						  =    False
				UpperCAmelCase_    :				int						  =    False
				UpperCAmelCase_    :				int						  =    (
				    (
				        FunnelModel,
				        FunnelForMaskedLM,
				        FunnelForPreTraining,
				        FunnelForQuestionAnswering,
				        FunnelForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				UpperCAmelCase_    :				Tuple						  =    (
				    {
				        "feature-extraction": (FunnelBaseModel, FunnelModel),
				        "fill-mask": FunnelForMaskedLM,
				        "question-answering": FunnelForQuestionAnswering,
				        "text-classification": FunnelForSequenceClassification,
				        "token-classification": FunnelForTokenClassification,
				        "zero-shot": FunnelForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__=False      )   ->   Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     super()._prepare_for_class(lowercase__		,     lowercase__		,     return_labels=lowercase__      )

									if return_labels:
														if model_class in get_values(lowercase__      ):
																			SCREAMING_SNAKE_CASE           =     torch.zeros(
																			    (self.model_tester.batch_size, self.model_tester.seq_length)		,     dtype=torch.long		,     device=lowercase__      )
									return inputs_dict


				def 							A (    self      )   ->   Tuple:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     FunnelModelTester(self      )
									SCREAMING_SNAKE_CASE           =     ConfigTester(self		,     config_class=lowercase__      )


				def 							A (    self      )   ->   Dict:
									"""simple docstring"""







									self.config_tester.run_common_tests()


				def 							A (    self      )   ->   List[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowercase__      )


				def 							A (    self      )   ->   Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*lowercase__      )


				def 							A (    self      )   ->   Dict:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*lowercase__      )


				def 							A (    self      )   ->   List[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*lowercase__      )


				def 							A (    self      )   ->   Dict:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*lowercase__      )




				def 							A (    self		,     lowercase__      )   ->   int:
									"""simple docstring"""







									if hasattr(lowercase__		,     'weight'      ) and module.weight is not None:
														module.weight.data.fill_(3      )
									if hasattr(lowercase__		,     'bias'      ) and module.bias is not None:
														module.bias.data.fill_(3      )

									for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
														if hasattr(lowercase__		,     lowercase__      ) and getattr(lowercase__		,     lowercase__      ) is not None:
																			SCREAMING_SNAKE_CASE           =     getattr(lowercase__		,     lowercase__      )
																			weight.data.fill_(3      )






@require_torch
class 					UpperCamelCase     (	__magic_name__  ,				unittest.TestCase				):

				"""simple docstring"""



				UpperCAmelCase_    :				Optional[Any]						  =    False
				UpperCAmelCase_    :				Tuple						  =    False
				UpperCAmelCase_    :				Optional[int]						  =    (
				    (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
				)


				def 							A (    self      )   ->   int:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     FunnelModelTester(self		,     base=lowercase__      )
									SCREAMING_SNAKE_CASE           =     ConfigTester(self		,     config_class=lowercase__      )


				def 							A (    self      )   ->   Union[str, Any]:
									"""simple docstring"""







									self.config_tester.run_common_tests()


				def 							A (    self      )   ->   int:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_base_model(*lowercase__      )


				def 							A (    self      )   ->   Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*lowercase__      )


				def 							A (    self      )   ->   Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_multiple_choice(*lowercase__      )


				def 							A (    self      )   ->   List[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs_for_common()
									SCREAMING_SNAKE_CASE           =     True

									for model_class in self.all_model_classes:
														if model_class.__name__ == "FunnelBaseModel":
																			continue
														SCREAMING_SNAKE_CASE           =     model_class(lowercase__      )
														model.to(lowercase__      )
														model.train()
														SCREAMING_SNAKE_CASE           =     self._prepare_for_class(lowercase__		,     lowercase__		,     return_labels=lowercase__      )
														SCREAMING_SNAKE_CASE           =     model(**lowercase__      ).loss
														loss.backward()




				def 							A (    self		,     lowercase__      )   ->   Optional[int]:
									"""simple docstring"""







									if hasattr(lowercase__		,     'weight'      ) and module.weight is not None:
														module.weight.data.fill_(3      )
									if hasattr(lowercase__		,     'bias'      ) and module.bias is not None:
														module.bias.data.fill_(3      )

									for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
														if hasattr(lowercase__		,     lowercase__      ) and getattr(lowercase__		,     lowercase__      ) is not None:
																			SCREAMING_SNAKE_CASE           =     getattr(lowercase__		,     lowercase__      )
																			weight.data.fill_(3      )






@require_torch
@require_sentencepiece
@require_tokenizers
class 					UpperCamelCase     (	unittest.TestCase				):

				"""simple docstring"""



				def 							A (    self      )   ->   Any:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     13
									SCREAMING_SNAKE_CASE           =     7
									SCREAMING_SNAKE_CASE           =     torch.arange(0		,     batch_size * sequence_length      ).long().reshape(lowercase__		,     lowercase__      )
									SCREAMING_SNAKE_CASE           =     [0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
									SCREAMING_SNAKE_CASE           =     torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]      )

									SCREAMING_SNAKE_CASE           =     FunnelModel.from_pretrained('sgugger/funnel-random-tiny'      )
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     token_type_ids=lowercase__      )[0].abs()

									SCREAMING_SNAKE_CASE           =     torch.tensor(2_344.8_352      )
									SCREAMING_SNAKE_CASE           =     torch.tensor(0.8_052      )
									self.assertTrue(torch.allclose(output.sum()		,     lowercase__		,     atol=1E-4      )      )
									self.assertTrue(torch.allclose(output.mean()		,     lowercase__		,     atol=1E-4      )      )

									SCREAMING_SNAKE_CASE           =     torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]      )
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     attention_mask=lowercase__		,     token_type_ids=lowercase__      )[0].abs()

									SCREAMING_SNAKE_CASE           =     torch.tensor(2_343.8_425      )
									SCREAMING_SNAKE_CASE           =     torch.tensor(0.8_049      )
									self.assertTrue(torch.allclose(output.sum()		,     lowercase__		,     atol=1E-4      )      )
									self.assertTrue(torch.allclose(output.mean()		,     lowercase__		,     atol=1E-4      )      )




				@slow
				def 							A (    self      )   ->   Dict:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     FunnelTokenizer.from_pretrained('huggingface/funnel-small'      )
									SCREAMING_SNAKE_CASE           =     FunnelModel.from_pretrained('huggingface/funnel-small'      )
									SCREAMING_SNAKE_CASE           =     tokenizer('Hello! I am the Funnel Transformer model.'		,     return_tensors='pt'      )
									SCREAMING_SNAKE_CASE           =     model(**lowercase__      )[0]

									SCREAMING_SNAKE_CASE           =     torch.tensor(235.7_246      )
									SCREAMING_SNAKE_CASE           =     torch.tensor(0.0_256      )
									self.assertTrue(torch.allclose(output.sum()		,     lowercase__		,     atol=1E-4      )      )
									self.assertTrue(torch.allclose(output.mean()		,     lowercase__		,     atol=1E-4      )      )


"""simple docstring"""




from dataclasses import dataclass
from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
from ..utils import BaseOutput, deprecate
from .attention import BasicTransformerBlock
from .embeddings import PatchEmbed
from .modeling_utils import ModelMixin




@dataclass
class 					UpperCamelCase     (	__magic_name__				):

				"""simple docstring"""



				UpperCAmelCase_    :				torch.FloatTensor






class 					UpperCamelCase     (	__magic_name__  ,				__magic_name__				):

				"""simple docstring"""



				@register_to_config
				def __init__(    self		,     lowercase__ = 16		,     lowercase__ = 88		,     lowercase__ = None		,     lowercase__ = None		,     lowercase__ = 1		,     lowercase__ = 0.0		,     lowercase__ = 32		,     lowercase__ = None		,     lowercase__ = False		,     lowercase__ = None		,     lowercase__ = None		,     lowercase__ = None		,     lowercase__ = "geglu"		,     lowercase__ = None		,     lowercase__ = False		,     lowercase__ = False		,     lowercase__ = False		,     lowercase__ = "layer_norm"		,     lowercase__ = True		,     )   ->   List[Any]:
									"""simple docstring"""







									super().__init__()
									SCREAMING_SNAKE_CASE           =     use_linear_projection
									SCREAMING_SNAKE_CASE           =     num_attention_heads
									SCREAMING_SNAKE_CASE           =     attention_head_dim
									SCREAMING_SNAKE_CASE           =     num_attention_heads * attention_head_dim

									# 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
									# Define whether input is continuous or discrete depending on configuration
									SCREAMING_SNAKE_CASE           =     (in_channels is not None) and (patch_size is None)
									SCREAMING_SNAKE_CASE           =     num_vector_embeds is not None
									SCREAMING_SNAKE_CASE           =     in_channels is not None and patch_size is not None

									if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
														SCREAMING_SNAKE_CASE           =     (
														    f'''The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or'''
														    ' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'
														    ' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'
														    ' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'
														    ' would be very nice if you could open a Pull request for the `transformer/config.json` file'
														)
														deprecate('norm_type!=num_embeds_ada_norm'		,     '1.0.0'		,     lowercase__		,     standard_warn=lowercase__      )
														SCREAMING_SNAKE_CASE           =     'ada_norm'

									if self.is_input_continuous and self.is_input_vectorized:
														raise ValueError(
														    f'''Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make'''
														    ' sure that either `in_channels` or `num_vector_embeds` is None.'      )
									elif self.is_input_vectorized and self.is_input_patches:
														raise ValueError(
														    f'''Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make'''
														    ' sure that either `num_vector_embeds` or `num_patches` is None.'      )
									elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
														raise ValueError(
														    f'''Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:'''
														    f''' {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None.'''      )

									# 2. Define input layers
									if self.is_input_continuous:
														SCREAMING_SNAKE_CASE           =     in_channels

														SCREAMING_SNAKE_CASE           =     torch.nn.GroupNorm(num_groups=lowercase__		,     num_channels=lowercase__		,     eps=1E-6		,     affine=lowercase__      )
														if use_linear_projection:
																			SCREAMING_SNAKE_CASE           =     nn.Linear(lowercase__		,     lowercase__      )
														else:
																			SCREAMING_SNAKE_CASE           =     nn.Convad(lowercase__		,     lowercase__		,     kernel_size=1		,     stride=1		,     padding=0      )
									elif self.is_input_vectorized:
														assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
														assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

														SCREAMING_SNAKE_CASE           =     sample_size
														SCREAMING_SNAKE_CASE           =     sample_size
														SCREAMING_SNAKE_CASE           =     num_vector_embeds
														SCREAMING_SNAKE_CASE           =     self.height * self.width

														SCREAMING_SNAKE_CASE           =     ImagePositionalEmbeddings(
														    num_embed=lowercase__		,     embed_dim=lowercase__		,     height=self.height		,     width=self.width      )
									elif self.is_input_patches:
														assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

														SCREAMING_SNAKE_CASE           =     sample_size
														SCREAMING_SNAKE_CASE           =     sample_size

														SCREAMING_SNAKE_CASE           =     patch_size
														SCREAMING_SNAKE_CASE           =     PatchEmbed(
														    height=lowercase__		,     width=lowercase__		,     patch_size=lowercase__		,     in_channels=lowercase__		,     embed_dim=lowercase__		,     )

									# 3. Define transformers blocks
									SCREAMING_SNAKE_CASE           =     nn.ModuleList(
									    [
									        BasicTransformerBlock(
									            lowercase__		,     lowercase__		,     lowercase__		,     dropout=lowercase__		,     cross_attention_dim=lowercase__		,     activation_fn=lowercase__		,     num_embeds_ada_norm=lowercase__		,     attention_bias=lowercase__		,     only_cross_attention=lowercase__		,     upcast_attention=lowercase__		,     norm_type=lowercase__		,     norm_elementwise_affine=lowercase__		,     )
									        for d in range(lowercase__      )
									    ]      )

									# 4. Define output layers
									SCREAMING_SNAKE_CASE           =     in_channels if out_channels is None else out_channels
									if self.is_input_continuous:
														# TODO: should use out_channels for continuous projections
														if use_linear_projection:
																			SCREAMING_SNAKE_CASE           =     nn.Linear(lowercase__		,     lowercase__      )
														else:
																			SCREAMING_SNAKE_CASE           =     nn.Convad(lowercase__		,     lowercase__		,     kernel_size=1		,     stride=1		,     padding=0      )
									elif self.is_input_vectorized:
														SCREAMING_SNAKE_CASE           =     nn.LayerNorm(lowercase__      )
														SCREAMING_SNAKE_CASE           =     nn.Linear(lowercase__		,     self.num_vector_embeds - 1      )
									elif self.is_input_patches:
														SCREAMING_SNAKE_CASE           =     nn.LayerNorm(lowercase__		,     elementwise_affine=lowercase__		,     eps=1E-6      )
														SCREAMING_SNAKE_CASE           =     nn.Linear(lowercase__		,     2 * inner_dim      )
														SCREAMING_SNAKE_CASE           =     nn.Linear(lowercase__		,     patch_size * patch_size * self.out_channels      )




				def 							A (    self		,     lowercase__		,     lowercase__ = None		,     lowercase__ = None		,     lowercase__ = None		,     lowercase__ = None		,     lowercase__ = None		,     lowercase__ = None		,     lowercase__ = True		,     )   ->   List[Any]:
									"""simple docstring"""







									if attention_mask is not None and attention_mask.ndim == 2:
														# assume that mask is expressed as:
														#   (1 = keep,      0 = discard)
														# convert mask into a bias that can be added to attention scores:
														#       (keep = +0,     discard = -10000.0)
														SCREAMING_SNAKE_CASE           =     (1 - attention_mask.to(hidden_states.dtype      )) * -10_000.0
														SCREAMING_SNAKE_CASE           =     attention_mask.unsqueeze(1      )

									# convert encoder_attention_mask to a bias the same way we do for attention_mask
									if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
														SCREAMING_SNAKE_CASE           =     (1 - encoder_attention_mask.to(hidden_states.dtype      )) * -10_000.0
														SCREAMING_SNAKE_CASE           =     encoder_attention_mask.unsqueeze(1      )

									# 1. Input
									if self.is_input_continuous:
														SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     hidden_states.shape
														SCREAMING_SNAKE_CASE           =     hidden_states

														SCREAMING_SNAKE_CASE           =     self.norm(lowercase__      )
														if not self.use_linear_projection:
																			SCREAMING_SNAKE_CASE           =     self.proj_in(lowercase__      )
																			SCREAMING_SNAKE_CASE           =     hidden_states.shape[1]
																			SCREAMING_SNAKE_CASE           =     hidden_states.permute(0		,     2		,     3		,     1      ).reshape(lowercase__		,     height * width		,     lowercase__      )
														else:
																			SCREAMING_SNAKE_CASE           =     hidden_states.shape[1]
																			SCREAMING_SNAKE_CASE           =     hidden_states.permute(0		,     2		,     3		,     1      ).reshape(lowercase__		,     height * width		,     lowercase__      )
																			SCREAMING_SNAKE_CASE           =     self.proj_in(lowercase__      )
									elif self.is_input_vectorized:
														SCREAMING_SNAKE_CASE           =     self.latent_image_embedding(lowercase__      )
									elif self.is_input_patches:
														SCREAMING_SNAKE_CASE           =     self.pos_embed(lowercase__      )

									# 2. Blocks
									for block in self.transformer_blocks:
														SCREAMING_SNAKE_CASE           =     block(
														    lowercase__		,     attention_mask=lowercase__		,     encoder_hidden_states=lowercase__		,     encoder_attention_mask=lowercase__		,     timestep=lowercase__		,     cross_attention_kwargs=lowercase__		,     class_labels=lowercase__		,     )

									# 3. Output
									if self.is_input_continuous:
														if not self.use_linear_projection:
																			SCREAMING_SNAKE_CASE           =     hidden_states.reshape(lowercase__		,     lowercase__		,     lowercase__		,     lowercase__      ).permute(0		,     3		,     1		,     2      ).contiguous()
																			SCREAMING_SNAKE_CASE           =     self.proj_out(lowercase__      )
														else:
																			SCREAMING_SNAKE_CASE           =     self.proj_out(lowercase__      )
																			SCREAMING_SNAKE_CASE           =     hidden_states.reshape(lowercase__		,     lowercase__		,     lowercase__		,     lowercase__      ).permute(0		,     3		,     1		,     2      ).contiguous()

														SCREAMING_SNAKE_CASE           =     hidden_states + residual
									elif self.is_input_vectorized:
														SCREAMING_SNAKE_CASE           =     self.norm_out(lowercase__      )
														SCREAMING_SNAKE_CASE           =     self.out(lowercase__      )
														# (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
														SCREAMING_SNAKE_CASE           =     logits.permute(0		,     2		,     1      )

														# log(p(x_0))
														SCREAMING_SNAKE_CASE           =     F.log_softmax(logits.double()		,     dim=1      ).float()
									elif self.is_input_patches:
														# TODO: cleanup!
														SCREAMING_SNAKE_CASE           =     self.transformer_blocks[0].norma.emb(
														    lowercase__		,     lowercase__		,     hidden_dtype=hidden_states.dtype      )
														SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     self.proj_out_a(F.silu(lowercase__      )      ).chunk(2		,     dim=1      )
														SCREAMING_SNAKE_CASE           =     self.norm_out(lowercase__      ) * (1 + scale[:, None]) + shift[:, None]
														SCREAMING_SNAKE_CASE           =     self.proj_out_a(lowercase__      )

														# unpatchify
														SCREAMING_SNAKE_CASE           =     SCREAMING_SNAKE_CASE           =     int(hidden_states.shape[1] ** 0.5      )
														SCREAMING_SNAKE_CASE           =     hidden_states.reshape(
														    shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)      )
														SCREAMING_SNAKE_CASE           =     torch.einsum('nhwpqc->nchpwq'		,     lowercase__      )
														SCREAMING_SNAKE_CASE           =     hidden_states.reshape(
														    shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)      )

									if not return_dict:
														return (output,)

									return TransformeraDModelOutput(sample=lowercase__      )
