
"""simple docstring"""




import logging
import sys
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

import librosa
import torch
from datasets import DatasetDict, load_dataset
from packaging import version
from torch import nn

from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaForPreTraining,
    is_apex_available,
    trainer_utils,
)
from transformers.models.wavaveca.modeling_wavaveca import _compute_mask_indices


if is_apex_available():
  from apex import amp

if version.parse(version.parse(torch.__version__).base_version) >= version.parse('1.6'):
  A__   :					str							=   True
  from torch.cuda.amp import autocast


A__   :					int							=   logging.getLogger(__name__)


@dataclass
class 			__magic_name__       :
       UpperCamelCase_   =		field(
           metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''} )
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''}   ,					)
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''Whether to freeze the feature extractor layers of the model.'''} )
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''Whether to log verbose messages or not.'''}   ,					)
       UpperCamelCase_   =		field(
           default=2.0   ,					metadata={'''help''': '''Maximum temperature for gumbel softmax.'''} )
       UpperCamelCase_   =		field(
           default=0.5   ,					metadata={'''help''': '''Minimum temperature for gumbel softmax.'''} )
       UpperCamelCase_   =		field(
           default=0.9_9_9_9_9_5   ,					metadata={'''help''': '''Decay of gumbel temperature during training.'''} )




def  _lowerCAmelCase		(		_UpperCamelCase ,				_UpperCamelCase						):






  """simple docstring"""



  logging.basicConfig(
      format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s''' ,				datefmt='''%m/%d/%Y %H:%M:%S''' ,				handlers=[logging.StreamHandler(sys.stdout						)] ,				)
  _lowercase:			Optional[Any]       =     logging.WARNING
  if model_args.verbose_logging:
    _lowercase:			List[str]       =     logging.DEBUG
  elif trainer_utils.is_main_process(training_args.local_rank						):
    _lowercase:			Optional[Any]       =     logging.INFO
  logger.setLevel(_UpperCamelCase						)


@dataclass
class 			__magic_name__       :
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''The name of the dataset to use (via the datasets library).'''} )
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''The configuration name of the dataset to use (via the datasets library).'''} )
       UpperCamelCase_   =		field(
           default='''train'''   ,					metadata={
               '''help''': '''The name of the training data set split to use (via the datasets library). Defaults to \'train\''''
           }   ,					)
       UpperCamelCase_   =		field(
           default='''validation'''   ,					metadata={
               '''help''': (
                   '''The name of the validation data set split to use (via the datasets library). Defaults to \'validation\''''
               )
           }   ,					)
       UpperCamelCase_   =		field(
           default='''file'''   ,					metadata={'''help''': '''Column in the dataset that contains speech file path. Defaults to \'file\''''}   ,					)
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''Overwrite the cached preprocessed datasets or not.'''} )
       UpperCamelCase_   =		field(
           default=1   ,					metadata={
               '''help''': '''The percentage of the train set used as validation set in case there\'s no validation split'''
           }   ,					)
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''The number of processes to use for the preprocessing.'''}   ,					)
       UpperCamelCase_   =		field(
           default=2_0.0   ,					metadata={'''help''': '''Filter audio files that are longer than `max_duration_in_seconds` seconds'''} )


@dataclass
class 			__magic_name__       :
       UpperCamelCase_   =		42
       UpperCamelCase_   =		42
       UpperCamelCase_   =		"longest"
       UpperCamelCase_   =		None
       UpperCamelCase_   =		None
       def __call__(     self ,				A_       )     ->	Dict[str, torch.Tensor]:





         """simple docstring"""




         _lowercase:			List[Any]       =     self.feature_extractor.pad(
             A_ ,				max_length=self.max_length ,				padding=self.padding ,				pad_to_multiple_of=self.pad_to_multiple_of ,				return_tensors='''pt''' ,				)
         _lowercase:			Tuple       =     self.model._get_feat_extract_output_lengths(batch['''input_values'''].shape[-1]       )

         _lowercase:			List[str]       =     batch['''input_values'''].shape[0]

         # make sure that no loss is computed on padded inputs
         if batch["attention_mask"] is not None:
           # compute real output lengths according to convolution formula
           _lowercase:			Any       =     self.model._get_feat_extract_output_lengths(batch['''attention_mask'''].sum(-1       )       ).to(
               torch.long       )

           _lowercase:			int       =     torch.zeros(
               (batch_size, mask_indices_seq_length) ,				dtype=torch.long ,				device=batch['''input_values'''].device       )

           # these two operations makes sure that all values
           # before the output lengths indices are attended to
           _lowercase:			int       =     1
           _lowercase:			Optional[Any]       =     attention_mask.flip([-1]       ).cumsum(-1       ).flip([-1]       ).bool()

         # sample randomly masked indices
         _lowercase:			Optional[int]       =     _compute_mask_indices(
             (batch_size, mask_indices_seq_length) ,				self.model.config.mask_time_prob ,				self.model.config.mask_time_length ,				attention_mask=A_ ,				min_masks=2 ,				)

         return batch


class 			__magic_name__       (       SCREAMING_SNAKE_CASE__ ):
       def __init__(     self ,				*A_ ,				A_=1 ,				A_=0 ,				A_=1.0 ,				**A_       )     ->	Tuple:





         """simple docstring"""




         super().__init__(*A_ ,				**A_       )
         _lowercase:			Any       =     0
         _lowercase:			List[str]       =     max_gumbel_temp
         _lowercase:			Union[str, Any]       =     min_gumbel_temp
         _lowercase:			Any       =     gumbel_temp_decay



       def 							lowercase_							(     self ,				A_ ,				A_       )     ->	torch.Tensor:





         """simple docstring"""




         model.train()
         _lowercase:			Tuple       =     self._prepare_inputs(A_       )

         if self.use_amp:
           with autocast():
             _lowercase:			Tuple       =     self.compute_loss(A_ ,				A_       )
         else:
           _lowercase:			Optional[int]       =     self.compute_loss(A_ ,				A_       )

         if self.args.n_gpu > 1 or self.deepspeed:
           if model.module.config.ctc_loss_reduction == "mean":
             _lowercase:			Tuple       =     loss.mean()
           elif model.module.config.ctc_loss_reduction == "sum":
             _lowercase:			Optional[Any]       =     loss.sum() / (inputs['''mask_time_indices''']).sum()
           else:
             raise ValueError(f'''{model.config.ctc_loss_reduction} is not valid. Choose one of [\'mean\', \'sum\']'''       )

         if self.args.gradient_accumulation_steps > 1:
           _lowercase:			Tuple       =     loss / self.args.gradient_accumulation_steps

         if self.use_amp:
           self.scaler.scale(A_       ).backward()
         elif self.use_apex:
           with amp.scale_loss(A_ ,				self.optimizer       ) as scaled_loss:
             scaled_loss.backward()
         elif self.deepspeed:
           self.deepspeed.backward(A_       )
         else:
           loss.backward()

         self.num_update_step += 1
         # make sure gumbel softmax temperature is decayed
         if self.args.n_gpu > 1 or self.deepspeed:
           model.module.set_gumbel_temperature(
               max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step ,				self.min_gumbel_temp       )       )
         else:
           model.set_gumbel_temperature(
               max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step ,				self.min_gumbel_temp       )       )

         return loss.detach()




