'''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    a__				:






     def __init__(			self		:		str  ,    a		:		List[str]  ,    ):





         """simple docstring"""



         __lowerCamelCase          =  parent
         __lowerCamelCase          =  13
         __lowerCamelCase          =  7
         __lowerCamelCase          =  True
         __lowerCamelCase          =  True
         __lowerCamelCase          =  False
         __lowerCamelCase          =  True
         __lowerCamelCase          =  99
         __lowerCamelCase          =  32
         __lowerCamelCase          =  2
         __lowerCamelCase          =  4
         __lowerCamelCase          =  37
         __lowerCamelCase          =  """gelu"""
         __lowerCamelCase          =  0.1
         __lowerCamelCase          =  0.1
         __lowerCamelCase          =  5_12
         __lowerCamelCase          =  16
         __lowerCamelCase          =  2
         __lowerCamelCase          =  0.02
         __lowerCamelCase          =  3
         __lowerCamelCase          =  4
         __lowerCamelCase          =  None






     def        SCREAMING_SNAKE_CASE__       (			self		:		Dict   ):





         """simple docstring"""



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

         __lowerCamelCase          =  None
         if self.use_input_mask:
             __lowerCamelCase          =  random_attention_mask([self.batch_size, self.seq_length]   )

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

         __lowerCamelCase          =  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        SCREAMING_SNAKE_CASE__       (			self		:		str  ,    a		:		Tuple  ,    a		:		Tuple  ,    a		:		Dict  ,    a		:		Any  ,    a		:		Any  ,    a		:		List[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  TFDistilBertModel(config=SCREAMING_SNAKE_CASE_   )
         __lowerCamelCase          =  {"""input_ids""": input_ids, """attention_mask""": input_mask}

         __lowerCamelCase          =  model(SCREAMING_SNAKE_CASE_   )

         __lowerCamelCase          =  [input_ids, input_mask]

         __lowerCamelCase          =  model(SCREAMING_SNAKE_CASE_   )

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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]  ,    a		:		int  ,    a		:		Optional[Any]  ,    a		:		Optional[int]  ,    a		:		Optional[int]  ,    a		:		Tuple  ,    a		:		Any   ):





         """simple docstring"""



         __lowerCamelCase          =  TFDistilBertForMaskedLM(config=SCREAMING_SNAKE_CASE_   )
         __lowerCamelCase          =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
         __lowerCamelCase          =  model(SCREAMING_SNAKE_CASE_   )
         self.parent.assertEqual(result.logits.shape  ,    (self.batch_size, self.seq_length, self.vocab_size)   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[Any]  ,    a		:		Optional[int]  ,    a		:		Optional[int]  ,    a		:		Tuple  ,    a		:		Optional[Any]  ,    a		:		Union[str, Any]  ,    a		:		Any   ):





         """simple docstring"""



         __lowerCamelCase          =  TFDistilBertForQuestionAnswering(config=SCREAMING_SNAKE_CASE_   )
         __lowerCamelCase          =  {
             """input_ids""": input_ids,
             """attention_mask""": input_mask,
         }
         __lowerCamelCase          =  model(SCREAMING_SNAKE_CASE_   )
         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        SCREAMING_SNAKE_CASE__       (			self		:		str  ,    a		:		Optional[int]  ,    a		:		Optional[int]  ,    a		:		Union[str, Any]  ,    a		:		Any  ,    a		:		List[str]  ,    a		:		Optional[int]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.num_labels
         __lowerCamelCase          =  TFDistilBertForSequenceClassification(SCREAMING_SNAKE_CASE_   )
         __lowerCamelCase          =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
         __lowerCamelCase          =  model(SCREAMING_SNAKE_CASE_   )
         self.parent.assertEqual(result.logits.shape  ,    (self.batch_size, self.num_labels)   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple  ,    a		:		Any  ,    a		:		List[Any]  ,    a		:		List[str]  ,    a		:		List[str]  ,    a		:		Any  ,    a		:		Union[str, Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.num_choices
         __lowerCamelCase          =  TFDistilBertForMultipleChoice(SCREAMING_SNAKE_CASE_   )
         __lowerCamelCase          =  tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE_  ,    1   )  ,    (1, self.num_choices, 1)   )
         __lowerCamelCase          =  tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE_  ,    1   )  ,    (1, self.num_choices, 1)   )
         __lowerCamelCase          =  {
             """input_ids""": multiple_choice_inputs_ids,
             """attention_mask""": multiple_choice_input_mask,
         }
         __lowerCamelCase          =  model(SCREAMING_SNAKE_CASE_   )
         self.parent.assertEqual(result.logits.shape  ,    (self.batch_size, self.num_choices)   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[str]  ,    a		:		Optional[int]  ,    a		:		Dict  ,    a		:		Optional[int]  ,    a		:		Optional[Any]  ,    a		:		List[Any]  ,    a		:		Optional[int]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.num_labels
         __lowerCamelCase          =  TFDistilBertForTokenClassification(SCREAMING_SNAKE_CASE_   )
         __lowerCamelCase          =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
         __lowerCamelCase          =  model(SCREAMING_SNAKE_CASE_   )
         self.parent.assertEqual(result.logits.shape  ,    (self.batch_size, self.seq_length, self.num_labels)   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple   ):





         """simple docstring"""



         __lowerCamelCase          =  self.prepare_config_and_inputs()
         (__lowerCamelCase)          =  config_and_inputs
         __lowerCamelCase          =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
         return config, inputs_dict




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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple   ):





         """simple docstring"""



         __lowerCamelCase          =  TFDistilBertModelTester(self   )
         __lowerCamelCase          =  ConfigTester(self  ,    config_class=SCREAMING_SNAKE_CASE_  ,    dim=37   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple   ):





         """simple docstring"""



         self.config_tester.run_common_tests()






     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple   ):





         """simple docstring"""



         __lowerCamelCase          =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_model(*SCREAMING_SNAKE_CASE_   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[str]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_masked_lm(*SCREAMING_SNAKE_CASE_   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         __lowerCamelCase          =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_question_answering(*SCREAMING_SNAKE_CASE_   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_sequence_classification(*SCREAMING_SNAKE_CASE_   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



         __lowerCamelCase          =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_multiple_choice(*SCREAMING_SNAKE_CASE_   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Union[str, Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_token_classification(*SCREAMING_SNAKE_CASE_   )






     @slow
     def        SCREAMING_SNAKE_CASE__       (			self		:		List[Any]   ):





         """simple docstring"""



         for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]   ):
             __lowerCamelCase          =  TFDistilBertModel.from_pretrained(SCREAMING_SNAKE_CASE_   )
             self.assertIsNotNone(SCREAMING_SNAKE_CASE_   )




@require_tf
class    a__				(       unittest.TestCase			):






     @slow
     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         __lowerCamelCase          =  TFDistilBertModel.from_pretrained('''distilbert-base-uncased'''   )
         __lowerCamelCase          =  tf.constant([[0, 1, 2, 3, 4, 5]]   )
         __lowerCamelCase          =  model(SCREAMING_SNAKE_CASE_   )[0]

         __lowerCamelCase          =  [1, 6, 7_68]
         self.assertEqual(output.shape  ,    SCREAMING_SNAKE_CASE_   )

         __lowerCamelCase          =  tf.constant(
             [
                 [
                     [0.19_26_18_85, -0.13_73_29_55, 0.4_11_97_99],
                     [0.22_15_01_56, -0.07_42_26_61, 0.39_03_72_04],
                     [0.22_75_60_18, -0.0_89_64_14, 0.3_70_14_67],
                 ]
             ]   )
         tf.debugging.assert_near(output[:, :3, :3]  ,    SCREAMING_SNAKE_CASE_  ,    atol=1e-4   )







'''simple docstring'''




import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin




class    a__				(       UpperCAmelCase__							,       UpperCAmelCase__			):






     @register_to_config
     def __init__(			self		:		Any  ,    *,
         a		:		int = 4  ,    a		:		int = 7_68  ,    a		:		int  ,    a		:		Optional[int]  ,    ):





         """simple docstring"""



         super().__init__()

         __lowerCamelCase          =  nn.Parameter(torch.zeros(a   )   )

         # parameters for additional clip time embeddings
         __lowerCamelCase          =  nn.Linear(a  ,    a   )
         __lowerCamelCase          =  nn.Linear(a  ,    a   )

         # parameters for encoder hidden states
         __lowerCamelCase          =  clip_extra_context_tokens
         __lowerCamelCase          =  nn.Linear(
             a  ,    self.clip_extra_context_tokens * cross_attention_dim   )
         __lowerCamelCase          =  nn.Linear(a  ,    a   )
         __lowerCamelCase          =  nn.LayerNorm(a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		int  ,    *, a		:		Union[str, Any]  ,    a		:		Any  ,    a		:		str  ,    a		:		Any   ):





         """simple docstring"""



         if do_classifier_free_guidance:
             # Add the classifier free guidance embeddings to the image embeddings
             __lowerCamelCase          =  image_embeddings.shape[0]
             __lowerCamelCase          =  self.learned_classifier_free_guidance_embeddings.unsqueeze(0   )
             __lowerCamelCase          =  classifier_free_guidance_embeddings.expand(
                 a  ,    -1   )
             __lowerCamelCase          =  torch.cat([classifier_free_guidance_embeddings, image_embeddings]  ,    dim=0   )

         # The image embeddings batch size and the text embeddings batch size are equal
         assert image_embeddings.shape[0] == prompt_embeds.shape[0]

         __lowerCamelCase          =  prompt_embeds.shape[0]

         # "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
         # adding CLIP embeddings to the existing timestep embedding, ...
         __lowerCamelCase          =  self.embedding_proj(a   )
         __lowerCamelCase          =  self.clip_image_embeddings_project_to_time_embeddings(a   )
         __lowerCamelCase          =  time_projected_image_embeddings + time_projected_prompt_embeds

         # ... and by projecting CLIP embeddings into four
         # extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
         __lowerCamelCase          =  self.clip_extra_context_tokens_proj(a   )
         __lowerCamelCase          =  clip_extra_context_tokens.reshape(a  ,    -1  ,    self.clip_extra_context_tokens   )
         __lowerCamelCase          =  clip_extra_context_tokens.permute(0  ,    2  ,    1   )

         __lowerCamelCase          =  self.encoder_hidden_states_proj(a   )
         __lowerCamelCase          =  self.text_encoder_hidden_states_norm(a   )
         __lowerCamelCase          =  torch.cat([clip_extra_context_tokens, text_encoder_hidden_states]  ,    dim=1   )

         return text_encoder_hidden_states, additive_clip_time_embeddings
