





"""simple docstring"""




from __future__ import annotations

import unittest

from transformers import DistilBertConfig, 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.models.distilbert.modeling_tf_distilbert import (
		    TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFDistilBertForMaskedLM,
		    TFDistilBertForMultipleChoice,
		    TFDistilBertForQuestionAnswering,
		    TFDistilBertForSequenceClassification,
		    TFDistilBertForTokenClassification,
		    TFDistilBertModel,
		)


class  __UpperCamelCase       :






							def __init__(     self   ,      lowerCAmelCase__   ,      )			->				List[str]:
								a   :	Tuple			=       parent
								a   :	str			=       13
								a   :	Any			=       7
								a   :	List[str]			=       True
								a   :	Optional[Any]			=       True
								a   :	Any			=       False
								a   :	Any			=       True
								a   :	Any			=       99
								a   :	Any			=       32
								a   :	Any			=       2
								a   :	List[Any]			=       4
								a   :	Union[str, Any]			=       37
								a   :	Dict			=       "gelu"
								a   :	List[Any]			=       0.1
								a   :	int			=       0.1
								a   :	Optional[int]			=       512
								a   :	str			=       16
								a   :	Dict			=       2
								a   :	Optional[int]			=       0.02
								a   :	List[Any]			=       3
								a   :	Union[str, Any]			=       4
								a   :	Optional[int]			=       None






							def 					__a     (     self   )			->				Dict:
								a   :	Optional[Any]			=       ids_tensor([self.batch_size, self.seq_length]   ,      self.vocab_size   )

								a   :	Any			=       None
								if self.use_input_mask:
									a   :	Optional[Any]			=       random_attention_mask([self.batch_size, self.seq_length]   )

								a   :	List[str]			=       None
								a   :	List[Any]			=       None
								a   :	Dict			=       None
								if self.use_labels:
									a   :	Any			=       ids_tensor([self.batch_size]   ,      self.type_sequence_label_size   )
									a   :	Optional[int]			=       ids_tensor([self.batch_size, self.seq_length]   ,      self.num_labels   )
									a   :	Union[str, Any]			=       ids_tensor([self.batch_size]   ,      self.num_choices   )

								a   :	Union[str, Any]			=       DistilBertConfig(
								    vocab_size=self.vocab_size   ,      dim=self.hidden_size   ,      n_layers=self.num_hidden_layers   ,      n_heads=self.num_attention_heads   ,      hidden_dim=self.intermediate_size   ,      hidden_act=self.hidden_act   ,      dropout=self.hidden_dropout_prob   ,      attention_dropout=self.attention_probs_dropout_prob   ,      max_position_embeddings=self.max_position_embeddings   ,      initializer_range=self.initializer_range   ,      )

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






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				Optional[int]:
								a   :	List[Any]			=       TFDistilBertModel(config=lowerCAmelCase__   )
								a   :	List[Any]			=       {"input_ids": input_ids, "attention_mask": input_mask}

								a   :	Tuple			=       model(lowerCAmelCase__   )

								a   :	List[Any]			=       [input_ids, input_mask]

								a   :	Tuple			=       model(lowerCAmelCase__   )

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






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				int:
								a   :	List[Any]			=       TFDistilBertForMaskedLM(config=lowerCAmelCase__   )
								a   :	List[str]			=       {"input_ids": input_ids, "attention_mask": input_mask}
								a   :	Union[str, Any]			=       model(lowerCAmelCase__   )
								self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.vocab_size)   )






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				List[str]:
								a   :	Any			=       TFDistilBertForQuestionAnswering(config=lowerCAmelCase__   )
								a   :	int			=       {
								    "input_ids": input_ids,
								    "attention_mask": input_mask,
								}
								a   :	Tuple			=       model(lowerCAmelCase__   )
								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   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				str:
								a   :	Dict			=       self.num_labels
								a   :	Optional[int]			=       TFDistilBertForSequenceClassification(lowerCAmelCase__   )
								a   :	List[Any]			=       {"input_ids": input_ids, "attention_mask": input_mask}
								a   :	Any			=       model(lowerCAmelCase__   )
								self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.num_labels)   )






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				List[str]:
								a   :	Dict			=       self.num_choices
								a   :	Union[str, Any]			=       TFDistilBertForMultipleChoice(lowerCAmelCase__   )
								a   :	Union[str, Any]			=       tf.tile(tf.expand_dims(lowerCAmelCase__   ,      1   )   ,      (1, self.num_choices, 1)   )
								a   :	List[str]			=       tf.tile(tf.expand_dims(lowerCAmelCase__   ,      1   )   ,      (1, self.num_choices, 1)   )
								a   :	List[Any]			=       {
								    "input_ids": multiple_choice_inputs_ids,
								    "attention_mask": multiple_choice_input_mask,
								}
								a   :	List[str]			=       model(lowerCAmelCase__   )
								self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.num_choices)   )






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				Any:
								a   :	Dict			=       self.num_labels
								a   :	Tuple			=       TFDistilBertForTokenClassification(lowerCAmelCase__   )
								a   :	Tuple			=       {"input_ids": input_ids, "attention_mask": input_mask}
								a   :	Optional[int]			=       model(lowerCAmelCase__   )
								self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.num_labels)   )






							def 					__a     (     self   )			->				Optional[int]:
								a   :	List[Any]			=       self.prepare_config_and_inputs()
								((a), (a), (a), (a), (a), (a))   :	Union[str, Any]			=       config_and_inputs
								a   :	List[str]			=       {"input_ids": input_ids, "attention_mask": input_mask}
								return config, inputs_dict


