


"""simple docstring"""







import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
 import jax.numpy as jnp

 from transformers.models.albert.modeling_flax_albert import (
     FlaxAlbertForMaskedLM,
     FlaxAlbertForMultipleChoice,
     FlaxAlbertForPreTraining,
     FlaxAlbertForQuestionAnswering,
     FlaxAlbertForSequenceClassification,
     FlaxAlbertForTokenClassification,
     FlaxAlbertModel,
 )






class 						_A  (  unittest.TestCase  ):







    """simple docstring"""




    def __init__(					self				:    List[Any]				, __UpperCAmelCase				:    List[Any]				, __UpperCAmelCase				:    Optional[Any]=13				, __UpperCAmelCase				:    List[str]=7				, __UpperCAmelCase				:    Optional[int]=True				, __UpperCAmelCase				:    List[str]=True				, __UpperCAmelCase				:    List[Any]=True				, __UpperCAmelCase				:    List[str]=True				, __UpperCAmelCase				:    Optional[Any]=99				, __UpperCAmelCase				:    List[str]=32				, __UpperCAmelCase				:    List[Any]=5				, __UpperCAmelCase				:    Optional[int]=4				, __UpperCAmelCase				:    List[str]=37				, __UpperCAmelCase				:    List[Any]="gelu"				, __UpperCAmelCase				:    Any=0.1				, __UpperCAmelCase				:    Tuple=0.1				, __UpperCAmelCase				:    List[str]=512				, __UpperCAmelCase				:    List[str]=16				, __UpperCAmelCase				:    List[str]=2				, __UpperCAmelCase				:    int=0.02				, __UpperCAmelCase				:    Dict=4				, ):
      a						:		Any      							=			parent
      a						:		Any      							=			batch_size
      a						:		List[str]      							=			seq_length
      a						:		Optional[Any]      							=			is_training
      a						:		int      							=			use_attention_mask
      a						:		Optional[Any]      							=			use_token_type_ids
      a						:		Any      							=			use_labels
      a						:		Any      							=			vocab_size
      a						:		Any      							=			hidden_size
      a						:		Any      							=			num_hidden_layers
      a						:		Dict      							=			num_attention_heads
      a						:		str      							=			intermediate_size
      a						:		Any      							=			hidden_act
      a						:		str      							=			hidden_dropout_prob
      a						:		Union[str, Any]      							=			attention_probs_dropout_prob
      a						:		Tuple      							=			max_position_embeddings
      a						:		Dict      							=			type_vocab_size
      a						:		Dict      							=			type_sequence_label_size
      a						:		List[str]      							=			initializer_range
      a						:		List[Any]      							=			num_choices






    def 	__snake_case							(					self				:    str):
      a						:		Optional[int]      							=			ids_tensor([self.batch_size, self.seq_length]				, self.vocab_size)

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

      a						:		Union[str, Any]      							=			None
      if self.use_token_type_ids:
        a						:		Any      							=			ids_tensor([self.batch_size, self.seq_length]				, self.type_vocab_size)

      a						:		Union[str, Any]      							=			AlbertConfig(
          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				, is_decoder=__UpperCAmelCase				, initializer_range=self.initializer_range				, )

      return config, input_ids, token_type_ids, attention_mask






    def 	__snake_case							(					self				:    Optional[int]):
      a						:		Any      							=			self.prepare_config_and_inputs()
      a							,		a							,		a							,		a						:		str      							=			config_and_inputs
      a						:		Any      							=			{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
      return config, inputs_dict






@require_flax
class 						_A  (  _a    ,unittest.TestCase  ):







    """simple docstring"""




    UpperCAmelCase       :							Union[str, Any]	=       (
        (
            FlaxAlbertModel,
            FlaxAlbertForPreTraining,
            FlaxAlbertForMaskedLM,
            FlaxAlbertForMultipleChoice,
            FlaxAlbertForQuestionAnswering,
            FlaxAlbertForSequenceClassification,
            FlaxAlbertForTokenClassification,
            FlaxAlbertForQuestionAnswering,
        )
        if is_flax_available()
        else ()
    )






    def 	__snake_case							(					self				:    str):
      a						:		List[Any]      							=			FlaxAlbertModelTester(self)






    @slow
    def 	__snake_case							(					self				:    Any):
      for model_class_name in self.all_model_classes:
        a						:		Dict      							=			model_class_name.from_pretrained("albert-base-v2")
        a						:		List[str]      							=			model(np.ones((1, 1)))
        self.assertIsNotNone(__UpperCAmelCase)






@require_flax
class 						_A  (  unittest.TestCase  ):







    """simple docstring"""




    @slow
    def 	__snake_case							(					self				:    int):
      a						:		Optional[Any]      							=			FlaxAlbertModel.from_pretrained("albert-base-v2")
      a						:		Tuple      							=			np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]])
      a						:		Any      							=			np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
      a						:		Tuple      							=			model(__UpperCAmelCase				, attention_mask=__UpperCAmelCase)[0]
      a						:		Optional[Any]      							=			(1, 11, 768)
      self.assertEqual(output.shape				, __UpperCAmelCase)
      a						:		Optional[Any]      							=			np.array(
          [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]])

      self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]				, __UpperCAmelCase				, atol=1e-4))







