



import argparse
import logging
import os
import re

import tensorflow as tf

from transformers import (
    AutoConfig,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
    PushToHubCallback,
    TFAutoModelForMaskedLM,
    create_optimizer,
)


a      :  Any			     =			logging.getLogger(__name__)

a      :  str			     =			tf.data.AUTOTUNE







def  lowerCAmelCase_	():





			"""simple docstring"""


			UpperCAmelCase_:		Optional[int]							      =     argparse.ArgumentParser(description="""Train a masked language model on TPU."""  )
			parser.add_argument(
			    """--pretrained_model_config"""		,						type=lowerCAmelCase__		,						default="""roberta-base"""		,						help="""The model config to use. Note that we don't copy the model's weights, only the config!"""		,						)
			parser.add_argument(
			    """--tokenizer"""		,						type=lowerCAmelCase__		,						default="""unigram-tokenizer-wikitext"""		,						help="""The name of the tokenizer to load. We use the pretrained tokenizer to initialize the model's vocab size."""		,						)

			parser.add_argument(
			    """--per_replica_batch_size"""		,						type=lowerCAmelCase__		,						default=8		,						help="""Batch size per TPU core."""		,						)

			parser.add_argument(
			    """--no_tpu"""		,						action="""store_true"""		,						help="""If set, run on CPU and don't try to initialize a TPU. Useful for debugging on non-TPU instances."""		,						)

			parser.add_argument(
			    """--tpu_name"""		,						type=lowerCAmelCase__		,						help="""Name of TPU resource to initialize. Should be blank on Colab, and 'local' on TPU VMs."""		,						default="""local"""		,						)

			parser.add_argument(
			    """--tpu_zone"""		,						type=lowerCAmelCase__		,						help="""Google cloud zone that TPU resource is located in. Only used for non-Colab TPU nodes."""		,						)

			parser.add_argument(
			    """--gcp_project"""		,						type=lowerCAmelCase__		,						help="""Google cloud project name. Only used for non-Colab TPU nodes."""  )

			parser.add_argument(
			    """--bfloat16"""		,						action="""store_true"""		,						help="""Use mixed-precision bfloat16 for training. This is the recommended lower-precision format for TPU."""		,						)

			parser.add_argument(
			    """--train_dataset"""		,						type=lowerCAmelCase__		,						help="""Path to training dataset to load. If the path begins with `gs://`"""
			    """ then the dataset will be loaded from a Google Cloud Storage bucket."""		,						)

			parser.add_argument(
			    """--shuffle_buffer_size"""		,						type=lowerCAmelCase__		,						default=2**1_8		,						help="""Size of the shuffle buffer (in samples)"""		,						)

			parser.add_argument(
			    """--eval_dataset"""		,						type=lowerCAmelCase__		,						help="""Path to evaluation dataset to load. If the path begins with `gs://`"""
			    """ then the dataset will be loaded from a Google Cloud Storage bucket."""		,						)

			parser.add_argument(
			    """--num_epochs"""		,						type=lowerCAmelCase__		,						default=1		,						help="""Number of epochs to train for."""		,						)

			parser.add_argument(
			    """--learning_rate"""		,						type=lowerCAmelCase__		,						default=1e-4		,						help="""Learning rate to use for training."""		,						)

			parser.add_argument(
			    """--weight_decay_rate"""		,						type=lowerCAmelCase__		,						default=1e-3		,						help="""Weight decay rate to use for training."""		,						)

			parser.add_argument(
			    """--max_length"""		,						type=lowerCAmelCase__		,						default=5_1_2		,						help="""Maximum length of tokenized sequences. Should match the setting used in prepare_tfrecord_shards.py"""		,						)

			parser.add_argument(
			    """--mlm_probability"""		,						type=lowerCAmelCase__		,						default=0.15		,						help="""Fraction of tokens to mask during training."""		,						)

			parser.add_argument("""--output_dir"""		,						type=lowerCAmelCase__		,						required=lowerCAmelCase__		,						help="""Path to save model checkpoints to."""  )
			parser.add_argument("""--hub_model_id"""		,						type=lowerCAmelCase__		,						help="""Model ID to upload to on the Hugging Face Hub."""  )

			UpperCAmelCase_:		Union[str, Any]							      =     parser.parse_args()
			return args







