





"""simple docstring"""


from __future__ import annotations

import unittest

from transformers import DebertaConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import tensorflow as tf

		from transformers import (
		    TFDebertaForMaskedLM,
		    TFDebertaForQuestionAnswering,
		    TFDebertaForSequenceClassification,
		    TFDebertaForTokenClassification,
		    TFDebertaModel,
		)






class        lowercase__ :

			"""simple docstring"""







			def __init__(			self	,      _A	,      _A=1_3	,      _A=7	,      _A=True	,      _A=True	,      _A=True	,      _A=True	,      _A=9_9	,      _A=3_2	,      _A=2	,      _A=4	,      _A=3_7	,      _A="gelu"	,      _A=0.1	,      _A=0.1	,      _A=5_1_2	,      _A=1_6	,      _A=2	,      _A=0.02	,      _A=3	,      _A=4	,      _A=None	,      ):






									'''simple docstring'''







									UpperCamelCase							:		int							      = parent
									UpperCamelCase							:		List[Any]							      = 1_3
									UpperCamelCase							:		Union[str, Any]							      = 7
									UpperCamelCase							:		int							      = True
									UpperCamelCase							:		Any							      = True
									UpperCamelCase							:		Optional[int]							      = True
									UpperCamelCase							:		Optional[int]							      = True
									UpperCamelCase							:		Union[str, Any]							      = 9_9
									UpperCamelCase							:		List[Any]							      = 3_2
									UpperCamelCase							:		List[Any]							      = 2
									UpperCamelCase							:		List[Any]							      = 4
									UpperCamelCase							:		List[Any]							      = 3_7
									UpperCamelCase							:		int							      = """gelu"""
									UpperCamelCase							:		Optional[Any]							      = 0.1
									UpperCamelCase							:		Union[str, Any]							      = 0.1
									UpperCamelCase							:		List[str]							      = 5_1_2
									UpperCamelCase							:		Optional[Any]							      = 1_6
									UpperCamelCase							:		List[str]							      = False
									UpperCamelCase							:		int							      = -1
									UpperCamelCase							:		List[str]							      = True
									UpperCamelCase							:		Union[str, Any]							      = 2
									UpperCamelCase							:		Tuple							      = 0.02
									UpperCamelCase							:		Optional[int]							      = 3
									UpperCamelCase							:		Optional[int]							      = 4
									UpperCamelCase							:		Optional[Any]							      = None







			def      _a			(			self						):






									'''simple docstring'''







									UpperCamelCase							:		str							      = ids_tensor([self.batch_size, self.seq_length]	,      self.vocab_size						)

									UpperCamelCase							:		Tuple							      = None
									if self.use_input_mask:
															UpperCamelCase							:		Optional[int]							      = random_attention_mask([self.batch_size, self.seq_length]						)

									UpperCamelCase							:		List[str]							      = None
									if self.use_token_type_ids:
															UpperCamelCase							:		Dict							      = ids_tensor([self.batch_size, self.seq_length]	,      self.type_vocab_size						)

									UpperCamelCase							:		Tuple							      = None
									UpperCamelCase							:		int							      = None
									UpperCamelCase							:		List[str]							      = None
									if self.use_labels:
															UpperCamelCase							:		int							      = ids_tensor([self.batch_size]	,      self.type_sequence_label_size						)
															UpperCamelCase							:		List[Any]							      = ids_tensor([self.batch_size, self.seq_length]	,      self.num_labels						)

									UpperCamelCase							:		List[str]							      = DebertaConfig(
									    vocab_size=self.vocab_size	,      hidden_size=self.hidden_size	,      num_hidden_layers=self.num_hidden_layers	,      num_attention_heads=self.num_attention_heads	,      intermediate_size=self.intermediate_size	,      hidden_act=self.hidden_act	,      hidden_dropout_prob=self.hidden_dropout_prob	,      attention_probs_dropout_prob=self.attention_probs_dropout_prob	,      max_position_embeddings=self.max_position_embeddings	,      type_vocab_size=self.type_vocab_size	,      relative_attention=self.relative_attention	,      max_relative_positions=self.max_relative_positions	,      position_biased_input=self.position_biased_input	,      initializer_range=self.initializer_range	,      return_dict=_A	,      )

									return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







			def      _a			(			self	,      _A	,      _A	,      _A	,      _A	,      _A	,      _A	,      _A						):






									'''simple docstring'''







									UpperCamelCase							:		str							      = TFDebertaModel(config=_A						)
									UpperCamelCase							:		Optional[int]							      = {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}

									UpperCamelCase							:		List[str]							      = [input_ids, input_mask]
									UpperCamelCase							:		List[Any]							      = model(_A						)

									UpperCamelCase							:		Union[str, Any]							      = model(_A						)

									self.parent.assertEqual(result.last_hidden_state.shape	,      (self.batch_size, self.seq_length, self.hidden_size)						)







			def      _a			(			self	,      _A	,      _A	,      _A	,      _A	,      _A	,      _A	,      _A						):






									'''simple docstring'''







									UpperCamelCase							:		Optional[Any]							      = TFDebertaForMaskedLM(config=_A						)
									UpperCamelCase							:		Optional[int]							      = {
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									}
									UpperCamelCase							:		List[str]							      = model(_A						)
									self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.seq_length, self.vocab_size)						)







			def      _a			(			self	,      _A	,      _A	,      _A	,      _A	,      _A	,      _A	,      _A						):






									'''simple docstring'''







									UpperCamelCase							:		List[str]							      = self.num_labels
									UpperCamelCase							:		Any							      = TFDebertaForSequenceClassification(config=_A						)
									UpperCamelCase							:		Any							      = {
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									}

									UpperCamelCase							:		Dict							      = model(_A						)
									self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.num_labels)						)







			def      _a			(			self	,      _A	,      _A	,      _A	,      _A	,      _A	,      _A	,      _A						):






									'''simple docstring'''







									UpperCamelCase							:		List[str]							      = self.num_labels
									UpperCamelCase							:		Union[str, Any]							      = TFDebertaForTokenClassification(config=_A						)
									UpperCamelCase							:		List[Any]							      = {
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									}
									UpperCamelCase							:		Any							      = model(_A						)
									self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.seq_length, self.num_labels)						)







			def      _a			(			self	,      _A	,      _A	,      _A	,      _A	,      _A	,      _A	,      _A						):






									'''simple docstring'''







									UpperCamelCase							:		Any							      = TFDebertaForQuestionAnswering(config=_A						)
									UpperCamelCase							:		Tuple							      = {
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									}

									UpperCamelCase							:		List[Any]							      = model(_A						)
									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						):






									'''simple docstring'''







									UpperCamelCase							:		Optional[Any]							      = self.prepare_config_and_inputs()
									(
									    UpperCamelCase	
									)							:		Dict							      = config_and_inputs
									UpperCamelCase							:		int							      = {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
									return config, inputs_dict






@require_tf
class        lowercase__ (      __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE				,    unittest.TestCase  ):

			"""simple docstring"""
			__lowerCAmelCase    :							Union[str, Any]								=     (
			    (
			        TFDebertaModel,
			        TFDebertaForMaskedLM,
			        TFDebertaForQuestionAnswering,
			        TFDebertaForSequenceClassification,
			        TFDebertaForTokenClassification,
			    )
			    if is_tf_available()
			    else ()
			)
			__lowerCAmelCase    :							Dict								=     (
			    {
			        """feature-extraction""": TFDebertaModel,
			        """fill-mask""": TFDebertaForMaskedLM,
			        """question-answering""": TFDebertaForQuestionAnswering,
			        """text-classification""": TFDebertaForSequenceClassification,
			        """token-classification""": TFDebertaForTokenClassification,
			        """zero-shot""": TFDebertaForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)

			__lowerCAmelCase    :							List[Any]								=     False
			__lowerCAmelCase    :							str								=     False







			def      _a			(			self						):






									'''simple docstring'''







									UpperCamelCase							:		List[Any]							      = TFDebertaModelTester(self						)
									UpperCamelCase							:		Union[str, Any]							      = ConfigTester(self	,      config_class=_A	,      hidden_size=3_7						)







			def      _a			(			self						):






									'''simple docstring'''







									self.config_tester.run_common_tests()







			def      _a			(			self						):






									'''simple docstring'''







									UpperCamelCase							:		Union[str, Any]							      = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_A						)







			def      _a			(			self						):






									'''simple docstring'''







									UpperCamelCase							:		str							      = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*_A						)







			def      _a			(			self						):






									'''simple docstring'''







									UpperCamelCase							:		int							      = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*_A						)







			def      _a			(			self						):






									'''simple docstring'''







									UpperCamelCase							:		str							      = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*_A						)







			def      _a			(			self						):






									'''simple docstring'''







									UpperCamelCase							:		Optional[Any]							      = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*_A						)







			@slow
			def      _a			(			self						):






									'''simple docstring'''







									UpperCamelCase							:		Optional[Any]							      = TFDebertaModel.from_pretrained("""kamalkraj/deberta-base"""						)
									self.assertIsNotNone(_A						)






