


import unittest

from transformers import MobileBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, 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,
		    MobileBertForMaskedLM,
		    MobileBertForMultipleChoice,
		    MobileBertForNextSentencePrediction,
		    MobileBertForPreTraining,
		    MobileBertForQuestionAnswering,
		    MobileBertForSequenceClassification,
		    MobileBertForTokenClassification,
		    MobileBertModel,
		)







class        __lowercase       :



		"""simple docstring"""

		def __init__( self		,      A		,      A=13		,      A=7		,      A=True		,      A=True		,      A=True		,      A=True		,      A=99		,      A=64		,      A=32		,      A=5		,      A=4		,      A=37		,      A="gelu"		,      A=0.1		,      A=0.1		,      A=5_12		,      A=16		,      A=2		,      A=0.02		,      A=3		,      A=4		,      A=None		,      )  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  parent
					lowerCamelCase    		=  batch_size
					lowerCamelCase    		=  seq_length
					lowerCamelCase    		=  is_training
					lowerCamelCase    		=  use_input_mask
					lowerCamelCase    		=  use_token_type_ids
					lowerCamelCase    		=  use_labels
					lowerCamelCase    		=  vocab_size
					lowerCamelCase    		=  hidden_size
					lowerCamelCase    		=  embedding_size
					lowerCamelCase    		=  num_hidden_layers
					lowerCamelCase    		=  num_attention_heads
					lowerCamelCase    		=  intermediate_size
					lowerCamelCase    		=  hidden_act
					lowerCamelCase    		=  hidden_dropout_prob
					lowerCamelCase    		=  attention_probs_dropout_prob
					lowerCamelCase    		=  max_position_embeddings
					lowerCamelCase    		=  type_vocab_size
					lowerCamelCase    		=  type_sequence_label_size
					lowerCamelCase    		=  initializer_range
					lowerCamelCase    		=  num_labels
					lowerCamelCase    		=  num_choices
					lowerCamelCase    		=  scope







		def 					__A   ( self	)  -> List[str]:







					'''simple docstring'''







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

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

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

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

					lowerCamelCase    		=  self.get_config()

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







		def 					__A   ( self	)  -> Dict:







					'''simple docstring'''







					return MobileBertConfig(
					    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		,      embedding_size=self.embedding_size		,      hidden_act=self.hidden_act		,      hidden_dropout_prob=self.hidden_dropout_prob		,      attention_probs_dropout_prob=self.attention_probs_dropout_prob		,      max_position_embeddings=self.max_position_embeddings		,      type_vocab_size=self.type_vocab_size		,      is_decoder=A		,      initializer_range=self.initializer_range		,      )







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> Tuple:







					'''simple docstring'''







					lowerCamelCase    		=  MobileBertModel(config=A	)
					model.to(A	)
					model.eval()
					lowerCamelCase    		=  model(A		,      attention_mask=A		,      token_type_ids=A	)
					lowerCamelCase    		=  model(A		,      token_type_ids=A	)
					lowerCamelCase    		=  model(A	)

					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 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  MobileBertForMaskedLM(config=A	)
					model.to(A	)
					model.eval()
					lowerCamelCase    		=  model(A		,      attention_mask=A		,      token_type_ids=A		,      labels=A	)
					self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.vocab_size)	)







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> Any:







					'''simple docstring'''







					lowerCamelCase    		=  MobileBertForNextSentencePrediction(config=A	)
					model.to(A	)
					model.eval()
					lowerCamelCase    		=  model(
					    A		,      attention_mask=A		,      token_type_ids=A		,      labels=A		,      )
					self.parent.assertEqual(result.logits.shape		,      (self.batch_size, 2)	)







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  MobileBertForPreTraining(config=A	)
					model.to(A	)
					model.eval()
					lowerCamelCase    		=  model(
					    A		,      attention_mask=A		,      token_type_ids=A		,      labels=A		,      next_sentence_label=A		,      )
					self.parent.assertEqual(result.prediction_logits.shape		,      (self.batch_size, self.seq_length, self.vocab_size)	)
					self.parent.assertEqual(result.seq_relationship_logits.shape		,      (self.batch_size, 2)	)







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  MobileBertForQuestionAnswering(config=A	)
					model.to(A	)
					model.eval()
					lowerCamelCase    		=  model(
					    A		,      attention_mask=A		,      token_type_ids=A		,      start_positions=A		,      end_positions=A		,      )
					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 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> str:







					'''simple docstring'''







					lowerCamelCase    		=  self.num_labels
					lowerCamelCase    		=  MobileBertForSequenceClassification(A	)
					model.to(A	)
					model.eval()
					lowerCamelCase    		=  model(A		,      attention_mask=A		,      token_type_ids=A		,      labels=A	)
					self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_labels)	)







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> Dict:







					'''simple docstring'''







					lowerCamelCase    		=  self.num_labels
					lowerCamelCase    		=  MobileBertForTokenClassification(config=A	)
					model.to(A	)
					model.eval()
					lowerCamelCase    		=  model(A		,      attention_mask=A		,      token_type_ids=A		,      labels=A	)
					self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.num_labels)	)







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> Union[str, Any]:







					'''simple docstring'''







					lowerCamelCase    		=  self.num_choices
					lowerCamelCase    		=  MobileBertForMultipleChoice(config=A	)
					model.to(A	)
					model.eval()
					lowerCamelCase    		=  input_ids.unsqueeze(1	).expand(-1		,      self.num_choices		,      -1	).contiguous()
					lowerCamelCase    		=  token_type_ids.unsqueeze(1	).expand(-1		,      self.num_choices		,      -1	).contiguous()
					lowerCamelCase    		=  input_mask.unsqueeze(1	).expand(-1		,      self.num_choices		,      -1	).contiguous()
					lowerCamelCase    		=  model(
					    A		,      attention_mask=A		,      token_type_ids=A		,      labels=A		,      )
					self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_choices)	)







		def 					__A   ( self	)  -> List[str]:







					'''simple docstring'''







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