@require_tf
class  __UpperCamelCase       (		a__	, a__	, unittest.TestCase		):
							lowerCamelCase     :      Optional[Any]         =(
							    (
							        TFDistilBertModel,
							        TFDistilBertForMaskedLM,
							        TFDistilBertForQuestionAnswering,
							        TFDistilBertForSequenceClassification,
							        TFDistilBertForTokenClassification,
							        TFDistilBertForMultipleChoice,
							    )
							    if is_tf_available()
							    else None
							)
							lowerCamelCase     :      Any         =(
							    {
							        """feature-extraction""": TFDistilBertModel,
							        """fill-mask""": TFDistilBertForMaskedLM,
							        """question-answering""": TFDistilBertForQuestionAnswering,
							        """text-classification""": TFDistilBertForSequenceClassification,
							        """token-classification""": TFDistilBertForTokenClassification,
							        """zero-shot""": TFDistilBertForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							lowerCamelCase     :      List[str]         =False
							lowerCamelCase     :      Optional[int]         =False






							def 					__a     (     self   )			->				str:
								a   :	Optional[Any]			=       TFDistilBertModelTester(self   )
								a   :	List[str]			=       ConfigTester(self   ,      config_class=lowerCAmelCase__   ,      dim=37   )






							def 					__a     (     self   )			->				Any:
								self.config_tester.run_common_tests()






							def 					__a     (     self   )			->				Union[str, Any]:
								a   :	Optional[int]			=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_model(*lowerCAmelCase__   )






							def 					__a     (     self   )			->				Dict:
								a   :	Optional[int]			=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_masked_lm(*lowerCAmelCase__   )






							def 					__a     (     self   )			->				List[str]:
								a   :	str			=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_question_answering(*lowerCAmelCase__   )






							def 					__a     (     self   )			->				Tuple:
								a   :	Dict			=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowerCAmelCase__   )






							def 					__a     (     self   )			->				Dict:
								a   :	Optional[int]			=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowerCAmelCase__   )






							def 					__a     (     self   )			->				Optional[Any]:
								a   :	Any			=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_distilbert_for_token_classification(*lowerCAmelCase__   )






							@slow
							def 					__a     (     self   )			->				Union[str, Any]:
								for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]   ):
									a   :	Tuple			=       TFDistilBertModel.from_pretrained(lowerCAmelCase__   )
									self.assertIsNotNone(lowerCAmelCase__   )


@require_tf
class  __UpperCamelCase       (		unittest.TestCase		):






							@slow
							def 					__a     (     self   )			->				Tuple:
								a   :	List[Any]			=       TFDistilBertModel.from_pretrained("distilbert-base-uncased"   )
								a   :	Dict			=       tf.constant([[0, 1, 2, 3, 4, 5]]   )
								a   :	int			=       model(lowerCAmelCase__   )[0]

								a   :	int			=       [1, 6, 768]
								self.assertEqual(output.shape   ,      lowerCAmelCase__   )

								a   :	List[Any]			=       tf.constant(
								    [
								        [
								            [0.19_261_885, -0.13_732_955, 0.4_119_799],
								            [0.22_150_156, -0.07_422_661, 0.39_037_204],
								            [0.22_756_018, -0.0_896_414, 0.3_701_467],
								        ]
								    ]   )
								tf.debugging.assert_near(output[:, :3, :3]   ,      lowerCAmelCase__   ,      atol=1E-4   )









