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






						'''simple docstring'''






						def __init__( self   ,       _lowercase   ,       _lowercase=13   ,       _lowercase=7   ,       _lowercase=True   ,       _lowercase=True   ,       _lowercase=True   ,       _lowercase=True   ,       _lowercase=99   ,       _lowercase=32   ,       _lowercase=5   ,       _lowercase=4   ,       _lowercase=37   ,       _lowercase="gelu"   ,       _lowercase=0.1   ,       _lowercase=0.1   ,       _lowercase=512   ,       _lowercase=16   ,       _lowercase=2   ,       _lowercase=0.02   ,       _lowercase=4   ,       ):

												"""simple docstring"""



												_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  _lowercase     ( self				):

												"""simple docstring"""



												_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=_lowercase   ,       initializer_range=self.initializer_range   ,       )

												return config, input_ids, token_type_ids, attention_mask




						def  _lowercase     ( self				):

												"""simple docstring"""



												_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   UpperCAmelCase_					(		_SCREAMING_SNAKE_CASE	,							unittest.TestCase			):






						'''simple docstring'''






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




						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			FlaxAlbertModelTester(self				)




						@slow
						def  _lowercase     ( self				):

												"""simple docstring"""



												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(_lowercase				)






@require_flax
class   UpperCAmelCase_					(		unittest.TestCase			):






						'''simple docstring'''






						@slow
						def  _lowercase     ( self				):

												"""simple docstring"""



												_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(_lowercase   ,       attention_mask=_lowercase				)[0]
												_lowerCAmelCase  						=			(1, 11, 768)
												self.assertEqual(output.shape   ,       _lowercase				)
												_lowerCAmelCase  						=			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]   ,       _lowercase   ,       atol=1e-4				)				)


'''simple docstring'''



import logging

from transformers import PretrainedConfig


_lowercase    =		logging.getLogger(__name__)


_lowercase    =		{
    """bertabs-finetuned-cnndm""": """https://huggingface.co/remi/bertabs-finetuned-cnndm-extractive-abstractive-summarization/resolve/main/config.json""",
}




class   UpperCAmelCase_					(		_SCREAMING_SNAKE_CASE			):






						'''simple docstring'''






						_lowercase		:				Optional[Any]									=       '''bertabs'''




						def __init__( self   ,       _lowercase=30_522   ,       _lowercase=512   ,       _lowercase=6   ,       _lowercase=512   ,       _lowercase=8   ,       _lowercase=512   ,       _lowercase=0.2   ,       _lowercase=6   ,       _lowercase=768   ,       _lowercase=8   ,       _lowercase=2_048   ,       _lowercase=0.2   ,       **_lowercase   ,       ):

												"""simple docstring"""



												super().__init__(**_lowercase				)

												_lowerCAmelCase  						=			vocab_size
												_lowerCAmelCase  						=			max_pos

												_lowerCAmelCase  						=			enc_layers
												_lowerCAmelCase  						=			enc_hidden_size
												_lowerCAmelCase  						=			enc_heads
												_lowerCAmelCase  						=			enc_ff_size
												_lowerCAmelCase  						=			enc_dropout

												_lowerCAmelCase  						=			dec_layers
												_lowerCAmelCase  						=			dec_hidden_size
												_lowerCAmelCase  						=			dec_heads
												_lowerCAmelCase  						=			dec_ff_size
												_lowerCAmelCase  						=			dec_dropout