@require_torch
class        __lowercase       (							a_			,							a_			,							unittest.TestCase  ):



		"""simple docstring"""

		UpperCamelCase     :					str			       =		(
		    (
		        MobileBertModel,
		        MobileBertForMaskedLM,
		        MobileBertForMultipleChoice,
		        MobileBertForNextSentencePrediction,
		        MobileBertForPreTraining,
		        MobileBertForQuestionAnswering,
		        MobileBertForSequenceClassification,
		        MobileBertForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		UpperCamelCase     :					Optional[int]			       =		(
		    {
		        "feature-extraction": MobileBertModel,
		        "fill-mask": MobileBertForMaskedLM,
		        "question-answering": MobileBertForQuestionAnswering,
		        "text-classification": MobileBertForSequenceClassification,
		        "token-classification": MobileBertForTokenClassification,
		        "zero-shot": MobileBertForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		UpperCamelCase     :					int			       =		True







		def 					__A   ( self		,      A		,      A		,      A=False	)  -> Any:







					'''simple docstring'''







					lowerCamelCase    		=  super()._prepare_for_class(A		,      A		,      return_labels=A	)

					if return_labels:
								if model_class in get_values(A	):
											lowerCamelCase    		=  torch.zeros(
											    (self.model_tester.batch_size, self.model_tester.seq_length)		,      dtype=torch.long		,      device=A	)
											lowerCamelCase    		=  torch.zeros(
											    self.model_tester.batch_size		,      dtype=torch.long		,      device=A	)
					return inputs_dict







		def 					__A   ( self	)  -> Union[str, Any]:







					'''simple docstring'''







					lowerCamelCase    		=  MobileBertModelTester(self	)
					lowerCamelCase    		=  ConfigTester(self		,      config_class=A		,      hidden_size=37	)







		def 					__A   ( self	)  -> Tuple:







					'''simple docstring'''







					self.config_tester.run_common_tests()







		def 					__A   ( self	)  -> int:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_mobilebert_model(*A	)







		def 					__A   ( self	)  -> Any:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_mobilebert_for_masked_lm(*A	)







		def 					__A   ( self	)  -> Any:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_mobilebert_for_multiple_choice(*A	)







		def 					__A   ( self	)  -> int:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*A	)







		def 					__A   ( self	)  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_mobilebert_for_pretraining(*A	)







		def 					__A   ( self	)  -> Tuple:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_mobilebert_for_question_answering(*A	)







		def 					__A   ( self	)  -> Optional[int]:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_mobilebert_for_sequence_classification(*A	)







		def 					__A   ( self	)  -> str:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_mobilebert_for_token_classification(*A	)






def 	__lowerCamelCase						(							lowerCamelCase__		:						Optional[Any]      ):





			'''simple docstring'''




			return torch.tensor(
			    lowerCamelCase__  ,							dtype=torch.long  ,							device=lowerCamelCase__  ,							)


UpperCAmelCase   :    Optional[Any]     =		1e-3







@require_torch
@require_sentencepiece
@require_tokenizers
class        __lowercase       (							unittest.TestCase  ):



		"""simple docstring"""

		@slow
		def 					__A   ( self	)  -> int:







					'''simple docstring'''







					lowerCamelCase    		=  MobileBertModel.from_pretrained("""google/mobilebert-uncased"""	).to(A	)
					lowerCamelCase    		=  _long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]	)
					with torch.no_grad():
								lowerCamelCase    		=  model(A	)[0]
					lowerCamelCase    		=  torch.Size((1, 9, 5_12)	)
					self.assertEqual(output.shape		,      A	)
					lowerCamelCase    		=  torch.tensor(
					    [
					        [
					            [-2.4_7_3_6_5_2_6e0_7, 8.2_6_9_1_6_5_6e0_4, 1.6_5_2_1_8_3_8e0_5],
					            [-5.7_5_4_1_7_0_4e-0_1, 3.9_0_5_6_0_2_2e0_0, 4.4_0_1_1_5_0_7e0_0],
					            [2.6_0_4_7_3_5_9e0_0, 1.5_6_7_7_6_5_2e0_0, -1.7_3_2_4_1_8_8e-0_1],
					        ]
					    ]		,      device=A		,      )

					# MobileBERT results range from 10e0 to 10e8. Even a 0.0000001% difference with a value of 10e8 results in a
					# ~1 difference, it's therefore not a good idea to measure using addition.
					# Here, we instead divide the expected result with the result in order to obtain ~1. We then check that the
					# result is held between bounds: 1 - TOLERANCE < expected_result / result < 1 + TOLERANCE
					lowerCamelCase    		=  torch.all((expected_slice / output[..., :3, :3]) >= 1 - TOLERANCE	)
					lowerCamelCase    		=  torch.all((expected_slice / output[..., :3, :3]) <= 1 + TOLERANCE	)

					self.assertTrue(lower_bound and upper_bound	)





import logging
from dataclasses import dataclass, field
from typing import Optional

from seqaseq_trainer import arg_to_scheduler

from transformers import TrainingArguments


UpperCAmelCase   :    Optional[Any]     =		logging.getLogger(__name__)







@dataclass
class        __lowercase       (							a_  ):



		"""simple docstring"""

		UpperCamelCase     :					Optional[float]			       =		field(
		    default=0.0			,							metadata={"help": "The label smoothing epsilon to apply (if not zero)."}  )
		UpperCamelCase     :					bool			       =		field(default=a_			,							metadata={"help": "Whether to SortishSamler or not."}  )
		UpperCamelCase     :					bool			       =		field(
		    default=a_			,							metadata={"help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."}  )
		UpperCamelCase     :					bool			       =		field(default=a_			,							metadata={"help": "whether to use adafactor"}  )
		UpperCamelCase     :					Optional[float]			       =		field(
		    default=a_			,							metadata={"help": "Encoder layer dropout probability. Goes into model.config."}  )
		UpperCamelCase     :					Optional[float]			       =		field(
		    default=a_			,							metadata={"help": "Decoder layer dropout probability. Goes into model.config."}  )
		UpperCamelCase     :					Optional[float]			       =		field(default=a_			,							metadata={"help": "Dropout probability. Goes into model.config."}  )
		UpperCamelCase     :					Optional[float]			       =		field(
		    default=a_			,							metadata={"help": "Attention dropout probability. Goes into model.config."}  )
		UpperCamelCase     :					Optional[str]			       =		field(
		    default="linear"			,							metadata={"help": f'''Which lr scheduler to use. Selected in {sorted(arg_to_scheduler.keys()  )}'''}			,							)

