






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





	"""simple docstring"""




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





	def    lowerCamelCase(self							):
		A_		:						Dict	    =   ids_tensor([self.batch_size, self.seq_length]     ,     self.vocab_size							)

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

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

		A_		:						List[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=lowerCAmelCase_     ,     initializer_range=self.initializer_range     ,     )

		return config, input_ids, token_type_ids, attention_mask





	def    lowerCamelCase(self							):
		A_		:						List[Any]	    =   self.prepare_config_and_inputs()
		A_    ,     A_    ,     A_    ,     A_		:						List[str]	    =   config_and_inputs
		A_		:						List[Any]	    =   {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
		return config, inputs_dict






@require_flax
class      SCREAMING_SNAKE_CASE  (				_SCREAMING_SNAKE_CASE       ,       unittest.TestCase						):





	"""simple docstring"""




	_A       :				List[Any]    =				(
	    (
	        FlaxAlbertModel,
	        FlaxAlbertForPreTraining,
	        FlaxAlbertForMaskedLM,
	        FlaxAlbertForMultipleChoice,
	        FlaxAlbertForQuestionAnswering,
	        FlaxAlbertForSequenceClassification,
	        FlaxAlbertForTokenClassification,
	        FlaxAlbertForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)





	def    lowerCamelCase(self							):
		A_		:						Optional[int]	    =   FlaxAlbertModelTester(self							)





	@slow
	def    lowerCamelCase(self							):
		for model_class_name in self.all_model_classes:
			A_		:						Optional[Any]	    =   model_class_name.from_pretrained("""albert-base-v2"""							)
			A_		:						Optional[Any]	    =   model(np.ones((1, 1)							)							)
			self.assertIsNotNone(lowerCAmelCase_							)






@require_flax
class      SCREAMING_SNAKE_CASE  (				unittest.TestCase						):





	"""simple docstring"""




	@slow
	def    lowerCamelCase(self							):
		A_		:						List[str]	    =   FlaxAlbertModel.from_pretrained("""albert-base-v2"""							)
		A_		:						Optional[Any]	    =   np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]							)
		A_		:						int	    =   np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]							)
		A_		:						List[str]	    =   model(lowerCAmelCase_     ,     attention_mask=lowerCAmelCase_							)[0]
		A_		:						int	    =   (1, 11, 768)
		self.assertEqual(output.shape     ,     lowerCAmelCase_							)
		A_		:						List[Any]	    =   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]     ,     lowerCAmelCase_     ,     atol=1e-4							)							)









"""simple docstring"""
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_lowerCAmelCase         =							{
    "configuration_luke": ["LUKE_PRETRAINED_CONFIG_ARCHIVE_MAP", "LukeConfig"],
    "tokenization_luke": ["LukeTokenizer"],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							_lowerCAmelCase         =							[
							    "LUKE_PRETRAINED_MODEL_ARCHIVE_LIST",
							    "LukeForEntityClassification",
							    "LukeForEntityPairClassification",
							    "LukeForEntitySpanClassification",
							    "LukeForMultipleChoice",
							    "LukeForQuestionAnswering",
							    "LukeForSequenceClassification",
							    "LukeForTokenClassification",
							    "LukeForMaskedLM",
							    "LukeModel",
							    "LukePreTrainedModel",
							]


if TYPE_CHECKING:
							from .configuration_luke import LUKE_PRETRAINED_CONFIG_ARCHIVE_MAP, LukeConfig
							from .tokenization_luke import LukeTokenizer

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_luke import (
														    LUKE_PRETRAINED_MODEL_ARCHIVE_LIST,
														    LukeForEntityClassification,
														    LukeForEntityPairClassification,
														    LukeForEntitySpanClassification,
														    LukeForMaskedLM,
														    LukeForMultipleChoice,
														    LukeForQuestionAnswering,
														    LukeForSequenceClassification,
														    LukeForTokenClassification,
														    LukeModel,
														    LukePreTrainedModel,
														)

else:
							import sys

							_lowerCAmelCase         =							_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

