
"""simple docstring"""







import math
import os
import re
import sys
import unittest
from pathlib import Path
from typing import Tuple
from unittest.mock import patch

from parameterized import parameterized

from transformers.testing_utils import (
    CaptureStderr,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_torch_dist_unique_port,
    require_apex,
    require_bitsandbytes,
    require_fairscale,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    require_torch_non_multi_gpu,
    slow,
)
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import set_seed


_SCREAMING_SNAKE_CASE	       =      os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(F"""{bindir}/../../examples/pytorch/translation"""):
				from run_translation import main  # noqa


set_seed(42)
_SCREAMING_SNAKE_CASE	       =      "sshleifer/student_marian_en_ro_6_1"
_SCREAMING_SNAKE_CASE	       =      "sshleifer/tiny-mbart"







@require_torch
class 							__magic_name__     (		_lowerCamelCase   ):





					def 							lowerCAmelCase      (				self   :							Tuple							,						snake_case_   :							List[Any]=False							,						snake_case_   :							List[Any]=None							,						snake_case_   :							Union[str, Any]=True							,						snake_case_   :							int=True							,						snake_case_   :							int=True							,						snake_case_   :							List[str]=True							,						):
												__snake_case     			=       self.run_trainer(
												    eval_steps=1							,						max_len=12							,						model_name=A__							,						num_train_epochs=1							,						distributed=A__							,						extra_args_str=A__							,						predict_with_generate=A__							,						do_train=A__							,						do_eval=A__							,						do_predict=A__							,						)
												__snake_case     			=       TrainerState.load_from_json(os.path.join(A__							,						"trainer_state.json"    )    ).log_history

												if not do_eval:
																			return

												__snake_case     			=       [log for log in logs if "eval_loss" in log.keys()]

												__snake_case     			=       eval_metrics[0]
												if predict_with_generate:
																			assert "eval_bleu" in first_step_stats

																			__snake_case     			=       eval_metrics[-1]
																			assert isinstance(last_step_stats["eval_bleu"]							,						A__    )
																			assert not math.isnan(float(last_step_stats["eval_loss"]    )    ), "eval_loss must not be `nan`"





					@require_torch_non_multi_gpu
					def 							lowerCAmelCase      (				self   :							Tuple    ):
												self.run_seqaseq_quick()





					@require_torch_multi_gpu
					def 							lowerCAmelCase      (				self   :							List[str]    ):
												self.run_seqaseq_quick(distributed=A__    )





					@require_torch_multi_gpu
					def 							lowerCAmelCase      (				self   :							Tuple    ):
												self.run_seqaseq_quick(distributed=A__    )





					@unittest.skip("Requires an update of the env running those tests"    )
					@require_torch_multi_gpu
					@require_fairscale
					def 							lowerCAmelCase      (				self   :							Optional[Any]    ):
												self.run_seqaseq_quick(distributed=A__							,						extra_args_str="--sharded_ddp simple"    )





					@unittest.skip("Requires an update of the env running those tests"    )
					@require_torch_multi_gpu
					@require_fairscale
					def 							lowerCAmelCase      (				self   :							str    ):
												self.run_seqaseq_quick(distributed=A__							,						extra_args_str="--sharded_ddp simple --fp16"    )





					@unittest.skip("Requires an update of the env running those tests"    )
					@require_torch_multi_gpu
					@require_fairscale
					def 							lowerCAmelCase      (				self   :							List[Any]    ):
												self.run_seqaseq_quick(distributed=A__							,						extra_args_str="--sharded_ddp zero_dp_2"							,						predict_with_generate=A__    )





					@unittest.skip("Requires an update of the env running those tests"    )
					@require_torch_multi_gpu
					@require_fairscale
					def 							lowerCAmelCase      (				self   :							str    ):
												self.run_seqaseq_quick(
												    distributed=A__							,						extra_args_str="--sharded_ddp zero_dp_2 --fp16"							,						predict_with_generate=A__    )





					@require_apex
					@require_torch_gpu
					def 							lowerCAmelCase      (				self   :							List[Any]    ):
												# XXX: apex breaks the trainer if it's run twice e.g. run_seq2seq.main() from the same
												# program and it breaks other tests that run from the same pytest worker, therefore until this is
												# sorted out it must be run only in an external program, that is distributed=True in this
												# test and only under one or more gpus - if we want cpu will need to make a special test
												#
												# specifically to the problem traced it to self.optimizer.step() - if it's run 2nd time via
												# 2nd main() call it botches the future eval.
												#
												self.run_seqaseq_quick(distributed=A__							,						extra_args_str="--fp16 --fp16_backend=apex"    )
												# test 2nd time - was getting eval_loss': nan'
												# to reproduce the problem set distributed=False
												self.run_seqaseq_quick(distributed=A__							,						extra_args_str="--fp16 --fp16_backend=apex"    )





					@parameterized.expand(["base", "low", "high", "mixed"]    )
					@require_torch_multi_gpu
					def 							lowerCAmelCase      (				self   :							int							,						snake_case_   :							List[Any]    ):
												# as each sub-test is slow-ish split into multiple sub-tests to avoid CI timeout
												__snake_case     			=       {
												    # test with the default log_level - should be info and thus log info once
												    "base": {"extra_args_str": "", "n_matches": 1},
												    # test with low log_level and log_level_replica - should be noisy on all processes
												    # now the info string should appear twice on 2 processes
												    "low": {"extra_args_str": "--log_level debug --log_level_replica debug", "n_matches": 2},
												    # test with high log_level and low log_level_replica
												    # now the info string should appear once only on the replica
												    "high": {"extra_args_str": "--log_level error --log_level_replica debug", "n_matches": 1},
												    # test with high log_level and log_level_replica - should be quiet on all processes
												    "mixed": {"extra_args_str": "--log_level error --log_level_replica error", "n_matches": 0},
												}

												__snake_case     			=       experiments[experiment_id]
												__snake_case     			=       {"distributed": True, "predict_with_generate": False, "do_eval": False, "do_predict": False}
												__snake_case     			=       "Running training"
												with CaptureStderr() as cl:
																			self.run_seqaseq_quick(**A__							,						extra_args_str=data["extra_args_str"]    )
												__snake_case     			=       len(re.findall(A__							,						cl.err    )    )
												self.assertEqual(A__							,						data["n_matches"]    )





					@slow
					def 							lowerCAmelCase      (				self   :							Dict    ):
												__snake_case     			=       self.run_trainer(
												    eval_steps=2							,						max_len=128							,						model_name=A__							,						learning_rate=3e-4							,						num_train_epochs=10							,						distributed=A__							,						)

												# Check metrics
												__snake_case     			=       TrainerState.load_from_json(os.path.join(A__							,						"trainer_state.json"    )    ).log_history
												__snake_case     			=       [log for log in logs if "eval_loss" in log.keys()]
												__snake_case     			=       eval_metrics[0]
												__snake_case     			=       eval_metrics[-1]

												assert first_step_stats["eval_loss"] > last_step_stats["eval_loss"], "model learned nothing"
												assert isinstance(last_step_stats["eval_bleu"]							,						A__    )

												# test if do_predict saves generations and metrics
												__snake_case     			=       os.listdir(A__    )
												__snake_case     			=       {os.path.basename(A__    ) for p in contents}
												assert "generated_predictions.txt" in contents
												assert "predict_results.json" in contents





					@slow
					@require_bitsandbytes
					def 							lowerCAmelCase      (				self   :							List[str]    ):
												from transformers.training_args import OptimizerNames

												def train_and_return_metrics(snake_case_   :							Dict    ) -> Tuple[int, float]:
																			__snake_case     			=       "--skip_memory_metrics 0"

																			__snake_case     			=       self.run_trainer(
																			    max_len=128							,						model_name=A__							,						learning_rate=3e-4							,						num_train_epochs=1							,						optim=A__							,						distributed=A__							,						extra_args_str=A__							,						do_eval=A__							,						do_predict=A__							,						n_gpus_to_use=1							,						)

																			# Check metrics
																			__snake_case     			=       TrainerState.load_from_json(Path(A__							,						"trainer_state.json"    )    ).log_history
																			__snake_case     			=       int(logs[0]["train_mem_gpu_peaked_delta"] / 2**20    )
																			__snake_case     			=       int(logs[0]["train_mem_gpu_alloc_delta"] / 2**20    )

																			__snake_case     			=       logs[0]["train_loss"]
																			return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

												__snake_case			,   __snake_case			,   __snake_case     			=       train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value    )
												__snake_case			,   __snake_case			,   __snake_case     			=       train_and_return_metrics(OptimizerNames.ADAMW_BNB.value    )

												__snake_case     			=       gpu_alloc_mem_orig - gpu_alloc_mem_bnb

												__snake_case     			=       gpu_peak_mem_orig + gpu_alloc_mem_orig
												__snake_case     			=       gpu_peak_mem_bnb + gpu_alloc_mem_bnb
												__snake_case     			=       gpu_total_mem_orig - gpu_total_mem_bnb

												# sshleifer/student_marian_en_ro_6_1 has 54M parameter, 29M of which is `nn.Embedding` which
												# doesn't get quantized and remains in fp32. Therefore we only have 25M parameters quantized
												# in 2 bytes and the diff in optim memory usage is derived as so:
												#
												# - normal 25*8=~200MB (8 bytes per param)
												# - bnb    25*2= ~50MB (2 bytes per param)
												#
												# Thus we should expect ~150MB total memory saved.
												#
												# Peak memory should be the same - the total should be different by about that same margin
												#
												# After leaving a small margin to accommodate for differences between gpus let's check
												# that we have at least 120MB in savings
												__snake_case     			=       120

												# uncomment the following if this test starts failing - requires py38 for a new print feature
												# gpu_peak_mem_diff = gpu_peak_mem_orig - gpu_peak_mem_bnb
												# print(f"{gpu_alloc_mem_orig=}MB {gpu_peak_mem_orig=}MB {gpu_alloc_mem_orig+gpu_peak_mem_orig=}MB")
												# print(f" {gpu_alloc_mem_bnb=}MB  {gpu_peak_mem_bnb=}MB  {gpu_alloc_mem_bnb+gpu_peak_mem_bnb=}MB")
												# print(f"{gpu_alloc_mem_diff=}MB")
												# print(f"{gpu_peak_mem_diff=}MB")
												# print(f"{gpu_total_mem_orig=}MB, {gpu_total_mem_bnb=}MB")
												# print(f"{gpu_total_mem_diff=}MB, {gpu_total_mem_diff=}MB")

												self.assertGreater(
												    A__							,						A__							,						"should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got"
												    F''' a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and'''
												    F''' gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB'''							,						)

												self.assertGreater(
												    A__							,						A__							,						"should use ~150MB less total gpu memory with BNB, compared to without it for this model but got"
												    F''' a difference of {gpu_total_mem_diff}MB, with gpu_total_mem_orig={gpu_total_mem_orig}MB and'''
												    F''' gpu_total_mem_bnb={gpu_total_mem_bnb}MB'''							,						)

												self.assertEqual(
												    A__							,						A__							,						F'''loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}'''    )





					def 							lowerCAmelCase      (				self   :							Optional[int]							,						snake_case_   :							Any							,						snake_case_   :							List[Any]							,						snake_case_   :							List[Any]							,						snake_case_   :							Tuple = 3e-3							,						snake_case_   :							Tuple = "adafactor"							,						snake_case_   :							Optional[int] = False							,						snake_case_   :							int = None							,						snake_case_   :							Optional[int] = 0							,						snake_case_   :							Any = True							,						snake_case_   :							Dict = True							,						snake_case_   :							Dict = True							,						snake_case_   :							Union[str, Any] = True							,						snake_case_   :							Any = None							,						):
												__snake_case     			=       self.test_file_dir / "../fixtures/tests_samples/wmt_en_ro"
												__snake_case     			=       self.get_auto_remove_tmp_dir()
												__snake_case     			=       F'''\n            --model_name_or_path {model_name}\n            --train_file {data_dir}/train.json\n            --validation_file {data_dir}/val.json\n            --test_file {data_dir}/test.json\n            --output_dir {output_dir}\n            --overwrite_output_dir\n            --max_train_samples 8\n            --max_source_length {max_len}\n            --max_target_length {max_len}\n            --do_train\n            --num_train_epochs {str(A__    )}\n            --per_device_train_batch_size 4\n            --learning_rate {learning_rate}\n            --warmup_steps 8\n            --logging_steps 0\n            --logging_strategy no\n            --save_steps {str(A__    )}\n            --group_by_length\n            --label_smoothing_factor 0.1\n            --target_lang ro_RO\n            --source_lang en_XX\n        '''.split()

												__snake_case     			=       F'''\n            --do_eval\n            --per_device_eval_batch_size 4\n            --max_eval_samples 8\n            --val_max_target_length {max_len}\n            --evaluation_strategy steps\n            --eval_steps {str(A__    )}\n        '''.split()

												__snake_case     			=       "\n            --do_predict\n        ".split()

												__snake_case     			=       []
												if do_train:
																			args += args_train

												if do_eval:
																			args += args_eval

												if do_predict:
																			args += args_predict

												if predict_with_generate:
																			args += "--predict_with_generate".split()

												if do_train:
																			if optim == "adafactor":
																										args += "--adafactor".split()
																			else:
																										args += F'''--optim {optim}'''.split()

												if extra_args_str is not None:
																			args += extra_args_str.split()

												if distributed:
																			if n_gpus_to_use is None:
																										__snake_case     			=       get_gpu_count()
																			__snake_case     			=       get_torch_dist_unique_port()
																			__snake_case     			=       F'''\n                -m torch.distributed.run\n                --nproc_per_node={n_gpus_to_use}\n                --master_port={master_port}\n                {self.examples_dir_str}/pytorch/translation/run_translation.py\n            '''.split()
																			__snake_case     			=       [sys.executable] + distributed_args + args
																			# keep for quick debug
																			# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
																			execute_subprocess_async(A__							,						env=self.get_env()    )
												else:
																			__snake_case     			=       ["run_translation.py"] + args
																			with patch.object(A__							,						"argv"							,						A__    ):
																										main()

												return output_dir


