

from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
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 TF_MODEL_FOR_PRETRAINING_MAPPING
       from transformers.models.albert.modeling_tf_albert import (
           TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
           TFAlbertForMaskedLM,
           TFAlbertForMultipleChoice,
           TFAlbertForPreTraining,
           TFAlbertForQuestionAnswering,
           TFAlbertForSequenceClassification,
           TFAlbertForTokenClassification,
           TFAlbertModel,
       )



class 					__snake_case     :

  def __init__(    self  ,						_A  ,						_A=13  ,						_A=7  ,						_A=True  ,						_A=True  ,						_A=True  ,						_A=True  ,						_A=99  ,						_A=16  ,						_A=32  ,						_A=2  ,						_A=4  ,						_A=37  ,						_A="gelu"  ,						_A=0.1  ,						_A=0.1  ,						_A=512  ,						_A=16  ,						_A=2  ,						_A=0.0_2  ,						_A=3  ,						_A=4  ,						_A=None  ,						):
      SCREAMING_SNAKE_CASE_													=  parent
      SCREAMING_SNAKE_CASE_													=  13
      SCREAMING_SNAKE_CASE_													=  7
      SCREAMING_SNAKE_CASE_													=  True
      SCREAMING_SNAKE_CASE_													=  True
      SCREAMING_SNAKE_CASE_													=  True
      SCREAMING_SNAKE_CASE_													=  True
      SCREAMING_SNAKE_CASE_													=  99
      SCREAMING_SNAKE_CASE_													=  16
      SCREAMING_SNAKE_CASE_													=  32
      SCREAMING_SNAKE_CASE_													=  2
      SCREAMING_SNAKE_CASE_													=  4
      SCREAMING_SNAKE_CASE_													=  37
      SCREAMING_SNAKE_CASE_													=  'gelu'
      SCREAMING_SNAKE_CASE_													=  0.1
      SCREAMING_SNAKE_CASE_													=  0.1
      SCREAMING_SNAKE_CASE_													=  512
      SCREAMING_SNAKE_CASE_													=  16
      SCREAMING_SNAKE_CASE_													=  2
      SCREAMING_SNAKE_CASE_													=  0.0_2
      SCREAMING_SNAKE_CASE_													=  3
      SCREAMING_SNAKE_CASE_													=  4
      SCREAMING_SNAKE_CASE_													=  None

  def   lowerCAmelCase__   (    self):
      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_													=  random_attention_mask([self.batch_size, self.seq_length])

      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_													=  AlbertConfig(
          vocab_size=self.vocab_size  ,						hidden_size=self.hidden_size  ,						embedding_size=self.embedding_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  ,						initializer_range=self.initializer_range  ,						)

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

  def   lowerCAmelCase__   (    self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A):
      SCREAMING_SNAKE_CASE_													=  TFAlbertModel(config=_A)
      # inputs = {'input_ids': input_ids,
      #           'attention_mask': input_mask,
      #           'token_type_ids': token_type_ids}
      # sequence_output, pooled_output = model(**inputs)
      SCREAMING_SNAKE_CASE_													=  {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
      SCREAMING_SNAKE_CASE_													=  model(_A)

      SCREAMING_SNAKE_CASE_													=  [input_ids, input_mask]
      SCREAMING_SNAKE_CASE_													=  model(_A)

      SCREAMING_SNAKE_CASE_													=  model(_A)

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

  def   lowerCAmelCase__   (    self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A):
      SCREAMING_SNAKE_CASE_													=  self.num_labels
      SCREAMING_SNAKE_CASE_													=  TFAlbertForPreTraining(config=_A)
      SCREAMING_SNAKE_CASE_													=  {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
      SCREAMING_SNAKE_CASE_													=  model(_A)
      self.parent.assertEqual(result.prediction_logits.shape  ,						(self.batch_size, self.seq_length, self.vocab_size))
      self.parent.assertEqual(result.sop_logits.shape  ,						(self.batch_size, self.num_labels))

  def   lowerCAmelCase__   (    self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A):
      SCREAMING_SNAKE_CASE_													=  TFAlbertForMaskedLM(config=_A)
      SCREAMING_SNAKE_CASE_													=  {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
      SCREAMING_SNAKE_CASE_													=  model(_A)
      self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length, self.vocab_size))

  def   lowerCAmelCase__   (    self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A):
      SCREAMING_SNAKE_CASE_													=  self.num_labels
      SCREAMING_SNAKE_CASE_													=  TFAlbertForSequenceClassification(config=_A)
      SCREAMING_SNAKE_CASE_													=  {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
      SCREAMING_SNAKE_CASE_													=  model(_A)
      self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_labels))

  def   lowerCAmelCase__   (    self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A):
      SCREAMING_SNAKE_CASE_													=  TFAlbertForQuestionAnswering(config=_A)
      SCREAMING_SNAKE_CASE_													=  {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
      SCREAMING_SNAKE_CASE_													=  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   lowerCAmelCase__   (    self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A):
      SCREAMING_SNAKE_CASE_													=  self.num_choices
      SCREAMING_SNAKE_CASE_													=  TFAlbertForMultipleChoice(config=_A)
      SCREAMING_SNAKE_CASE_													=  tf.tile(tf.expand_dims(_A  ,						1)  ,						(1, self.num_choices, 1))
      SCREAMING_SNAKE_CASE_													=  tf.tile(tf.expand_dims(_A  ,						1)  ,						(1, self.num_choices, 1))
      SCREAMING_SNAKE_CASE_													=  tf.tile(tf.expand_dims(_A  ,						1)  ,						(1, self.num_choices, 1))
      SCREAMING_SNAKE_CASE_													=  {
          'input_ids': multiple_choice_inputs_ids,
          'attention_mask': multiple_choice_input_mask,
          'token_type_ids': multiple_choice_token_type_ids,
      }
      SCREAMING_SNAKE_CASE_													=  model(_A)
      self.parent.assertListEqual(list(result['logits'].shape)  ,						[self.batch_size, self.num_choices])

  def   lowerCAmelCase__   (    self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A):
      SCREAMING_SNAKE_CASE_													=  self.num_labels
      SCREAMING_SNAKE_CASE_													=  TFAlbertForTokenClassification(config=_A)
      SCREAMING_SNAKE_CASE_													=  {
          'input_ids': input_ids,
          'attention_mask': input_mask,
          'token_type_ids': token_type_ids,
      }
      SCREAMING_SNAKE_CASE_													=  model(_A)
      self.parent.assertListEqual(list(result['logits'].shape)  ,						[self.batch_size, self.seq_length, self.num_labels])


  def   lowerCAmelCase__   (    self):
      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_
      )							,	
      )													=  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_tf