from ...processing_utils import ProcessorMixin


class      lowerCamelCase    (SCREAMING_SNAKE_CASE__    ):
  """simple docstring"""







  lowerCamelCase__   =  ['''image_processor''', '''feature_extractor''']
  lowerCamelCase__   =  '''TvltImageProcessor'''
  lowerCamelCase__   =  '''TvltFeatureExtractor'''







  def __init__(							self       :				List[str]     ,   __magic_name__       :				Any     ,   __magic_name__       :				Any      )     ->		int:
        super().__init__(image_processor=__magic_name__     ,   feature_extractor=__magic_name__      )

        SCREAMING_SNAKE_CASE_      =     image_processor
        SCREAMING_SNAKE_CASE_      =     feature_extractor







  def __call__(							self       :				List[str]     ,   __magic_name__       :				Optional[Any]=None     ,   __magic_name__       :				Union[str, Any]=None     ,   __magic_name__       :				int=None     ,   __magic_name__       :				str=None     ,   __magic_name__       :				Any=False     ,   __magic_name__       :				int=False     ,   *__magic_name__       :				int     ,   **__magic_name__       :				Any     ,   )     ->		List[Any]:

        if images is None and audio is None:
              raise ValueError("You need to specify either an `images` or `audio` input to process."      )

        SCREAMING_SNAKE_CASE_      =     None
        if images is not None:
              SCREAMING_SNAKE_CASE_      =     self.image_processor(__magic_name__     ,   mask_pixel=__magic_name__     ,   *__magic_name__     ,   **__magic_name__      )
        if images_mixed is not None:
              SCREAMING_SNAKE_CASE_      =     self.image_processor(__magic_name__     ,   is_mixed=__magic_name__     ,   *__magic_name__     ,   **__magic_name__      )
        if audio is not None:
              SCREAMING_SNAKE_CASE_      =     self.feature_extractor(
                  __magic_name__     ,   *__magic_name__     ,   sampling_rate=__magic_name__     ,   mask_audio=__magic_name__     ,   **__magic_name__      )

        SCREAMING_SNAKE_CASE_      =     {}
        if audio is not None:
              output_dict.update(__magic_name__      )
        if images is not None:
              output_dict.update(__magic_name__      )
        if images_mixed_dict is not None:
              output_dict.update(__magic_name__      )
        return output_dict







  @property
  def       __A   (							self       :				Optional[Any]      )     ->		int:
        SCREAMING_SNAKE_CASE_      =     self.image_processor.model_input_names
        SCREAMING_SNAKE_CASE_      =     self.feature_extractor.model_input_names
        return list(dict.fromkeys(image_processor_input_names + feature_extractor_input_names      )      )


