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






			'''simple docstring'''

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

			def  lowerCAmelCase__		(    self	:    Union[str, Any] ):
							UpperCamelCase_:							List[Any]		=  ids_tensor([self.batch_size, self.seq_length]  ,		self.vocab_size )

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

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

							UpperCamelCase_:							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=snake_case_  ,		initializer_range=self.initializer_range  ,		)

							return config, input_ids, token_type_ids, attention_mask




			def  lowerCAmelCase__		(    self	:    str ):
							UpperCamelCase_:							str		=  self.prepare_config_and_inputs()
							UpperCamelCase_,				UpperCamelCase_,				UpperCamelCase_,				UpperCamelCase_:							int		=  config_and_inputs
							UpperCamelCase_:							List[str]		=  {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
							return config, inputs_dict




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






			'''simple docstring'''

			__UpperCamelCase :  str    =	(
			    (
			        FlaxAlbertModel,
			        FlaxAlbertForPreTraining,
			        FlaxAlbertForMaskedLM,
			        FlaxAlbertForMultipleChoice,
			        FlaxAlbertForQuestionAnswering,
			        FlaxAlbertForSequenceClassification,
			        FlaxAlbertForTokenClassification,
			        FlaxAlbertForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)

			def  lowerCAmelCase__		(    self	:    Union[str, Any] ):
							UpperCamelCase_:							Optional[Any]		=  FlaxAlbertModelTester(self )




			@slow
			def  lowerCAmelCase__		(    self	:    str ):
							for model_class_name in self.all_model_classes:
											UpperCamelCase_:							Optional[int]		=  model_class_name.from_pretrained("""albert-base-v2""" )
											UpperCamelCase_:							List[str]		=  model(np.ones((1, 1) ) )
											self.assertIsNotNone(snake_case_ )




@require_flax
class     _UpperCamelCase			(   unittest.TestCase					):






			'''simple docstring'''

			@slow
			def  lowerCAmelCase__		(    self	:    int ):
							UpperCamelCase_:							str		=  FlaxAlbertModel.from_pretrained("""albert-base-v2""" )
							UpperCamelCase_:							Any		=  np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]] )
							UpperCamelCase_:							Optional[int]		=  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] )
							UpperCamelCase_:							Optional[Any]		=  model(snake_case_  ,		attention_mask=snake_case_ )[0]
							UpperCamelCase_:							Tuple		=  (1, 11, 768)
							self.assertEqual(output.shape  ,		snake_case_ )
							UpperCamelCase_:							List[str]		=  np.array(
							    [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]] )

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








from typing import TYPE_CHECKING

from ...utils import _LazyModule


lowerCamelCase_  :	Dict 				=					{"""tokenization_byt5""": ["""ByT5Tokenizer"""]}


if TYPE_CHECKING:
							from .tokenization_byta import ByTaTokenizer
else:
							import sys

							lowerCamelCase_  :	Any 				=					_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)