"""simple docstring"""







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 							__magic_name__     :





					def __init__(				self   :							Optional[int]							,						snake_case_   :							List[str]							,						snake_case_   :							str=13							,						snake_case_   :							Optional[int]=7							,						snake_case_   :							Optional[Any]=True							,						snake_case_   :							Union[str, Any]=True							,						snake_case_   :							Tuple=True							,						snake_case_   :							int=True							,						snake_case_   :							int=99							,						snake_case_   :							Optional[int]=64							,						snake_case_   :							Dict=32							,						snake_case_   :							Dict=5							,						snake_case_   :							List[str]=4							,						snake_case_   :							List[Any]=37							,						snake_case_   :							Optional[Any]="gelu"							,						snake_case_   :							Optional[Any]=0.1							,						snake_case_   :							List[Any]=0.1							,						snake_case_   :							Union[str, Any]=512							,						snake_case_   :							int=16							,						snake_case_   :							List[str]=2							,						snake_case_   :							Union[str, Any]=0.02							,						snake_case_   :							Union[str, Any]=3							,						snake_case_   :							str=4							,						snake_case_   :							int=None							,						):
												__snake_case     			=       parent
												__snake_case     			=       batch_size
												__snake_case     			=       seq_length
												__snake_case     			=       is_training
												__snake_case     			=       use_input_mask
												__snake_case     			=       use_token_type_ids
												__snake_case     			=       use_labels
												__snake_case     			=       vocab_size
												__snake_case     			=       hidden_size
												__snake_case     			=       embedding_size
												__snake_case     			=       num_hidden_layers
												__snake_case     			=       num_attention_heads
												__snake_case     			=       intermediate_size
												__snake_case     			=       hidden_act
												__snake_case     			=       hidden_dropout_prob
												__snake_case     			=       attention_probs_dropout_prob
												__snake_case     			=       max_position_embeddings
												__snake_case     			=       type_vocab_size
												__snake_case     			=       type_sequence_label_size
												__snake_case     			=       initializer_range
												__snake_case     			=       num_labels
												__snake_case     			=       num_choices
												__snake_case     			=       scope





					def 							lowerCAmelCase      (				self   :							Union[str, Any]    ):
												__snake_case     			=       ids_tensor([self.batch_size, self.seq_length]							,						self.vocab_size    )

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

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

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

												__snake_case     			=       self.get_config()

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





					def 							lowerCAmelCase      (				self   :							Union[str, Any]    ):
												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=snake_case_							,						initializer_range=self.initializer_range							,						)





					def 							lowerCAmelCase      (				self   :							Optional[Any]							,						snake_case_   :							Any							,						snake_case_   :							Optional[Any]							,						snake_case_   :							Tuple							,						snake_case_   :							Any							,						snake_case_   :							str							,						snake_case_   :							Dict							,						snake_case_   :							List[str]    ):
												__snake_case     			=       MobileBertModel(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_    )
												__snake_case     			=       model(snake_case_							,						token_type_ids=snake_case_    )
												__snake_case     			=       model(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 							lowerCAmelCase      (				self   :							List[str]							,						snake_case_   :							List[Any]							,						snake_case_   :							Any							,						snake_case_   :							Optional[int]							,						snake_case_   :							int							,						snake_case_   :							int							,						snake_case_   :							str							,						snake_case_   :							str    ):
												__snake_case     			=       MobileBertForMaskedLM(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_    )
												self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.seq_length, self.vocab_size)    )





					def 							lowerCAmelCase      (				self   :							Optional[Any]							,						snake_case_   :							Optional[int]							,						snake_case_   :							Union[str, Any]							,						snake_case_   :							List[str]							,						snake_case_   :							int							,						snake_case_   :							int							,						snake_case_   :							int							,						snake_case_   :							Any    ):
												__snake_case     			=       MobileBertForNextSentencePrediction(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(
												    snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_							,						)
												self.parent.assertEqual(result.logits.shape							,						(self.batch_size, 2)    )





					def 							lowerCAmelCase      (				self   :							Optional[int]							,						snake_case_   :							List[str]							,						snake_case_   :							List[str]							,						snake_case_   :							str							,						snake_case_   :							Optional[Any]							,						snake_case_   :							List[str]							,						snake_case_   :							List[str]							,						snake_case_   :							Dict    ):
												__snake_case     			=       MobileBertForPreTraining(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(
												    snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_							,						next_sentence_label=snake_case_							,						)
												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 							lowerCAmelCase      (				self   :							Any							,						snake_case_   :							List[str]							,						snake_case_   :							int							,						snake_case_   :							Any							,						snake_case_   :							Union[str, Any]							,						snake_case_   :							Dict							,						snake_case_   :							int							,						snake_case_   :							int    ):
												__snake_case     			=       MobileBertForQuestionAnswering(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(
												    snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						start_positions=snake_case_							,						end_positions=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 							lowerCAmelCase      (				self   :							str							,						snake_case_   :							Tuple							,						snake_case_   :							int							,						snake_case_   :							str							,						snake_case_   :							Dict							,						snake_case_   :							Optional[int]							,						snake_case_   :							str							,						snake_case_   :							List[str]    ):
												__snake_case     			=       self.num_labels
												__snake_case     			=       MobileBertForSequenceClassification(snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_    )
												self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.num_labels)    )





					def 							lowerCAmelCase      (				self   :							Optional[Any]							,						snake_case_   :							Tuple							,						snake_case_   :							str							,						snake_case_   :							Union[str, Any]							,						snake_case_   :							Tuple							,						snake_case_   :							Optional[int]							,						snake_case_   :							Tuple							,						snake_case_   :							List[Any]    ):
												__snake_case     			=       self.num_labels
												__snake_case     			=       MobileBertForTokenClassification(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       model(snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_    )
												self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.seq_length, self.num_labels)    )





					def 							lowerCAmelCase      (				self   :							Tuple							,						snake_case_   :							Union[str, Any]							,						snake_case_   :							Tuple							,						snake_case_   :							Optional[Any]							,						snake_case_   :							Any							,						snake_case_   :							Any							,						snake_case_   :							Dict							,						snake_case_   :							Union[str, Any]    ):
												__snake_case     			=       self.num_choices
												__snake_case     			=       MobileBertForMultipleChoice(config=snake_case_    )
												model.to(snake_case_    )
												model.eval()
												__snake_case     			=       input_ids.unsqueeze(1    ).expand(-1							,						self.num_choices							,						-1    ).contiguous()
												__snake_case     			=       token_type_ids.unsqueeze(1    ).expand(-1							,						self.num_choices							,						-1    ).contiguous()
												__snake_case     			=       input_mask.unsqueeze(1    ).expand(-1							,						self.num_choices							,						-1    ).contiguous()
												__snake_case     			=       model(
												    snake_case_							,						attention_mask=snake_case_							,						token_type_ids=snake_case_							,						labels=snake_case_							,						)
												self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.num_choices)    )





					def 							lowerCAmelCase      (				self   :							List[Any]    ):
												__snake_case     			=       self.prepare_config_and_inputs()
												(
												    (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   (
												    __snake_case
												)			,   
												)     			=       config_and_inputs
												__snake_case     			=       {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
												return config, inputs_dict







@require_torch
class 							__magic_name__     (		lowercase__      ,		lowercase__      ,		unittest.TestCase   ):
					_SCREAMING_SNAKE_CASE       :					Dict     							=   (
					    (
					        MobileBertModel,
					        MobileBertForMaskedLM,
					        MobileBertForMultipleChoice,
					        MobileBertForNextSentencePrediction,
					        MobileBertForPreTraining,
					        MobileBertForQuestionAnswering,
					        MobileBertForSequenceClassification,
					        MobileBertForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					_SCREAMING_SNAKE_CASE       :					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 {}
					)
					_SCREAMING_SNAKE_CASE       :					int     							=   True





					def 							lowerCAmelCase      (				self   :							List[str]							,						snake_case_   :							List[str]							,						snake_case_   :							str							,						snake_case_   :							str=False    ):
												__snake_case     			=       super()._prepare_for_class(snake_case_							,						snake_case_							,						return_labels=snake_case_    )

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





					def 							lowerCAmelCase      (				self   :							Optional[Any]    ):
												__snake_case     			=       MobileBertModelTester(self    )
												__snake_case     			=       ConfigTester(self							,						config_class=snake_case_							,						hidden_size=37    )





					def 							lowerCAmelCase      (				self   :							List[Any]    ):
												self.config_tester.run_common_tests()





					def 							lowerCAmelCase      (				self   :							Tuple    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_model(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							Optional[Any]    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_masked_lm(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							Tuple    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_multiple_choice(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							Any    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							Any    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_pretraining(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							List[str]    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_question_answering(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							List[Any]    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_sequence_classification(*snake_case_    )





					def 							lowerCAmelCase      (				self   :							str    ):
												__snake_case     			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_token_classification(*snake_case_    )




def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE )  ->     List[Any]:







							"""simple docstring"""


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


_SCREAMING_SNAKE_CASE	       =      1E-3







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





					@slow
					def 							lowerCAmelCase      (				self   :							Tuple    ):
												__snake_case     			=       MobileBertModel.from_pretrained("google/mobilebert-uncased"    ).to(snake_case_    )
												__snake_case     			=       _long_tensor([[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]    )
												with torch.no_grad():
																			__snake_case     			=       model(snake_case_    )[0]
												__snake_case     			=       torch.Size((1, 9, 512)    )
												self.assertEqual(output.shape							,						snake_case_    )
												__snake_case     			=       torch.tensor(
												    [
												        [
												            [-2.4_73_65_26e07, 8.2_69_16_56e04, 1.6_52_18_38e05],
												            [-5.7_54_17_04e-01, 3.9_05_60_22e00, 4.4_01_15_07e00],
												            [2.6_04_73_59e00, 1.5_67_76_52e00, -1.7_32_41_88e-01],
												        ]
												    ]							,						device=snake_case_							,						)

												# 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
												__snake_case     			=       torch.all((expected_slice / output[..., :3, :3]) >= 1 - TOLERANCE    )
												__snake_case     			=       torch.all((expected_slice / output[..., :3, :3]) <= 1 + TOLERANCE    )

												self.assertTrue(lower_bound and upper_bound    )