def  _lowerCAmelCase		(		):






  """simple docstring"""



  _lowercase:			str       =     HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)						)

  _lowercase ,				_lowercase ,				_lowercase:			Dict       =     parser.parse_args_into_dataclasses()
  configure_logger(_UpperCamelCase ,				_UpperCamelCase						)

  # Downloading and loading a dataset from the hub.
  _lowercase:			List[str]       =     load_dataset(data_args.dataset_name ,				data_args.dataset_config_name ,				cache_dir=model_args.cache_dir						)

  if "validation" not in datasets.keys():
    # make sure only "validation" and "train" keys remain"
    _lowercase:			str       =     DatasetDict()
    _lowercase:			str       =     load_dataset(
        data_args.dataset_name ,				data_args.dataset_config_name ,				split=f'''{data_args.train_split_name}[:{data_args.validation_split_percentage}%]''' ,				cache_dir=model_args.cache_dir ,				)
    _lowercase:			List[Any]       =     load_dataset(
        data_args.dataset_name ,				data_args.dataset_config_name ,				split=f'''{data_args.train_split_name}[{data_args.validation_split_percentage}%:]''' ,				cache_dir=model_args.cache_dir ,				)
  else:
    # make sure only "validation" and "train" keys remain"
    _lowercase:			Optional[int]       =     DatasetDict()
    _lowercase:			Any       =     load_dataset(
        data_args.dataset_name ,				data_args.dataset_config_name ,				split='''validation''' ,				cache_dir=model_args.cache_dir ,				)
    _lowercase:			List[Any]       =     load_dataset(
        data_args.dataset_name ,				data_args.dataset_config_name ,				split=f'''{data_args.train_split_name}''' ,				cache_dir=model_args.cache_dir ,				)

  # only normalized-inputs-training is supported
  _lowercase:			List[Any]       =     WavaVecaFeatureExtractor.from_pretrained(
      model_args.model_name_or_path ,				cache_dir=model_args.cache_dir ,				do_normalize=_UpperCamelCase						)

  def prepare_dataset(_UpperCamelCase						):
    # check that all files have the correct sampling rate
    _lowercase ,				_lowercase:			Union[str, Any]       =     librosa.load(batch[data_args.speech_file_column] ,				sr=feature_extractor.sampling_rate						)
    return batch

  # load audio files into numpy arrays
  _lowercase:			List[Any]       =     datasets.map(
      _UpperCamelCase ,				num_proc=data_args.preprocessing_num_workers ,				remove_columns=datasets['''train'''].column_names						)

  # filter audio files that are too long
  _lowercase:			Tuple       =     vectorized_datasets.filter(
      lambda _UpperCamelCase						: len(data['''speech''']						) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate						)						)

  def normalize(_UpperCamelCase						):
    return feature_extractor(batch['''speech'''] ,				sampling_rate=feature_extractor.sampling_rate						)

  # normalize and transform to `BatchFeatures`
  _lowercase:			Optional[Any]       =     vectorized_datasets.map(
      _UpperCamelCase ,				batched=_UpperCamelCase ,				num_proc=data_args.preprocessing_num_workers ,				load_from_cache_file=not data_args.overwrite_cache ,				remove_columns=vectorized_datasets['''train'''].column_names ,				)

  # pretraining is only supported for "newer" stable layer norm architecture
  # apply_spec_augment has to be True, mask_feature_prob has to be 0.0
  _lowercase:			Tuple       =     WavaVecaConfig.from_pretrained(
      model_args.model_name_or_path ,				cache_dir=model_args.cache_dir ,				gradient_checkpointing=training_args.gradient_checkpointing ,				)

  if not config.do_stable_layer_norm or config.feat_extract_norm != "layer":
    raise ValueError(
        '''PreTraining is only supported for ``config.do_stable_layer_norm=True`` and'''
        ''' ``config.feat_extract_norm=\'layer\''''						)

  _lowercase:			Optional[Any]       =     WavaVecaForPreTraining(_UpperCamelCase						)

  _lowercase:			List[str]       =     DataCollatorForWavaVecaPretraining(model=_UpperCamelCase ,				feature_extractor=_UpperCamelCase						)

  _lowercase:			Any       =     WavaVecaPreTrainer(
      model=_UpperCamelCase ,				data_collator=_UpperCamelCase ,				args=_UpperCamelCase ,				train_dataset=vectorized_datasets['''train'''] ,				eval_dataset=vectorized_datasets['''validation'''] ,				tokenizer=_UpperCamelCase ,				max_gumbel_temp=model_args.max_gumbel_temperature ,				min_gumbel_temp=model_args.min_gumbel_temperature ,				gumbel_temp_decay=model_args.gumbel_temperature_decay ,				)
  trainer.train()


