


from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
	import numpy
	import tensorflow as tf

	from transformers.models.esm.modeling_tf_esm import (
	    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TFEsmForMaskedLM,
	    TFEsmForSequenceClassification,
	    TFEsmForTokenClassification,
	    TFEsmModel,
	)



class 				A       :
					'''simple docstring'''






					def __init__(self       :		Tuple ,   _UpperCAmelCase       :		Tuple ,   )		->  int:




							"""simple docstring"""


							lowercase__	     =						parent
							lowercase__	     =						13
							lowercase__	     =						7
							lowercase__	     =						True
							lowercase__	     =						True
							lowercase__	     =						True
							lowercase__	     =						99
							lowercase__	     =						32
							lowercase__	     =						2
							lowercase__	     =						4
							lowercase__	     =						37
							lowercase__	     =						"""gelu"""
							lowercase__	     =						0.1
							lowercase__	     =						0.1
							lowercase__	     =						512
							lowercase__	     =						16
							lowercase__	     =						2
							lowercase__	     =						0.02
							lowercase__	     =						3
							lowercase__	     =						4
							lowercase__	     =						None

					def     lowerCamelCase__	(self       :		Any						)		->  Optional[Any]:




							"""simple docstring"""


							lowercase__	     =						ids_tensor([self.batch_size, self.seq_length] ,   self.vocab_size						)

							lowercase__	     =						None
							if self.use_input_mask:
									lowercase__	     =						random_attention_mask([self.batch_size, self.seq_length]						)

							lowercase__	     =						None
							lowercase__	     =						None
							lowercase__	     =						None
							if self.use_labels:
									lowercase__	     =						ids_tensor([self.batch_size] ,   self.type_sequence_label_size						)
									lowercase__	     =						ids_tensor([self.batch_size, self.seq_length] ,   self.num_labels						)
									lowercase__	     =						ids_tensor([self.batch_size] ,   self.num_choices						)

							lowercase__	     =						EsmConfig(
							    vocab_size=self.vocab_size ,   hidden_size=self.hidden_size ,   num_hidden_layers=self.num_hidden_layers ,   pad_token_id=1 ,   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 ,   initializer_range=self.initializer_range ,   )

							return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels

					def     lowerCamelCase__	(self       :		str						)		->  int:




							"""simple docstring"""


							(
							    lowercase__							
							)	     =						self.prepare_config_and_inputs()

							lowercase__	     =						True
							lowercase__	     =						floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
							lowercase__	     =						ids_tensor([self.batch_size, self.seq_length] ,   vocab_size=2						)

							return (
							    config,
							    input_ids,
							    input_mask,
							    sequence_labels,
							    token_labels,
							    choice_labels,
							    encoder_hidden_states,
							    encoder_attention_mask,
							)

					def     lowerCamelCase__	(self       :		Optional[Any] ,   _UpperCAmelCase       :		str ,   _UpperCAmelCase       :		int ,   _UpperCAmelCase       :		Optional[int] ,   _UpperCAmelCase       :		Any ,   _UpperCAmelCase       :		int ,   _UpperCAmelCase       :		Union[str, Any]						)		->  List[str]:




							"""simple docstring"""


							lowercase__	     =						TFEsmModel(config=SCREAMING_SNAKE_CASE_						)
							lowercase__	     =						{"""input_ids""": input_ids, """attention_mask""": input_mask}
							lowercase__	     =						model(SCREAMING_SNAKE_CASE_						)

							lowercase__	     =						[input_ids, input_mask]
							lowercase__	     =						model(SCREAMING_SNAKE_CASE_						)

							lowercase__	     =						model(SCREAMING_SNAKE_CASE_						)

							self.parent.assertEqual(result.last_hidden_state.shape ,   (self.batch_size, self.seq_length, self.hidden_size)						)

					def     lowerCamelCase__	(self       :		Dict ,   _UpperCAmelCase       :		List[Any] ,   _UpperCAmelCase       :		Dict ,   _UpperCAmelCase       :		Optional[int] ,   _UpperCAmelCase       :		str ,   _UpperCAmelCase       :		Union[str, Any] ,   _UpperCAmelCase       :		int ,   _UpperCAmelCase       :		Optional[int] ,   _UpperCAmelCase       :		str ,   )		->  Optional[Any]:




							"""simple docstring"""


							lowercase__	     =						True

							lowercase__	     =						TFEsmModel(config=SCREAMING_SNAKE_CASE_						)
							lowercase__	     =						{
							    """input_ids""": input_ids,
							    """attention_mask""": input_mask,
							    """encoder_hidden_states""": encoder_hidden_states,
							    """encoder_attention_mask""": encoder_attention_mask,
							}
							lowercase__	     =						model(SCREAMING_SNAKE_CASE_						)

							lowercase__	     =						[input_ids, input_mask]
							lowercase__	     =						model(SCREAMING_SNAKE_CASE_ ,   encoder_hidden_states=SCREAMING_SNAKE_CASE_						)

							# Also check the case where encoder outputs are not passed
							lowercase__	     =						model(SCREAMING_SNAKE_CASE_ ,   attention_mask=SCREAMING_SNAKE_CASE_						)

							self.parent.assertEqual(result.last_hidden_state.shape ,   (self.batch_size, self.seq_length, self.hidden_size)						)

					def     lowerCamelCase__	(self       :		Union[str, Any] ,   _UpperCAmelCase       :		List[str] ,   _UpperCAmelCase       :		Tuple ,   _UpperCAmelCase       :		List[str] ,   _UpperCAmelCase       :		int ,   _UpperCAmelCase       :		Tuple ,   _UpperCAmelCase       :		Any						)		->  List[str]:




							"""simple docstring"""


							lowercase__	     =						TFEsmForMaskedLM(config=SCREAMING_SNAKE_CASE_						)
							lowercase__	     =						model([input_ids, input_mask]						)
							self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.seq_length, self.vocab_size)						)

					def     lowerCamelCase__	(self       :		Optional[Any] ,   _UpperCAmelCase       :		Tuple ,   _UpperCAmelCase       :		List[Any] ,   _UpperCAmelCase       :		Optional[int] ,   _UpperCAmelCase       :		Optional[Any] ,   _UpperCAmelCase       :		Tuple ,   _UpperCAmelCase       :		Any						)		->  Union[str, Any]:




							"""simple docstring"""


							lowercase__	     =						self.num_labels
							lowercase__	     =						TFEsmForTokenClassification(config=SCREAMING_SNAKE_CASE_						)
							lowercase__	     =						{"""input_ids""": input_ids, """attention_mask""": input_mask}
							lowercase__	     =						model(SCREAMING_SNAKE_CASE_						)
							self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.seq_length, self.num_labels)						)


					def     lowerCamelCase__	(self       :		Any						)		->  List[str]:




							"""simple docstring"""


							lowercase__	     =						self.prepare_config_and_inputs()
							(
							    lowercase__							
							)	     =						config_and_inputs
							lowercase__	     =						{"""input_ids""": input_ids, """attention_mask""": input_mask}
							return config, inputs_dict