@require_tf
class        lowercase__ (      unittest.TestCase  ):

			"""simple docstring"""







			@unittest.skip(reason="""Model not available yet"""						)
			def      _a			(			self						):






									'''simple docstring'''







									pass







			@slow
			def      _a			(			self						):






									'''simple docstring'''







									UpperCamelCase							:		List[str]							      = TFDebertaModel.from_pretrained("""kamalkraj/deberta-base"""						)
									UpperCamelCase							:		Union[str, Any]							      = tf.constant([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]]						)
									UpperCamelCase							:		Tuple							      = tf.constant([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]						)
									UpperCamelCase							:		str							      = model(_A	,      attention_mask=_A						)[0]

									UpperCamelCase							:		Union[str, Any]							      = tf.constant(
									    [
									        [
									            [-0.59_85_58_96, -0.80_55_25_66, -0.8_46_21_35],
									            [1.4_48_40_25, -0.93_48_37_94, -0.80_59_30_85],
									            [0.3_12_27_41, 0.00_31_60_59, -1.4_13_13_77],
									        ]
									    ]						)
									tf.debugging.assert_near(output[:, 1:4, 1:4]	,      _A	,      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        lowercase__ (      __SCREAMING_SNAKE_CASE  ):

			"""simple docstring"""
			__lowerCAmelCase    :							torch.FloatTensor






class        lowercase__ (      __SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE  ):

			"""simple docstring"""







			@register_to_config
			def __init__(			self	,      _A = 1_6	,      _A = 8_8	,      _A = None	,      _A = None	,      _A = 1	,      _A = 0.0	,      _A = 3_2	,      _A = None	,      _A = False	,      _A = None	,      _A = None	,      _A = None	,      _A = "geglu"	,      _A = None	,      _A = False	,      _A = False	,      _A = False	,      _A = "layer_norm"	,      _A = True	,      ):






									'''simple docstring'''







									super().__init__()
									UpperCamelCase							:		int							      = use_linear_projection
									UpperCamelCase							:		str							      = num_attention_heads
									UpperCamelCase							:		Union[str, Any]							      = attention_head_dim
									UpperCamelCase							:		List[Any]							      = 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
									UpperCamelCase							:		int							      = (in_channels is not None) and (patch_size is None)
									UpperCamelCase							:		str							      = num_vector_embeds is not None
									UpperCamelCase							:		Tuple							      = in_channels is not None and patch_size is not None

									if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
															UpperCamelCase							:		Tuple							      = (
															    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"""	,      _A	,      standard_warn=_A						)
															UpperCamelCase							:		Optional[int]							      = """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:
															UpperCamelCase							:		Union[str, Any]							      = in_channels

															UpperCamelCase							:		List[str]							      = torch.nn.GroupNorm(num_groups=_A	,      num_channels=_A	,      eps=1e-6	,      affine=_A						)
															if use_linear_projection:
																					UpperCamelCase							:		Dict							      = nn.Linear(_A	,      _A						)
															else:
																					UpperCamelCase							:		int							      = nn.Convad(_A	,      _A	,      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"

															UpperCamelCase							:		Tuple							      = sample_size
															UpperCamelCase							:		str							      = sample_size
															UpperCamelCase							:		int							      = num_vector_embeds
															UpperCamelCase							:		Any							      = self.height * self.width

															UpperCamelCase							:		Any							      = ImagePositionalEmbeddings(
															    num_embed=_A	,      embed_dim=_A	,      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"

															UpperCamelCase							:		Optional[Any]							      = sample_size
															UpperCamelCase							:		int							      = sample_size

															UpperCamelCase							:		int							      = patch_size
															UpperCamelCase							:		Optional[Any]							      = PatchEmbed(
															    height=_A	,      width=_A	,      patch_size=_A	,      in_channels=_A	,      embed_dim=_A	,      )

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

									# 4. Define output layers
									UpperCamelCase							:		Union[str, Any]							      = 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:
																					UpperCamelCase							:		Optional[int]							      = nn.Linear(_A	,      _A						)
															else:
																					UpperCamelCase							:		str							      = nn.Convad(_A	,      _A	,      kernel_size=1	,      stride=1	,      padding=0						)
									elif self.is_input_vectorized:
															UpperCamelCase							:		str							      = nn.LayerNorm(_A						)
															UpperCamelCase							:		Tuple							      = nn.Linear(_A	,      self.num_vector_embeds - 1						)
									elif self.is_input_patches:
															UpperCamelCase							:		Any							      = nn.LayerNorm(_A	,      elementwise_affine=_A	,      eps=1e-6						)
															UpperCamelCase							:		Optional[Any]							      = nn.Linear(_A	,      2 * inner_dim						)
															UpperCamelCase							:		List[Any]							      = nn.Linear(_A	,      patch_size * patch_size * self.out_channels						)







			def      _a			(			self	,      _A	,      _A = None	,      _A = None	,      _A = None	,      _A = None	,      _A = None	,      _A = None	,      _A = True	,      ):






									'''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)
															UpperCamelCase							:		List[str]							      = (1 - attention_mask.to(hidden_states.dtype						)) * -1_00_00.0
															UpperCamelCase							:		Any							      = 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:
															UpperCamelCase							:		Union[str, Any]							      = (1 - encoder_attention_mask.to(hidden_states.dtype						)) * -1_00_00.0
															UpperCamelCase							:		Dict							      = encoder_attention_mask.unsqueeze(1						)

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

															UpperCamelCase							:		List[Any]							      = self.norm(_A						)
															if not self.use_linear_projection:
																					UpperCamelCase							:		Union[str, Any]							      = self.proj_in(_A						)
																					UpperCamelCase							:		Any							      = hidden_states.shape[1]
																					UpperCamelCase							:		Optional[int]							      = hidden_states.permute(0	,      2	,      3	,      1						).reshape(_A	,      height * width	,      _A						)
															else:
																					UpperCamelCase							:		Optional[int]							      = hidden_states.shape[1]
																					UpperCamelCase							:		Optional[Any]							      = hidden_states.permute(0	,      2	,      3	,      1						).reshape(_A	,      height * width	,      _A						)
																					UpperCamelCase							:		Optional[int]							      = self.proj_in(_A						)
									elif self.is_input_vectorized:
															UpperCamelCase							:		Dict							      = self.latent_image_embedding(_A						)
									elif self.is_input_patches:
															UpperCamelCase							:		str							      = self.pos_embed(_A						)

									# 2. Blocks
									for block in self.transformer_blocks:
															UpperCamelCase							:		Union[str, Any]							      = block(
															    _A	,      attention_mask=_A	,      encoder_hidden_states=_A	,      encoder_attention_mask=_A	,      timestep=_A	,      cross_attention_kwargs=_A	,      class_labels=_A	,      )

									# 3. Output
									if self.is_input_continuous:
															if not self.use_linear_projection:
																					UpperCamelCase							:		Optional[int]							      = hidden_states.reshape(_A	,      _A	,      _A	,      _A						).permute(0	,      3	,      1	,      2						).contiguous()
																					UpperCamelCase							:		Tuple							      = self.proj_out(_A						)
															else:
																					UpperCamelCase							:		List[Any]							      = self.proj_out(_A						)
																					UpperCamelCase							:		str							      = hidden_states.reshape(_A	,      _A	,      _A	,      _A						).permute(0	,      3	,      1	,      2						).contiguous()

															UpperCamelCase							:		int							      = hidden_states + residual
									elif self.is_input_vectorized:
															UpperCamelCase							:		List[str]							      = self.norm_out(_A						)
															UpperCamelCase							:		Dict							      = self.out(_A						)
															# (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
															UpperCamelCase							:		Tuple							      = logits.permute(0	,      2	,      1						)

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

															# unpatchify
															UpperCamelCase							:		Tuple							      = int(hidden_states.shape[1] ** 0.5						)
															UpperCamelCase							:		Dict							      = hidden_states.reshape(
															    shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)						)
															UpperCamelCase							:		Optional[Any]							      = torch.einsum("""nhwpqc->nchpwq"""	,      _A						)
															UpperCamelCase							:		Union[str, Any]							      = 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=_A						)