def  lowerCAmelCase_	(lowerCAmelCase__:		Tuple  ):





			"""simple docstring"""


			try:
						if args.tpu_name:
									UpperCAmelCase_:		Optional[Any]							      =     tf.distribute.cluster_resolver.TPUClusterResolver(
									    args.tpu_name		,						zone=args.tpu_zone		,						project=args.gcp_project  )
						else:
									UpperCAmelCase_:		str							      =     tf.distribute.cluster_resolver.TPUClusterResolver()
			except ValueError:
						raise RuntimeError(
						    """Couldn't connect to TPU! Most likely you need to specify --tpu_name, --tpu_zone, or """
						    """--gcp_project. When running on a TPU VM, use --tpu_name local."""  )

			tf.config.experimental_connect_to_cluster(lowerCAmelCase__  )
			tf.tpu.experimental.initialize_tpu_system(lowerCAmelCase__  )

			return tpu







def  lowerCAmelCase_	(lowerCAmelCase__:		List[str]  ):





			"""simple docstring"""


			UpperCAmelCase_:		str							      =     0
			for file in file_list:
						UpperCAmelCase_:		Optional[Any]							      =     file.split("""/"""  )[-1]
						UpperCAmelCase_:		int							      =     re.search(r"""-\d+-(\d+)\.tfrecord"""		,						lowerCAmelCase__  ).group(1  )
						UpperCAmelCase_:		str							      =     int(lowerCAmelCase__  )
						num_samples += sample_count

			return num_samples







def  lowerCAmelCase_	(lowerCAmelCase__:		Union[str, Any]		,						lowerCAmelCase__:		Tuple		,						lowerCAmelCase__:		str		,						lowerCAmelCase__:		List[Any]		,						lowerCAmelCase__:		List[Any]		,						lowerCAmelCase__:		Dict=None  ):





			"""simple docstring"""


			UpperCAmelCase_:		Tuple							      =     count_samples(lowerCAmelCase__  )
			UpperCAmelCase_:		List[Any]							      =     tf.data.Dataset.from_tensor_slices(lowerCAmelCase__  )
			if shuffle:
						UpperCAmelCase_:		Any							      =     dataset.shuffle(len(lowerCAmelCase__  )  )
			UpperCAmelCase_:		Union[str, Any]							      =     tf.data.TFRecordDataset(lowerCAmelCase__		,						num_parallel_reads=lowerCAmelCase__  )
			# TF can't infer the total sample count because it doesn't read all the records yet, so we assert it here
			UpperCAmelCase_:		Dict							      =     dataset.apply(tf.data.experimental.assert_cardinality(lowerCAmelCase__  )  )
			UpperCAmelCase_:		Tuple							      =     dataset.map(lowerCAmelCase__		,						num_parallel_calls=lowerCAmelCase__  )
			if shuffle:
						assert shuffle_buffer_size is not None
						UpperCAmelCase_:		List[str]							      =     dataset.shuffle(args.shuffle_buffer_size  )
			UpperCAmelCase_:		str							      =     dataset.batch(lowerCAmelCase__		,						drop_remainder=lowerCAmelCase__  )
			UpperCAmelCase_:		List[str]							      =     dataset.map(lowerCAmelCase__		,						num_parallel_calls=lowerCAmelCase__  )
			UpperCAmelCase_:		Dict							      =     dataset.prefetch(lowerCAmelCase__  )
			return dataset