@require_tf
class 				A       (			_lowerCAmelCase					,    _lowerCAmelCase					,    unittest.TestCase				):
					'''simple docstring'''






					A__     =					(
					    (
					        TFEsmModel,
					        TFEsmForMaskedLM,
					        TFEsmForSequenceClassification,
					        TFEsmForTokenClassification,
					    )
					    if is_tf_available()
					    else ()
					)
					A__     =					(
					    {
					        '''feature-extraction''': TFEsmModel,
					        '''fill-mask''': TFEsmForMaskedLM,
					        '''text-classification''': TFEsmForSequenceClassification,
					        '''token-classification''': TFEsmForTokenClassification,
					        '''zero-shot''': TFEsmForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					A__     =					False
					A__     =					False

					def     lowerCamelCase__	(self       :		Optional[int]						)		->  str:




							"""simple docstring"""


							lowercase__	     =						TFEsmModelTester(self						)
							lowercase__	     =						ConfigTester(self ,   config_class=SCREAMING_SNAKE_CASE_ ,   hidden_size=37						)

					def     lowerCamelCase__	(self       :		Optional[Any]						)		->  int:




							"""simple docstring"""


							self.config_tester.run_common_tests()

					def     lowerCamelCase__	(self       :		Tuple						)		->  List[str]:




							"""simple docstring"""


							lowercase__	     =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_						)

					def     lowerCamelCase__	(self       :		List[str]						)		->  Tuple:




							"""simple docstring"""


							lowercase__	     =						self.model_tester.prepare_config_and_inputs_for_decoder()
							self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE_						)

					def     lowerCamelCase__	(self       :		Tuple						)		->  str:




							"""simple docstring"""


							lowercase__	     =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_						)

					def     lowerCamelCase__	(self       :		Dict						)		->  List[str]:




							"""simple docstring"""


							lowercase__	     =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_						)

					@slow
					def     lowerCamelCase__	(self       :		List[str]						)		->  Any:




							"""simple docstring"""


							for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowercase__	     =						TFEsmModel.from_pretrained(SCREAMING_SNAKE_CASE_						)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE_						)

					@unittest.skip("""Protein models do not support embedding resizing."""						)
					def     lowerCamelCase__	(self       :		List[Any]						)		->  Dict:




							"""simple docstring"""


							pass

					@unittest.skip("""Protein models do not support embedding resizing."""						)
					def     lowerCamelCase__	(self       :		Optional[int]						)		->  Tuple:




							"""simple docstring"""


							pass


					def     lowerCamelCase__	(self       :		Optional[Any]						)		->  Dict:




							"""simple docstring"""


							lowercase__	     =						self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
									lowercase__	     =						model_class(SCREAMING_SNAKE_CASE_						)
									assert isinstance(model.get_input_embeddings() ,   tf.keras.layers.Layer						)
									if model_class is TFEsmForMaskedLM:
											# Output embedding test differs from the main test because they're a matrix, not a layer
											lowercase__	     =						model.get_bias()
											assert isinstance(SCREAMING_SNAKE_CASE_ ,   SCREAMING_SNAKE_CASE_						)
											for k, v in name.items():
													assert isinstance(SCREAMING_SNAKE_CASE_ ,   tf.Variable						)
									else:
											lowercase__	     =						model.get_output_embeddings()
											assert x is None
											lowercase__	     =						model.get_bias()
											assert name is None