"""simple docstring"""




import flax.linen as nn
import jax
import jax.numpy as jnp


class  __UpperCamelCase       (		nn.Module		):
							lowerCamelCase     :      int
							lowerCamelCase     :      jnp.dtype         =jnp.floataa






							def 					__a     (     self   )			->				Tuple:
								a   :	str			=       nn.Conv(
								    self.out_channels   ,      kernel_size=(3, 3)   ,      strides=(1, 1)   ,      padding=((1, 1), (1, 1))   ,      dtype=self.dtype   ,      )






							def __call__(     self   ,      lowerCAmelCase__   )			->				Optional[Any]:
								a, a, a, a   :	List[str]			=       hidden_states.shape
								a   :	List[Any]			=       jax.image.resize(
								    lowerCAmelCase__   ,      shape=(batch, height * 2, width * 2, channels)   ,      method="nearest"   ,      )
								a   :	List[str]			=       self.conv(lowerCAmelCase__   )
								return hidden_states


class  __UpperCamelCase       (		nn.Module		):
							lowerCamelCase     :      int
							lowerCamelCase     :      jnp.dtype         =jnp.floataa






							def 					__a     (     self   )			->				Dict:
								a   :	Optional[Any]			=       nn.Conv(
								    self.out_channels   ,      kernel_size=(3, 3)   ,      strides=(2, 2)   ,      padding=((1, 1), (1, 1))   ,      dtype=self.dtype   ,      )






							def __call__(     self   ,      lowerCAmelCase__   )			->				Tuple:
								# pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
								# hidden_states = jnp.pad(hidden_states, pad_width=pad)
								a   :	Tuple			=       self.conv(lowerCAmelCase__   )
								return hidden_states


class  __UpperCamelCase       (		nn.Module		):
							lowerCamelCase     :      int
							lowerCamelCase     :      int         =None
							lowerCamelCase     :      float         =0.0
							lowerCamelCase     :      bool         =None
							lowerCamelCase     :      jnp.dtype         =jnp.floataa






							def 					__a     (     self   )			->				int:
								a   :	Dict			=       self.in_channels if self.out_channels is None else self.out_channels

								a   :	Union[str, Any]			=       nn.GroupNorm(num_groups=32   ,      epsilon=1E-5   )
								a   :	List[Any]			=       nn.Conv(
								    lowerCAmelCase__   ,      kernel_size=(3, 3)   ,      strides=(1, 1)   ,      padding=((1, 1), (1, 1))   ,      dtype=self.dtype   ,      )

								a   :	List[Any]			=       nn.Dense(lowerCAmelCase__   ,      dtype=self.dtype   )

								a   :	Union[str, Any]			=       nn.GroupNorm(num_groups=32   ,      epsilon=1E-5   )
								a   :	Optional[int]			=       nn.Dropout(self.dropout_prob   )
								a   :	Dict			=       nn.Conv(
								    lowerCAmelCase__   ,      kernel_size=(3, 3)   ,      strides=(1, 1)   ,      padding=((1, 1), (1, 1))   ,      dtype=self.dtype   ,      )

								a   :	Union[str, Any]			=       self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

								a   :	List[str]			=       None
								if use_nin_shortcut:
									a   :	Optional[Any]			=       nn.Conv(
									    lowerCAmelCase__   ,      kernel_size=(1, 1)   ,      strides=(1, 1)   ,      padding="VALID"   ,      dtype=self.dtype   ,      )






							def __call__(     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__=True   )			->				str:
								a   :	int			=       hidden_states
								a   :	Tuple			=       self.norma(lowerCAmelCase__   )
								a   :	Any			=       nn.swish(lowerCAmelCase__   )
								a   :	int			=       self.conva(lowerCAmelCase__   )

								a   :	int			=       self.time_emb_proj(nn.swish(lowerCAmelCase__   )   )
								a   :	Tuple			=       jnp.expand_dims(jnp.expand_dims(lowerCAmelCase__   ,      1   )   ,      1   )
								a   :	Dict			=       hidden_states + temb

								a   :	str			=       self.norma(lowerCAmelCase__   )
								a   :	List[Any]			=       nn.swish(lowerCAmelCase__   )
								a   :	List[str]			=       self.dropout(lowerCAmelCase__   ,      lowerCAmelCase__   )
								a   :	List[str]			=       self.conva(lowerCAmelCase__   )

								if self.conv_shortcut is not None:
									a   :	Tuple			=       self.conv_shortcut(lowerCAmelCase__   )

								return hidden_states + residual