if __name__ == "__main__":
  main()


"""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 ,				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_=512 ,				A_=16 ,				A_=2 ,				A_=0.02 ,				A_=3 ,				A_=4 ,				A_=None ,				)     ->	Dict:





         """simple docstring"""




         _lowercase:			List[Any]       =     parent
         _lowercase:			str       =     batch_size
         _lowercase:			List[str]       =     seq_length
         _lowercase:			Any       =     is_training
         _lowercase:			Union[str, Any]       =     use_input_mask
         _lowercase:			Any       =     use_token_type_ids
         _lowercase:			int       =     use_labels
         _lowercase:			Any       =     vocab_size
         _lowercase:			Optional[int]       =     hidden_size
         _lowercase:			Optional[int]       =     embedding_size
         _lowercase:			Union[str, Any]       =     num_hidden_layers
         _lowercase:			Optional[int]       =     num_attention_heads
         _lowercase:			Optional[int]       =     intermediate_size
         _lowercase:			List[str]       =     hidden_act
         _lowercase:			Dict       =     hidden_dropout_prob
         _lowercase:			Any       =     attention_probs_dropout_prob
         _lowercase:			Tuple       =     max_position_embeddings
         _lowercase:			List[str]       =     type_vocab_size
         _lowercase:			Any       =     type_sequence_label_size
         _lowercase:			int       =     initializer_range
         _lowercase:			Any       =     num_labels
         _lowercase:			int       =     num_choices
         _lowercase:			List[str]       =     scope
       def 							lowercase_							(     self       )     ->	Optional[Any]:





         """simple docstring"""




         _lowercase:			Tuple       =     ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size       )

         _lowercase:			Optional[Any]       =     None
         if self.use_input_mask:
           _lowercase:			Any       =     random_attention_mask([self.batch_size, self.seq_length]       )

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

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

         _lowercase:			str       =     self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
       def 							lowercase_							(     self       )     ->	List[Any]:





         """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 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	Any:





         """simple docstring"""




         _lowercase:			Tuple       =     MobileBertModel(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			Tuple       =     model(A_ ,				attention_mask=A_ ,				token_type_ids=A_       )
         _lowercase:			Dict       =     model(A_ ,				token_type_ids=A_       )
         _lowercase:			Optional[int]       =     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 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	int:





         """simple docstring"""




         _lowercase:			int       =     MobileBertForMaskedLM(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			Optional[int]       =     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 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	int:





         """simple docstring"""




         _lowercase:			Any       =     MobileBertForNextSentencePrediction(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			int       =     model(
             A_ ,				attention_mask=A_ ,				token_type_ids=A_ ,				labels=A_ ,				)
         self.parent.assertEqual(result.logits.shape ,				(self.batch_size, 2)       )
       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	Optional[int]:





         """simple docstring"""




         _lowercase:			str       =     MobileBertForPreTraining(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			List[str]       =     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 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	List[str]:





         """simple docstring"""




         _lowercase:			Optional[Any]       =     MobileBertForQuestionAnswering(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			Dict       =     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 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	List[Any]:





         """simple docstring"""




         _lowercase:			Optional[Any]       =     self.num_labels
         _lowercase:			int       =     MobileBertForSequenceClassification(A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			List[Any]       =     model(A_ ,				attention_mask=A_ ,				token_type_ids=A_ ,				labels=A_       )
         self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_labels)       )
       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	Tuple:





         """simple docstring"""




         _lowercase:			Dict       =     self.num_labels
         _lowercase:			Optional[Any]       =     MobileBertForTokenClassification(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			Optional[int]       =     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 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	List[str]:





         """simple docstring"""




         _lowercase:			Union[str, Any]       =     self.num_choices
         _lowercase:			List[Any]       =     MobileBertForMultipleChoice(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			Any       =     input_ids.unsqueeze(1       ).expand(-1 ,				self.num_choices ,				-1       ).contiguous()
         _lowercase:			Tuple       =     token_type_ids.unsqueeze(1       ).expand(-1 ,				self.num_choices ,				-1       ).contiguous()
         _lowercase:			List[str]       =     input_mask.unsqueeze(1       ).expand(-1 ,				self.num_choices ,				-1       ).contiguous()
         _lowercase:			Any       =     model(
             A_ ,				attention_mask=A_ ,				token_type_ids=A_ ,				labels=A_ ,				)
         self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_choices)       )



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





         """simple docstring"""




         _lowercase:			Any       =     self.prepare_config_and_inputs()
         (
             (
             _lowercase
         ) ,				(
             _lowercase
         ) ,				(
             _lowercase
         ) ,				(
             _lowercase
         ) ,				(
             _lowercase
         ) ,				(
             _lowercase
         ) ,				(
             _lowercase
         ) ,				
         ):			Dict       =     config_and_inputs
         _lowercase:			int       =     {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
         return config, inputs_dict


@require_torch
class 			__magic_name__       (       SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__   ,					unittest.TestCase ):
       UpperCamelCase_   =		(
           (
               MobileBertModel,
               MobileBertForMaskedLM,
               MobileBertForMultipleChoice,
               MobileBertForNextSentencePrediction,
               MobileBertForPreTraining,
               MobileBertForQuestionAnswering,
               MobileBertForSequenceClassification,
               MobileBertForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       UpperCamelCase_   =		(
           {
               '''feature-extraction''': MobileBertModel,
               '''fill-mask''': MobileBertForMaskedLM,
               '''question-answering''': MobileBertForQuestionAnswering,
               '''text-classification''': MobileBertForSequenceClassification,
               '''token-classification''': MobileBertForTokenClassification,
               '''zero-shot''': MobileBertForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       UpperCamelCase_   =		True
       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_=False       )     ->	str:





         """simple docstring"""




         _lowercase:			Optional[Any]       =     super()._prepare_for_class(A_ ,				A_ ,				return_labels=A_       )

         if return_labels:
           if model_class in get_values(A_       ):
             _lowercase:			Tuple       =     torch.zeros(
                 (self.model_tester.batch_size, self.model_tester.seq_length) ,				dtype=torch.long ,				device=A_       )
             _lowercase:			Tuple       =     torch.zeros(
                 self.model_tester.batch_size ,				dtype=torch.long ,				device=A_       )
         return inputs_dict
       def 							lowercase_							(     self       )     ->	Tuple:





         """simple docstring"""




         _lowercase:			Optional[Any]       =     MobileBertModelTester(self       )
         _lowercase:			str       =     ConfigTester(self ,				config_class=A_ ,				hidden_size=37       )
       def 							lowercase_							(     self       )     ->	Any:





         """simple docstring"""




         self.config_tester.run_common_tests()
       def 							lowercase_							(     self       )     ->	Union[str, Any]:





         """simple docstring"""




         _lowercase:			str       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_model(*A_       )
       def 							lowercase_							(     self       )     ->	Optional[Any]:





         """simple docstring"""




         _lowercase:			Union[str, Any]       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_masked_lm(*A_       )
       def 							lowercase_							(     self       )     ->	Any:





         """simple docstring"""




         _lowercase:			List[str]       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_multiple_choice(*A_       )
       def 							lowercase_							(     self       )     ->	Union[str, Any]:





         """simple docstring"""




         _lowercase:			Any       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*A_       )
       def 							lowercase_							(     self       )     ->	Optional[Any]:





         """simple docstring"""




         _lowercase:			Optional[Any]       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_pretraining(*A_       )
       def 							lowercase_							(     self       )     ->	Tuple:





         """simple docstring"""




         _lowercase:			List[Any]       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_question_answering(*A_       )
       def 							lowercase_							(     self       )     ->	List[str]:





         """simple docstring"""




         _lowercase:			str       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_sequence_classification(*A_       )



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





         """simple docstring"""




         _lowercase:			Dict       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_token_classification(*A_       )




def  _lowerCAmelCase		(		_UpperCamelCase						):






  """simple docstring"""



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


A__   :					Optional[Any]							=   1e-3







@require_torch
@require_sentencepiece
@require_tokenizers
class 			__magic_name__       (       unittest.TestCase ):
       @slow
       def 							lowercase_							(     self       )     ->	int:





         """simple docstring"""




         _lowercase:			Optional[int]       =     MobileBertModel.from_pretrained('''google/mobilebert-uncased'''       ).to(A_       )
         _lowercase:			str       =     _long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]       )
         with torch.no_grad():
           _lowercase:			Dict       =     model(A_       )[0]
         _lowercase:			List[Any]       =     torch.Size((1, 9, 512)       )
         self.assertEqual(output.shape ,				A_       )
         _lowercase:			Tuple       =     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=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
         _lowercase:			Tuple       =     torch.all((expected_slice / output[..., :3, :3]) >= 1 - TOLERANCE       )
         _lowercase:			Dict       =     torch.all((expected_slice / output[..., :3, :3]) <= 1 + TOLERANCE       )

         self.assertTrue(lower_bound and upper_bound       )