@require_tf
class 				A       (			unittest.TestCase				):
					'''simple docstring'''






					@slow
					def     lowerCamelCase__	(self       :		int						)		->  str:




							"""simple docstring"""


							lowercase__	     =						TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""						)

							lowercase__	     =						tf.constant([[0, 1, 2, 3, 4, 5]]						)
							lowercase__	     =						model(SCREAMING_SNAKE_CASE_						)[0]
							lowercase__	     =						[1, 6, 33]
							self.assertEqual(list(output.numpy().shape						) ,   SCREAMING_SNAKE_CASE_						)
							# compare the actual values for a slice.
							lowercase__	     =						tf.constant(
							    [
							        [
							            [8.921_518, -10.58_9814, -6.4_671_307],
							            [-6.3_967_156, -13.91_1377, -1.1_211_915],
							            [-7.781_247, -13.95_1557, -3.740_592],
							        ]
							    ]						)
							self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() ,   expected_slice.numpy() ,   atol=1E-2						)						)


					@slow
					def     lowerCamelCase__	(self       :		Optional[Any]						)		->  List[Any]:




							"""simple docstring"""


							lowercase__	     =						TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""						)

							lowercase__	     =						tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]						)
							lowercase__	     =						model(SCREAMING_SNAKE_CASE_						)[0]
							# compare the actual values for a slice.
							lowercase__	     =						tf.constant(
							    [
							        [
							            [0.14_443_092, 0.54_125_327, 0.3_247_739],
							            [0.30_340_484, 0.00_526_676, 0.31_077_722],
							            [0.32_278_043, -0.24_987_096, 0.3_414_628],
							        ]
							    ]						)
							self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() ,   expected_slice.numpy() ,   atol=1E-4						)						)






