



import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
	import jax
	import jax.numpy as jnp

	from transformers.modeling_flax_pytorch_utils import (
	    convert_pytorch_state_dict_to_flax,
	    load_flax_weights_in_pytorch_model,
	)
	from transformers.models.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
	import torch







class 				_a				:

			def __init__(self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_=14,   SCREAMING_SNAKE_CASE_=7,   SCREAMING_SNAKE_CASE_=True,   SCREAMING_SNAKE_CASE_=True,   SCREAMING_SNAKE_CASE_=False,   SCREAMING_SNAKE_CASE_=True,   SCREAMING_SNAKE_CASE_=99,   SCREAMING_SNAKE_CASE_=32,   SCREAMING_SNAKE_CASE_=4,   SCREAMING_SNAKE_CASE_=4,   SCREAMING_SNAKE_CASE_=4,   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_=0.0_2,   )       ->   Tuple:
						UpperCAmelCase_:		Optional[Any]							      =     parent
						UpperCAmelCase_:		Any							      =     batch_size
						UpperCAmelCase_:		int							      =     seq_length
						UpperCAmelCase_:		Optional[int]							      =     is_training
						UpperCAmelCase_:		List[Any]							      =     use_input_mask
						UpperCAmelCase_:		Optional[int]							      =     use_token_type_ids
						UpperCAmelCase_:		Optional[int]							      =     use_labels
						UpperCAmelCase_:		Dict							      =     vocab_size
						UpperCAmelCase_:		Optional[Any]							      =     hidden_size
						UpperCAmelCase_:		Tuple							      =     rotary_dim
						UpperCAmelCase_:		List[Any]							      =     num_hidden_layers
						UpperCAmelCase_:		List[str]							      =     num_attention_heads
						UpperCAmelCase_:		str							      =     intermediate_size
						UpperCAmelCase_:		str							      =     hidden_act
						UpperCAmelCase_:		List[str]							      =     hidden_dropout_prob
						UpperCAmelCase_:		Tuple							      =     attention_probs_dropout_prob
						UpperCAmelCase_:		List[Any]							      =     max_position_embeddings
						UpperCAmelCase_:		int							      =     initializer_range
						UpperCAmelCase_:		str							      =     None
						UpperCAmelCase_:		str							      =     vocab_size - 1
						UpperCAmelCase_:		str							      =     vocab_size - 1
						UpperCAmelCase_:		int							      =     vocab_size - 1

			def      __snake_case       (self					)       ->   Union[str, 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_:		List[Any]							      =     random_attention_mask([self.batch_size, self.seq_length]					)

						UpperCAmelCase_:		Dict							      =     GPTJConfig(
						    vocab_size=self.vocab_size,   n_embd=self.hidden_size,   n_layer=self.num_hidden_layers,   n_head=self.num_attention_heads,   n_positions=self.max_position_embeddings,   use_cache=SCREAMING_SNAKE_CASE_,   bos_token_id=self.bos_token_id,   eos_token_id=self.eos_token_id,   pad_token_id=self.pad_token_id,   rotary_dim=self.rotary_dim,   )

						return (config, input_ids, input_mask)

			def      __snake_case       (self					)       ->   str:
						UpperCAmelCase_:		Optional[int]							      =     self.prepare_config_and_inputs()
						UpperCAmelCase_      ,							UpperCAmelCase_      ,							UpperCAmelCase_:		Tuple							      =     config_and_inputs
						UpperCAmelCase_:		Dict							      =     {"""input_ids""": input_ids, """attention_mask""": attention_mask}
						return config, inputs_dict

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   int:
						UpperCAmelCase_:		List[Any]							      =     20
						UpperCAmelCase_:		List[str]							      =     model_class_name(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		int							      =     model.init_cache(input_ids.shape[0],   SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Optional[int]							      =     jnp.ones((input_ids.shape[0], max_decoder_length),   dtype="""i4"""					)

						UpperCAmelCase_:		int							      =     jnp.broadcast_to(
						    jnp.arange(input_ids.shape[-1] - 1					)[None, :],   (input_ids.shape[0], input_ids.shape[-1] - 1)					)
						UpperCAmelCase_:		Union[str, Any]							      =     model(
						    input_ids[:, :-1],   attention_mask=SCREAMING_SNAKE_CASE_,   past_key_values=SCREAMING_SNAKE_CASE_,   position_ids=SCREAMING_SNAKE_CASE_,   )

						UpperCAmelCase_:		Optional[int]							      =     jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]],   dtype="""i4"""					)
						UpperCAmelCase_:		int							      =     model(
						    input_ids[:, -1:],   attention_mask=SCREAMING_SNAKE_CASE_,   past_key_values=outputs_cache.past_key_values,   position_ids=SCREAMING_SNAKE_CASE_,   )

						UpperCAmelCase_:		Optional[int]							      =     model(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Union[str, Any]							      =     np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])					)					)
						self.parent.assertTrue(diff < 1E-3,   msg=f'Max diff is {diff}'					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Dict:
						UpperCAmelCase_:		str							      =     20
						UpperCAmelCase_:		Tuple							      =     model_class_name(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		str							      =     jnp.concatenate(
						    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])					)],   axis=-1,   )

						UpperCAmelCase_:		str							      =     model.init_cache(input_ids.shape[0],   SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Optional[int]							      =     jnp.broadcast_to(
						    jnp.arange(input_ids.shape[-1] - 1					)[None, :],   (input_ids.shape[0], input_ids.shape[-1] - 1)					)

						UpperCAmelCase_:		Any							      =     model(
						    input_ids[:, :-1],   attention_mask=SCREAMING_SNAKE_CASE_,   past_key_values=SCREAMING_SNAKE_CASE_,   position_ids=SCREAMING_SNAKE_CASE_,   )
						UpperCAmelCase_:		int							      =     jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]],   dtype="""i4"""					)
						UpperCAmelCase_:		int							      =     model(
						    input_ids[:, -1:],   past_key_values=outputs_cache.past_key_values,   attention_mask=SCREAMING_SNAKE_CASE_,   position_ids=SCREAMING_SNAKE_CASE_,   )

						UpperCAmelCase_:		int							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Optional[Any]							      =     np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])					)					)
						self.parent.assertTrue(diff < 1E-3,   msg=f'Max diff is {diff}'					)