def  lowerCAmelCase_	(lowerCAmelCase__:		List[str]  ):





			"""simple docstring"""


			if not args.no_tpu:
						UpperCAmelCase_:		List[Any]							      =     initialize_tpu(lowerCAmelCase__  )
						UpperCAmelCase_:		Any							      =     tf.distribute.TPUStrategy(lowerCAmelCase__  )
			else:
						UpperCAmelCase_:		Dict							      =     tf.distribute.OneDeviceStrategy(device="""/gpu:0"""  )

			if args.bfloataa:
						tf.keras.mixed_precision.set_global_policy("""mixed_bfloat16"""  )

			UpperCAmelCase_:		List[Any]							      =     AutoTokenizer.from_pretrained(args.tokenizer  )
			UpperCAmelCase_:		Any							      =     AutoConfig.from_pretrained(args.pretrained_model_config  )
			UpperCAmelCase_:		int							      =     tokenizer.vocab_size

			UpperCAmelCase_:		List[str]							      =     tf.io.gfile.glob(os.path.join(args.train_dataset		,						"""*.tfrecord"""  )  )
			if not training_records:
						raise ValueError(F'No .tfrecord files found in {args.train_dataset}.'  )
			UpperCAmelCase_:		Dict							      =     tf.io.gfile.glob(os.path.join(args.eval_dataset		,						"""*.tfrecord"""  )  )
			if not eval_records:
						raise ValueError(F'No .tfrecord files found in {args.eval_dataset}.'  )

			UpperCAmelCase_:		str							      =     count_samples(lowerCAmelCase__  )

			UpperCAmelCase_:		Union[str, Any]							      =     num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
			UpperCAmelCase_:		Union[str, Any]							      =     steps_per_epoch * args.num_epochs

			with strategy.scope():
						UpperCAmelCase_:		Optional[int]							      =     TFAutoModelForMaskedLM.from_config(lowerCAmelCase__  )
						model(model.dummy_inputs  )  # Pass some dummy inputs through the model to ensure all the weights are built
						UpperCAmelCase_      ,							UpperCAmelCase_:		List[str]							      =     create_optimizer(
						    num_train_steps=lowerCAmelCase__		,						num_warmup_steps=total_train_steps // 2_0		,						init_lr=args.learning_rate		,						weight_decay_rate=args.weight_decay_rate		,						)

						# Transformers models compute the right loss for their task by default when labels are passed, and will
						# use this for training unless you specify your own loss function in compile().
						model.compile(optimizer=lowerCAmelCase__		,						metrics=["""accuracy"""]  )

			def decode_fn(lowerCAmelCase__:		Dict  ):
						UpperCAmelCase_:		Any							      =     {
						    """input_ids""": tf.io.FixedLenFeature(dtype=tf.intaa		,						shape=(args.max_length,)  ),
						    """attention_mask""": tf.io.FixedLenFeature(dtype=tf.intaa		,						shape=(args.max_length,)  ),
						}
						return tf.io.parse_single_example(lowerCAmelCase__		,						lowerCAmelCase__  )

			# Many of the data collators in Transformers are TF-compilable when return_tensors == "tf", so we can
			# use their methods in our data pipeline.
			UpperCAmelCase_:		str							      =     DataCollatorForLanguageModeling(
			    tokenizer=lowerCAmelCase__		,						mlm_probability=args.mlm_probability		,						mlm=lowerCAmelCase__		,						return_tensors="""tf"""  )

			def mask_with_collator(lowerCAmelCase__:		int  ):
						# TF really needs an isin() function
						UpperCAmelCase_:		int							      =     (
						    ~tf.cast(batch["""attention_mask"""]		,						tf.bool  )
						    | (batch["""input_ids"""] == tokenizer.cls_token_id)
						    | (batch["""input_ids"""] == tokenizer.sep_token_id)
						)
						UpperCAmelCase_      ,							UpperCAmelCase_:		str							      =     data_collator.tf_mask_tokens(
						    batch["""input_ids"""]		,						vocab_size=len(lowerCAmelCase__  )		,						mask_token_id=tokenizer.mask_token_id		,						special_tokens_mask=lowerCAmelCase__		,						)
						return batch

			UpperCAmelCase_:		Any							      =     args.per_replica_batch_size * strategy.num_replicas_in_sync

			UpperCAmelCase_:		List[Any]							      =     prepare_dataset(
			    lowerCAmelCase__		,						decode_fn=lowerCAmelCase__		,						mask_fn=lowerCAmelCase__		,						batch_size=lowerCAmelCase__		,						shuffle=lowerCAmelCase__		,						shuffle_buffer_size=args.shuffle_buffer_size		,						)

			UpperCAmelCase_:		Dict							      =     prepare_dataset(
			    lowerCAmelCase__		,						decode_fn=lowerCAmelCase__		,						mask_fn=lowerCAmelCase__		,						batch_size=lowerCAmelCase__		,						shuffle=lowerCAmelCase__		,						)

			UpperCAmelCase_:		List[Any]							      =     []
			if args.hub_model_id:
						callbacks.append(
						    PushToHubCallback(output_dir=args.output_dir		,						hub_model_id=args.hub_model_id		,						tokenizer=lowerCAmelCase__  )  )

			model.fit(
			    lowerCAmelCase__		,						validation_data=lowerCAmelCase__		,						epochs=args.num_epochs		,						callbacks=lowerCAmelCase__		,						)

			model.save_pretrained(args.output_dir  )


if __name__ == "__main__":
	a      :  Dict			     =			parse_args()
	main(args)











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