import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    MODEL_FOR_PRETRAINING_MAPPING,
	    AlbertForMaskedLM,
	    AlbertForMultipleChoice,
	    AlbertForPreTraining,
	    AlbertForQuestionAnswering,
	    AlbertForSequenceClassification,
	    AlbertForTokenClassification,
	    AlbertModel,
	)
	from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST







class 				_a				:

			def __init__(self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_=13,   SCREAMING_SNAKE_CASE_=7,   SCREAMING_SNAKE_CASE_=True,   SCREAMING_SNAKE_CASE_=True,   SCREAMING_SNAKE_CASE_=True,   SCREAMING_SNAKE_CASE_=True,   SCREAMING_SNAKE_CASE_=99,   SCREAMING_SNAKE_CASE_=16,   SCREAMING_SNAKE_CASE_=36,   SCREAMING_SNAKE_CASE_=6,   SCREAMING_SNAKE_CASE_=6,   SCREAMING_SNAKE_CASE_=6,   SCREAMING_SNAKE_CASE_=37,   SCREAMING_SNAKE_CASE_="gelu",   SCREAMING_SNAKE_CASE_=0.1,   SCREAMING_SNAKE_CASE_=0.1,   SCREAMING_SNAKE_CASE_=512,   SCREAMING_SNAKE_CASE_=16,   SCREAMING_SNAKE_CASE_=2,   SCREAMING_SNAKE_CASE_=0.0_2,   SCREAMING_SNAKE_CASE_=3,   SCREAMING_SNAKE_CASE_=4,   SCREAMING_SNAKE_CASE_=None,   )       ->   Tuple:
						UpperCAmelCase_:		Any							      =     parent
						UpperCAmelCase_:		Union[str, Any]							      =     batch_size
						UpperCAmelCase_:		int							      =     seq_length
						UpperCAmelCase_:		Any							      =     is_training
						UpperCAmelCase_:		str							      =     use_input_mask
						UpperCAmelCase_:		Dict							      =     use_token_type_ids
						UpperCAmelCase_:		Any							      =     use_labels
						UpperCAmelCase_:		Any							      =     vocab_size
						UpperCAmelCase_:		Dict							      =     embedding_size
						UpperCAmelCase_:		Dict							      =     hidden_size
						UpperCAmelCase_:		Dict							      =     num_hidden_layers
						UpperCAmelCase_:		Dict							      =     num_hidden_groups
						UpperCAmelCase_:		int							      =     num_attention_heads
						UpperCAmelCase_:		List[str]							      =     intermediate_size
						UpperCAmelCase_:		Any							      =     hidden_act
						UpperCAmelCase_:		Dict							      =     hidden_dropout_prob
						UpperCAmelCase_:		Any							      =     attention_probs_dropout_prob
						UpperCAmelCase_:		Tuple							      =     max_position_embeddings
						UpperCAmelCase_:		Tuple							      =     type_vocab_size
						UpperCAmelCase_:		List[Any]							      =     type_sequence_label_size
						UpperCAmelCase_:		Any							      =     initializer_range
						UpperCAmelCase_:		int							      =     num_labels
						UpperCAmelCase_:		Optional[int]							      =     num_choices
						UpperCAmelCase_:		Tuple							      =     scope

			def      __snake_case       (self					)       ->   List[Any]:
						UpperCAmelCase_:		Union[str, Any]							      =     ids_tensor([self.batch_size, self.seq_length],   self.vocab_size					)

						UpperCAmelCase_:		Tuple							      =     None
						if self.use_input_mask:
									UpperCAmelCase_:		Union[str, Any]							      =     random_attention_mask([self.batch_size, self.seq_length]					)

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

						UpperCAmelCase_:		Tuple							      =     None
						UpperCAmelCase_:		Dict							      =     None
						UpperCAmelCase_:		Tuple							      =     None
						if self.use_labels:
									UpperCAmelCase_:		Optional[int]							      =     ids_tensor([self.batch_size],   self.type_sequence_label_size					)
									UpperCAmelCase_:		Dict							      =     ids_tensor([self.batch_size, self.seq_length],   self.num_labels					)
									UpperCAmelCase_:		List[Any]							      =     ids_tensor([self.batch_size],   self.num_choices					)

						UpperCAmelCase_:		Optional[int]							      =     self.get_config()

						return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

			def      __snake_case       (self					)       ->   Optional[Any]:
						return 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,   initializer_range=self.initializer_range,   num_hidden_groups=self.num_hidden_groups,   )

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Optional[Any]:
						UpperCAmelCase_:		int							      =     AlbertModel(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Optional[int]							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		List[Any]							      =     model(SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Union[str, Any]							      =     model(SCREAMING_SNAKE_CASE_					)
						self.parent.assertEqual(result.last_hidden_state.shape,   (self.batch_size, self.seq_length, self.hidden_size)					)
						self.parent.assertEqual(result.pooler_output.shape,   (self.batch_size, self.hidden_size)					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   str:
						UpperCAmelCase_:		Any							      =     AlbertForPreTraining(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		List[Any]							      =     model(
						    SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   labels=SCREAMING_SNAKE_CASE_,   sentence_order_label=SCREAMING_SNAKE_CASE_,   )
						self.parent.assertEqual(result.prediction_logits.shape,   (self.batch_size, self.seq_length, self.vocab_size)					)
						self.parent.assertEqual(result.sop_logits.shape,   (self.batch_size, config.num_labels)					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   List[str]:
						UpperCAmelCase_:		Optional[int]							      =     AlbertForMaskedLM(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Optional[Any]							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   labels=SCREAMING_SNAKE_CASE_					)
						self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.seq_length, self.vocab_size)					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Tuple:
						UpperCAmelCase_:		Optional[Any]							      =     AlbertForQuestionAnswering(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Any							      =     model(
						    SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   start_positions=SCREAMING_SNAKE_CASE_,   end_positions=SCREAMING_SNAKE_CASE_,   )
						self.parent.assertEqual(result.start_logits.shape,   (self.batch_size, self.seq_length)					)
						self.parent.assertEqual(result.end_logits.shape,   (self.batch_size, self.seq_length)					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Any:
						UpperCAmelCase_:		Optional[Any]							      =     self.num_labels
						UpperCAmelCase_:		List[Any]							      =     AlbertForSequenceClassification(SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		str							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   labels=SCREAMING_SNAKE_CASE_					)
						self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.num_labels)					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   List[Any]:
						UpperCAmelCase_:		List[str]							      =     self.num_labels
						UpperCAmelCase_:		Union[str, Any]							      =     AlbertForTokenClassification(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		List[str]							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   labels=SCREAMING_SNAKE_CASE_					)
						self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.seq_length, self.num_labels)					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Dict:
						UpperCAmelCase_:		Dict							      =     self.num_choices
						UpperCAmelCase_:		Dict							      =     AlbertForMultipleChoice(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		str							      =     input_ids.unsqueeze(1					).expand(-1,   self.num_choices,   -1					).contiguous()
						UpperCAmelCase_:		str							      =     token_type_ids.unsqueeze(1					).expand(-1,   self.num_choices,   -1					).contiguous()
						UpperCAmelCase_:		List[Any]							      =     input_mask.unsqueeze(1					).expand(-1,   self.num_choices,   -1					).contiguous()
						UpperCAmelCase_:		Union[str, Any]							      =     model(
						    SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   labels=SCREAMING_SNAKE_CASE_,   )
						self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.num_choices)					)

			def      __snake_case       (self					)       ->   Union[str, Any]:
						UpperCAmelCase_:		Optional[Any]							      =     self.prepare_config_and_inputs()
						(
						    (
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							
						):		List[str]							      =     config_and_inputs
						UpperCAmelCase_:		Tuple							      =     {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
						return config, inputs_dict







@require_torch
class 				_a				(   _lowerCAmelCase ,		_lowerCAmelCase ,		unittest.TestCase					):
			A      		=	(
			    (
			        AlbertModel,
			        AlbertForPreTraining,
			        AlbertForMaskedLM,
			        AlbertForMultipleChoice,
			        AlbertForSequenceClassification,
			        AlbertForTokenClassification,
			        AlbertForQuestionAnswering,
			    )
			    if is_torch_available()
			    else ()
			)
			A      		=	(
			    {
			        '''feature-extraction''': AlbertModel,
			        '''fill-mask''': AlbertForMaskedLM,
			        '''question-answering''': AlbertForQuestionAnswering,
			        '''text-classification''': AlbertForSequenceClassification,
			        '''token-classification''': AlbertForTokenClassification,
			        '''zero-shot''': AlbertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			A      		=	True

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_=False					)       ->   Any:
						UpperCAmelCase_:		Tuple							      =     super()._prepare_for_class(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   return_labels=SCREAMING_SNAKE_CASE_					)

						if return_labels:
									if model_class in get_values(SCREAMING_SNAKE_CASE_					):
												UpperCAmelCase_:		Dict							      =     torch.zeros(
												    (self.model_tester.batch_size, self.model_tester.seq_length),   dtype=torch.long,   device=SCREAMING_SNAKE_CASE_					)
												UpperCAmelCase_:		Optional[int]							      =     torch.zeros(
												    self.model_tester.batch_size,   dtype=torch.long,   device=SCREAMING_SNAKE_CASE_					)
						return inputs_dict

			def      __snake_case       (self					)       ->   Optional[Any]:
						UpperCAmelCase_:		Any							      =     AlbertModelTester(self					)
						UpperCAmelCase_:		Union[str, Any]							      =     ConfigTester(self,   config_class=SCREAMING_SNAKE_CASE_,   hidden_size=37					)

			def      __snake_case       (self					)       ->   List[Any]:
						self.config_tester.run_common_tests()

			def      __snake_case       (self					)       ->   List[Any]:
						UpperCAmelCase_:		List[Any]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		int							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Optional[Any]:
						UpperCAmelCase_:		Tuple							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   str:
						UpperCAmelCase_:		Optional[Any]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		List[str]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   str:
						UpperCAmelCase_:		Optional[int]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Dict:
						UpperCAmelCase_:		Optional[Any]							      =     self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
									UpperCAmelCase_:		Union[str, Any]							      =     type
									self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_					)

			@slow
			def      __snake_case       (self					)       ->   List[Any]:
						for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									UpperCAmelCase_:		Dict							      =     AlbertModel.from_pretrained(SCREAMING_SNAKE_CASE_					)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE_					)







@require_torch
class 				_a				(   unittest.TestCase					):

			@slow
			def      __snake_case       (self					)       ->   Dict:
						UpperCAmelCase_:		str							      =     AlbertModel.from_pretrained("""albert-base-v2"""					)
						UpperCAmelCase_:		List[str]							      =     torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]					)
						UpperCAmelCase_:		List[str]							      =     torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]					)
						with torch.no_grad():
									UpperCAmelCase_:		Union[str, Any]							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_					)[0]
						UpperCAmelCase_:		List[str]							      =     torch.Size((1, 11, 768)					)
						self.assertEqual(output.shape,   SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Dict							      =     torch.tensor(
						    [[[-0.6_5_1_3, 1.5_0_3_5, -0.2_7_6_6], [-0.6_5_1_5, 1.5_0_4_6, -0.2_7_8_0], [-0.6_5_1_2, 1.5_0_4_9, -0.2_7_8_4]]]					)

						self.assertTrue(torch.allclose(output[:, 1:4, 1:4],   SCREAMING_SNAKE_CASE_,   atol=1E-4					)					)