@require_flax
class 				_a				(   _lowerCAmelCase ,		_lowerCAmelCase ,		unittest.TestCase					):
			A      		=	(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
			A      		=	(FlaxGPTJForCausalLM,) if is_flax_available() else ()

			def      __snake_case       (self					)       ->   Union[str, Any]:
						UpperCAmelCase_:		int							      =     FlaxGPTJModelTester(self					)

			def      __snake_case       (self					)       ->   List[str]:
						for model_class_name in self.all_model_classes:
									UpperCAmelCase_      ,							UpperCAmelCase_      ,							UpperCAmelCase_:		Dict							      =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.check_use_cache_forward(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   List[str]:
						for model_class_name in self.all_model_classes:
									UpperCAmelCase_      ,							UpperCAmelCase_      ,							UpperCAmelCase_:		str							      =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.check_use_cache_forward_with_attn_mask(
									    SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

			@tooslow
			def      __snake_case       (self					)       ->   List[str]:
						UpperCAmelCase_:		List[Any]							      =     GPTaTokenizer.from_pretrained("""gpt2""",   pad_token="""<|endoftext|>""",   padding_side="""left"""					)
						UpperCAmelCase_:		Any							      =     tokenizer(["""Hello this is a long string""", """Hey"""],   return_tensors="""np""",   padding=SCREAMING_SNAKE_CASE_,   truncation=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Union[str, Any]							      =     FlaxGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""					)
						UpperCAmelCase_:		Optional[int]							      =     False
						UpperCAmelCase_:		Optional[Any]							      =     model.config.eos_token_id

						UpperCAmelCase_:		List[str]							      =     jax.jit(model.generate					)

						UpperCAmelCase_:		str							      =     jit_generate(
						    inputs["""input_ids"""],   attention_mask=inputs["""attention_mask"""],   pad_token_id=tokenizer.pad_token_id					).sequences

						UpperCAmelCase_:		Optional[int]							      =     tokenizer.batch_decode(SCREAMING_SNAKE_CASE_,   skip_special_tokens=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Optional[int]							      =     [
						    """Hello this is a long string of text.\n\nI'm trying to get the text of the""",
						    """Hey, I'm a little late to the party. I'm going to""",
						]

						self.assertListEqual(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

			@is_pt_flax_cross_test
			def      __snake_case       (self					)       ->   Optional[int]:
						UpperCAmelCase_      ,							UpperCAmelCase_:		Union[str, Any]							      =     self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									with self.subTest(model_class.__name__					):
												# prepare inputs
												UpperCAmelCase_:		int							      =     self._prepare_for_class(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)
												UpperCAmelCase_:		Dict							      =     {k: torch.tensor(v.tolist()					) for k, v in prepared_inputs_dict.items()}

												# load corresponding PyTorch class
												UpperCAmelCase_:		Union[str, Any]							      =     model_class.__name__[4:]  # Skip the "Flax" at the beginning
												UpperCAmelCase_:		Tuple							      =     getattr(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

												UpperCAmelCase_      ,							UpperCAmelCase_:		Any							      =     pt_inputs["""input_ids"""].shape
												UpperCAmelCase_:		Union[str, Any]							      =     np.random.randint(0,   seq_length - 1,   size=(batch_size,)					)
												for batch_idx, start_index in enumerate(SCREAMING_SNAKE_CASE_					):
															UpperCAmelCase_:		Any							      =     0
															UpperCAmelCase_:		Any							      =     1
															UpperCAmelCase_:		Tuple							      =     0
															UpperCAmelCase_:		List[Any]							      =     1
												UpperCAmelCase_:		Tuple							      =     pt_model_class(SCREAMING_SNAKE_CASE_					).eval()
												UpperCAmelCase_:		Optional[int]							      =     model_class(SCREAMING_SNAKE_CASE_,   dtype=jnp.floataa					)

												UpperCAmelCase_:		Tuple							      =     convert_pytorch_state_dict_to_flax(pt_model.state_dict(),   SCREAMING_SNAKE_CASE_					)
												UpperCAmelCase_:		List[str]							      =     fx_state

												with torch.no_grad():
															UpperCAmelCase_:		str							      =     pt_model(**SCREAMING_SNAKE_CASE_					).to_tuple()

												UpperCAmelCase_:		Optional[int]							      =     fx_model(**SCREAMING_SNAKE_CASE_					).to_tuple()
												self.assertEqual(len(SCREAMING_SNAKE_CASE_					),   len(SCREAMING_SNAKE_CASE_					),   """Output lengths differ between Flax and PyTorch"""					)
												for fx_output, pt_output in zip(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					):
															self.assert_almost_equals(fx_output[:, -1],   pt_output[:, -1].numpy(),   4E-2					)

												with tempfile.TemporaryDirectory() as tmpdirname:
															pt_model.save_pretrained(SCREAMING_SNAKE_CASE_					)
															UpperCAmelCase_:		str							      =     model_class.from_pretrained(SCREAMING_SNAKE_CASE_,   from_pt=SCREAMING_SNAKE_CASE_					)

												UpperCAmelCase_:		Dict							      =     fx_model_loaded(**SCREAMING_SNAKE_CASE_					).to_tuple()
												self.assertEqual(
												    len(SCREAMING_SNAKE_CASE_					),   len(SCREAMING_SNAKE_CASE_					),   """Output lengths differ between Flax and PyTorch"""					)
												for fx_output_loaded, pt_output in zip(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					):
															self.assert_almost_equals(fx_output_loaded[:, -1],   pt_output[:, -1].numpy(),   4E-2					)

			@is_pt_flax_cross_test
			def      __snake_case       (self					)       ->   List[Any]:
						UpperCAmelCase_      ,							UpperCAmelCase_:		str							      =     self.model_tester.prepare_config_and_inputs_for_common()
						for model_class in self.all_model_classes:
									with self.subTest(model_class.__name__					):
												# prepare inputs
												UpperCAmelCase_:		Optional[Any]							      =     self._prepare_for_class(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)
												UpperCAmelCase_:		Optional[int]							      =     {k: torch.tensor(v.tolist()					) for k, v in prepared_inputs_dict.items()}

												# load corresponding PyTorch class
												UpperCAmelCase_:		Optional[int]							      =     model_class.__name__[4:]  # Skip the "Flax" at the beginning
												UpperCAmelCase_:		List[str]							      =     getattr(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

												UpperCAmelCase_:		List[str]							      =     pt_model_class(SCREAMING_SNAKE_CASE_					).eval()
												UpperCAmelCase_:		List[Any]							      =     model_class(SCREAMING_SNAKE_CASE_,   dtype=jnp.floataa					)

												UpperCAmelCase_:		Tuple							      =     load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE_,   fx_model.params					)
												UpperCAmelCase_      ,							UpperCAmelCase_:		List[Any]							      =     pt_inputs["""input_ids"""].shape
												UpperCAmelCase_:		List[Any]							      =     np.random.randint(0,   seq_length - 1,   size=(batch_size,)					)
												for batch_idx, start_index in enumerate(SCREAMING_SNAKE_CASE_					):
															UpperCAmelCase_:		Union[str, Any]							      =     0
															UpperCAmelCase_:		List[str]							      =     1
															UpperCAmelCase_:		int							      =     0
															UpperCAmelCase_:		Optional[int]							      =     1

												# make sure weights are tied in PyTorch
												pt_model.tie_weights()

												with torch.no_grad():
															UpperCAmelCase_:		int							      =     pt_model(**SCREAMING_SNAKE_CASE_					).to_tuple()

												UpperCAmelCase_:		Any							      =     fx_model(**SCREAMING_SNAKE_CASE_					).to_tuple()
												self.assertEqual(len(SCREAMING_SNAKE_CASE_					),   len(SCREAMING_SNAKE_CASE_					),   """Output lengths differ between Flax and PyTorch"""					)
												for fx_output, pt_output in zip(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					):
															self.assert_almost_equals(fx_output[:, -1],   pt_output[:, -1].numpy(),   4E-2					)

												with tempfile.TemporaryDirectory() as tmpdirname:
															fx_model.save_pretrained(SCREAMING_SNAKE_CASE_					)
															UpperCAmelCase_:		Tuple							      =     pt_model_class.from_pretrained(SCREAMING_SNAKE_CASE_,   from_flax=SCREAMING_SNAKE_CASE_					)

												with torch.no_grad():
															UpperCAmelCase_:		Dict							      =     pt_model_loaded(**SCREAMING_SNAKE_CASE_					).to_tuple()

												self.assertEqual(
												    len(SCREAMING_SNAKE_CASE_					),   len(SCREAMING_SNAKE_CASE_					),   """Output lengths differ between Flax and PyTorch"""					)
												for fx_output, pt_output in zip(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					):
															self.assert_almost_equals(fx_output[:, -1],   pt_output[:, -1].numpy(),   4E-2					)

			@tooslow
			def      __snake_case       (self					)       ->   Any:
						for model_class_name in self.all_model_classes:
									UpperCAmelCase_:		Tuple							      =     model_class_name.from_pretrained("""EleutherAI/gpt-j-6B"""					)
									UpperCAmelCase_:		int							      =     model(np.ones((1, 1)					)					)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE_					)











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






