




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 			__lowerCamelCase	(   unittest.TestCase   ):




    """simple docstring"""

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


    def        a				(	self					:					Any       )				->	Tuple:
         lowerCAmelCase__          =			ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size       )

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

         lowerCAmelCase__          =			None
         if self.use_token_type_ids:
              lowerCAmelCase__          =			ids_tensor([self.batch_size, self.seq_length]       ,     self.type_vocab_size       )

         lowerCAmelCase__          =			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=SCREAMING_SNAKE_CASE__       ,     initializer_range=self.initializer_range       ,     )

         return config, input_ids, token_type_ids, attention_mask


    def        a				(	self					:					Optional[Any]       )				->	Dict:
         lowerCAmelCase__          =			self.prepare_config_and_inputs()
         lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			config_and_inputs
         lowerCAmelCase__          =			{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
         return config, inputs_dict






@require_flax
class 			__lowerCamelCase	(   UpperCamelCase__     ,    unittest.TestCase   ):




    """simple docstring"""

    snake_case__ =				(
        (
            FlaxAlbertModel,
            FlaxAlbertForPreTraining,
            FlaxAlbertForMaskedLM,
            FlaxAlbertForMultipleChoice,
            FlaxAlbertForQuestionAnswering,
            FlaxAlbertForSequenceClassification,
            FlaxAlbertForTokenClassification,
            FlaxAlbertForQuestionAnswering,
        )
        if is_flax_available()
        else ()
    )


    def        a				(	self					:					Dict       )				->	Union[str, Any]:
         lowerCAmelCase__          =			FlaxAlbertModelTester(self       )


    @slow
    def        a				(	self					:					int       )				->	Any:
         for model_class_name in self.all_model_classes:
              lowerCAmelCase__          =			model_class_name.from_pretrained("albert-base-v2"       )
              lowerCAmelCase__          =			model(np.ones((1, 1)       )       )
              self.assertIsNotNone(SCREAMING_SNAKE_CASE__       )






@require_flax
class 			__lowerCamelCase	(   unittest.TestCase   ):




    """simple docstring"""

    @slow
    def        a				(	self					:					Any       )				->	Tuple:
         lowerCAmelCase__          =			FlaxAlbertModel.from_pretrained("albert-base-v2"       )
         lowerCAmelCase__          =			np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]       )
         lowerCAmelCase__          =			np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       )[0]
         lowerCAmelCase__          =			(1, 11, 768)
         self.assertEqual(output.shape       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			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]       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-4       )       )






import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin





class 			__lowerCamelCase	(   UpperCamelCase__   ):




    """simple docstring"""

    snake_case__ =				"Speech2TextFeatureExtractor"
    snake_case__ =				"Speech2TextTokenizer"


    def __init__(	self					:					int       ,     SCREAMING_SNAKE_CASE__					:					Dict       ,     SCREAMING_SNAKE_CASE__					:					Tuple       )				->	Optional[Any]:
         super().__init__(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.feature_extractor
         lowerCAmelCase__          =			False


    def __call__(	self					:					str       ,     *SCREAMING_SNAKE_CASE__					:					int       ,     **SCREAMING_SNAKE_CASE__					:					int       )				->	Optional[Any]:
         # For backward compatibility
         if self._in_target_context_manager:
              return self.current_processor(*SCREAMING_SNAKE_CASE__       ,     **SCREAMING_SNAKE_CASE__       )

         if "raw_speech" in kwargs:
              warnings.warn("Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead."       )
              lowerCAmelCase__          =			kwargs.pop("raw_speech"       )
         else:
              lowerCAmelCase__          =			kwargs.pop("audio"       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			kwargs.pop("sampling_rate"       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			kwargs.pop("text"       ,     SCREAMING_SNAKE_CASE__       )
         if len(SCREAMING_SNAKE_CASE__       ) > 0:
              lowerCAmelCase__          =			args[0]
              lowerCAmelCase__          =			args[1:]

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

         if audio is not None:
              lowerCAmelCase__          =			self.feature_extractor(SCREAMING_SNAKE_CASE__       ,     *SCREAMING_SNAKE_CASE__       ,     sampling_rate=SCREAMING_SNAKE_CASE__       ,     **SCREAMING_SNAKE_CASE__       )
         if text is not None:
              lowerCAmelCase__          =			self.tokenizer(SCREAMING_SNAKE_CASE__       ,     **SCREAMING_SNAKE_CASE__       )

         if text is None:
              return inputs
         elif audio is None:
              return encodings
         else:
              lowerCAmelCase__          =			encodings["input_ids"]
              return inputs


    def        a				(	self					:					str       ,     *SCREAMING_SNAKE_CASE__					:					Dict       ,     **SCREAMING_SNAKE_CASE__					:					int       )				->	Union[str, Any]:
         return self.tokenizer.batch_decode(*SCREAMING_SNAKE_CASE__       ,     **SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Dict       ,     *SCREAMING_SNAKE_CASE__					:					Optional[int]       ,     **SCREAMING_SNAKE_CASE__					:					Optional[int]       )				->	Optional[Any]:
         return self.tokenizer.decode(*SCREAMING_SNAKE_CASE__       ,     **SCREAMING_SNAKE_CASE__       )


    @contextmanager
    def        a				(	self					:					Union[str, Any]       )				->	Any:
         warnings.warn(
             "`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your "
             "labels by using the argument `text` of the regular `__call__` method (either in the same call as "
             "your audio inputs, or in a separate call."       )
         lowerCAmelCase__          =			True
         lowerCAmelCase__          =			self.tokenizer
         yield
         lowerCAmelCase__          =			self.feature_extractor
         lowerCAmelCase__          =			False