class 					__snake_case     (			lowerCAmelCase__   ,       lowerCAmelCase__   ,       unittest.TestCase				):
  __lowerCAmelCase   :		List[Any]         =     (
      (
          TFAlbertModel,
          TFAlbertForPreTraining,
          TFAlbertForMaskedLM,
          TFAlbertForSequenceClassification,
          TFAlbertForQuestionAnswering,
          TFAlbertForTokenClassification,
          TFAlbertForMultipleChoice,
      )
      if is_tf_available()
      else ()
  )
  __lowerCAmelCase   :		Dict         =     (
      {
          'feature-extraction': TFAlbertModel,
          'fill-mask': TFAlbertForMaskedLM,
          'question-answering': TFAlbertForQuestionAnswering,
          'text-classification': TFAlbertForSequenceClassification,
          'token-classification': TFAlbertForTokenClassification,
          'zero-shot': TFAlbertForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  __lowerCAmelCase   :		Dict         =     False
  __lowerCAmelCase   :		Union[str, Any]         =     False

  def   lowerCAmelCase__   (    self  ,						_A  ,						_A  ,						_A=False):
      SCREAMING_SNAKE_CASE_													=  super()._prepare_for_class(_A  ,						_A  ,						return_labels=_A)

      if return_labels:
          if model_class in get_values(_A):
              SCREAMING_SNAKE_CASE_													=  tf.zeros(self.model_tester.batch_size  ,						dtype=tf.intaa)

      return inputs_dict

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  TFAlbertModelTester(self)
      SCREAMING_SNAKE_CASE_													=  ConfigTester(self  ,						config_class=_A  ,						hidden_size=37)

  def   lowerCAmelCase__   (    self):
      self.config_tester.run_common_tests()

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_model(*_A)

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_for_pretraining(*_A)

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_for_masked_lm(*_A)

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_for_multiple_choice(*_A)

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_for_sequence_classification(*_A)

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_for_question_answering(*_A)


  @slow
  def   lowerCAmelCase__   (    self):
      for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          SCREAMING_SNAKE_CASE_													=  TFAlbertModel.from_pretrained(_A)
          self.assertIsNotNone(_A)






@require_tf
class 					__snake_case     (			unittest.TestCase				):

  @slow
  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  TFAlbertForPreTraining.from_pretrained('albert-base-v2')
      SCREAMING_SNAKE_CASE_													=  tf.constant([[0, 1, 2, 3, 4, 5]])
      SCREAMING_SNAKE_CASE_													=  model(_A)[0]

      SCREAMING_SNAKE_CASE_													=  [1, 6, 30000]
      self.assertEqual(output.shape  ,						_A)

      SCREAMING_SNAKE_CASE_													=  tf.constant(
          [
              [
                  [4.5_9_5_6_6_8, 0.7_4_4_6_2_7_5_4, -1.8_1_8_1_4_7],
                  [4.5_9_5_4_3_4_7, 0.7_4_5_4_1_8_4, -1.8_1_8_8_2_5_8],
                  [4.5_9_5_4_9_0_5, 0.7_4_4_8_2_3_5, -1.8_1_8_2_3_1_6],
              ]
          ])
      tf.debugging.assert_near(output[:, :3, :3]  ,						_A  ,						atol=1E-4)





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCamelCase__     :					Dict								=			{
    "configuration_donut_swin": ["DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP", "DonutSwinConfig"],
    "processing_donut": ["DonutProcessor"],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCamelCase__     :					Optional[int]								=			[
           "DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST",
           "DonutSwinModel",
           "DonutSwinPreTrainedModel",
       ]

try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCamelCase__     :					List[Any]								=			["DonutFeatureExtractor"]
       UpperCamelCase__     :					Dict								=			["DonutImageProcessor"]


if TYPE_CHECKING:
       from .configuration_donut_swin import DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, DonutSwinConfig
       from .processing_donut import DonutProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_donut_swin import (
                  DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
                  DonutSwinModel,
                  DonutSwinPreTrainedModel,
              )

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .feature_extraction_donut import DonutFeatureExtractor
              from .image_processing_donut import DonutImageProcessor

else:
       import sys

       UpperCamelCase__     :					Tuple								=			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


