



'''simple docstring'''





from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
 import tensorflow as tf

 from transformers import (
     TF_MODEL_FOR_PRETRAINING_MAPPING,
     TFMobileBertForMaskedLM,
     TFMobileBertForMultipleChoice,
     TFMobileBertForNextSentencePrediction,
     TFMobileBertForPreTraining,
     TFMobileBertForQuestionAnswering,
     TFMobileBertForSequenceClassification,
     TFMobileBertForTokenClassification,
     TFMobileBertModel,
 )


@require_tf
class 				snake_case (UpperCamelCase		,      UpperCamelCase		,      unittest.TestCase   ):
       lowerCAmelCase__  :str						   =					(
           (
               TFMobileBertModel,
               TFMobileBertForMaskedLM,
               TFMobileBertForNextSentencePrediction,
               TFMobileBertForPreTraining,
               TFMobileBertForQuestionAnswering,
               TFMobileBertForSequenceClassification,
               TFMobileBertForTokenClassification,
               TFMobileBertForMultipleChoice,
           )
           if is_tf_available()
           else ()
       )
       lowerCAmelCase__  :List[Any]						   =					(
           {
               "feature-extraction": TFMobileBertModel,
               "fill-mask": TFMobileBertForMaskedLM,
               "question-answering": TFMobileBertForQuestionAnswering,
               "text-classification": TFMobileBertForSequenceClassification,
               "token-classification": TFMobileBertForTokenClassification,
               "zero-shot": TFMobileBertForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )
       lowerCAmelCase__  :Optional[Any]						   =					False
       lowerCAmelCase__  :Any						   =					False



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_=False      )  ->							List[Any]:
           lowercase__		       =						super()._prepare_for_class(UpperCAmelCase_     ,UpperCAmelCase_     ,return_labels=UpperCAmelCase_      )

           if return_labels:
               if model_class in get_values(UpperCAmelCase_      ):
                   lowercase__		       =						tf.zeros(self.model_tester.batch_size     ,dtype=tf.intaa      )

           return inputs_dict


       class 				snake_case (UpperCamelCase   ):



              def __init__(				self     ,UpperCAmelCase_     ,UpperCAmelCase_=13     ,UpperCAmelCase_=7     ,UpperCAmelCase_=True     ,UpperCAmelCase_=True     ,UpperCAmelCase_=True     ,UpperCAmelCase_=True     ,UpperCAmelCase_=99     ,UpperCAmelCase_=32     ,UpperCAmelCase_=32     ,UpperCAmelCase_=2     ,UpperCAmelCase_=4     ,UpperCAmelCase_=37     ,UpperCAmelCase_="gelu"     ,UpperCAmelCase_=0.1     ,UpperCAmelCase_=0.1     ,UpperCAmelCase_=512     ,UpperCAmelCase_=16     ,UpperCAmelCase_=2     ,UpperCAmelCase_=0.02     ,UpperCAmelCase_=3     ,UpperCAmelCase_=4     ,UpperCAmelCase_=None     ,)  ->							Tuple:
                  lowercase__		       =						parent
                  lowercase__		       =						batch_size
                  lowercase__		       =						seq_length
                  lowercase__		       =						is_training
                  lowercase__		       =						use_input_mask
                  lowercase__		       =						use_token_type_ids
                  lowercase__		       =						use_labels
                  lowercase__		       =						vocab_size
                  lowercase__		       =						hidden_size
                  lowercase__		       =						num_hidden_layers
                  lowercase__		       =						num_attention_heads
                  lowercase__		       =						intermediate_size
                  lowercase__		       =						hidden_act
                  lowercase__		       =						hidden_dropout_prob
                  lowercase__		       =						attention_probs_dropout_prob
                  lowercase__		       =						max_position_embeddings
                  lowercase__		       =						type_vocab_size
                  lowercase__		       =						type_sequence_label_size
                  lowercase__		       =						initializer_range
                  lowercase__		       =						num_labels
                  lowercase__		       =						num_choices
                  lowercase__		       =						scope
                  lowercase__		       =						embedding_size



              def 							_a      (				self      )  ->							Tuple:
                  lowercase__		       =						ids_tensor([self.batch_size, self.seq_length]     ,self.vocab_size      )

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

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

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

                  lowercase__		       =						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     ,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     ,embedding_size=self.embedding_size     ,)

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



              def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Any:
                  lowercase__		       =						TFMobileBertModel(config=UpperCAmelCase_      )
                  lowercase__		       =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
                  lowercase__		       =						model(UpperCAmelCase_      )

                  lowercase__		       =						[input_ids, input_mask]
                  lowercase__		       =						model(UpperCAmelCase_      )

                  lowercase__		       =						model(UpperCAmelCase_      )

                  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     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Tuple:
                  lowercase__		       =						TFMobileBertForMaskedLM(config=UpperCAmelCase_      )
                  lowercase__		       =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
                  lowercase__		       =						model(UpperCAmelCase_      )
                  self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.seq_length, self.vocab_size)      )



              def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							List[str]:
                  lowercase__		       =						TFMobileBertForNextSentencePrediction(config=UpperCAmelCase_      )
                  lowercase__		       =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
                  lowercase__		       =						model(UpperCAmelCase_      )
                  self.parent.assertEqual(result.logits.shape     ,(self.batch_size, 2)      )



              def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							str:
                  lowercase__		       =						TFMobileBertForPreTraining(config=UpperCAmelCase_      )
                  lowercase__		       =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
                  lowercase__		       =						model(UpperCAmelCase_      )
                  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     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Dict:
                  lowercase__		       =						self.num_labels
                  lowercase__		       =						TFMobileBertForSequenceClassification(config=UpperCAmelCase_      )
                  lowercase__		       =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
                  lowercase__		       =						model(UpperCAmelCase_      )
                  self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.num_labels)      )



              def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Optional[int]:
                  lowercase__		       =						self.num_choices
                  lowercase__		       =						TFMobileBertForMultipleChoice(config=UpperCAmelCase_      )
                  lowercase__		       =						tf.tile(tf.expand_dims(UpperCAmelCase_     ,1      )     ,(1, self.num_choices, 1)      )
                  lowercase__		       =						tf.tile(tf.expand_dims(UpperCAmelCase_     ,1      )     ,(1, self.num_choices, 1)      )
                  lowercase__		       =						tf.tile(tf.expand_dims(UpperCAmelCase_     ,1      )     ,(1, self.num_choices, 1)      )
                  lowercase__		       =						{
                      "input_ids": multiple_choice_inputs_ids,
                      "attention_mask": multiple_choice_input_mask,
                      "token_type_ids": multiple_choice_token_type_ids,
                  }
                  lowercase__		       =						model(UpperCAmelCase_      )
                  self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.num_choices)      )



              def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Optional[Any]:
                  lowercase__		       =						self.num_labels
                  lowercase__		       =						TFMobileBertForTokenClassification(config=UpperCAmelCase_      )
                  lowercase__		       =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
                  lowercase__		       =						model(UpperCAmelCase_      )
                  self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.seq_length, self.num_labels)      )



              def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Union[str, Any]:
                  lowercase__		       =						TFMobileBertForQuestionAnswering(config=UpperCAmelCase_      )
                  lowercase__		       =						{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
                  lowercase__		       =						model(UpperCAmelCase_      )
                  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      )  ->							Union[str, Any]:
                  lowercase__		       =						self.prepare_config_and_inputs()
                  (
                      (
                      lowercase__
                  )       ,       (
                      lowercase__
                  )       ,       (
                      lowercase__
                  )       ,       (
                      lowercase__
                  )       ,       (
                      lowercase__
                  )       ,       (
                      lowercase__
                  )       ,       (
                      lowercase__
                  )       ,       
                  )		       =						config_and_inputs
                  lowercase__		       =						{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
                  return config, inputs_dict





       def 							_a      (				self      )  ->							Optional[Any]:
           lowercase__		       =						TFMobileBertModelTest.TFMobileBertModelTester(self      )
           lowercase__		       =						ConfigTester(self     ,config_class=UpperCAmelCase_     ,hidden_size=37      )



       def 							_a      (				self      )  ->							List[str]:
           self.config_tester.run_common_tests()



       def 							_a      (				self      )  ->							List[str]:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_model(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							str:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_masked_lm(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							List[Any]:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_multiple_choice(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Optional[Any]:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							int:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_pretraining(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Union[str, Any]:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_question_answering(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Union[str, Any]:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_sequence_classification(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_token_classification(*UpperCAmelCase_      )







       @slow
       def 							_a      (				self      )  ->							List[Any]:
           # for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           for model_name in ["google/mobilebert-uncased"]:
               lowercase__		       =						TFMobileBertModel.from_pretrained(UpperCAmelCase_      )
               self.assertIsNotNone(UpperCAmelCase_      )





@require_tf
class 				snake_case (unittest.TestCase   ):



       @slow
       def 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						TFMobileBertForPreTraining.from_pretrained("google/mobilebert-uncased"      )
           lowercase__		       =						tf.constant([[0, 1, 2, 3, 4, 5]]      )
           lowercase__		       =						model(UpperCAmelCase_      )[0]

           lowercase__		       =						[1, 6, 30_522]
           self.assertEqual(output.shape     ,UpperCAmelCase_      )

           lowercase__		       =						tf.constant(
               [
                   [
                       [-4.5_91_95_47, -9.24_82_95, -9.64_52_56],
                       [-6.7_30_61_75, -6.44_02_84, -6.6_05_28_37],
                       [-7.2_74_35_06, -6.7_84_79_15, -6.02_46_73],
                   ]
               ]      )
           tf.debugging.assert_near(output[:, :3, :3]     ,UpperCAmelCase_     ,atol=1E-4      )









'''simple docstring'''





import numpy as np
import torch
from torch.utils.data import Dataset

from utils import logger


class 				snake_case (UpperCamelCase   ):



       def __init__(				self     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							int:
           lowercase__		       =						params

           lowercase__		       =						np.array(UpperCAmelCase_      )
           lowercase__		       =						np.array([len(UpperCAmelCase_      ) for t in data]      )

           self.check()
           self.remove_long_sequences()
           self.remove_empty_sequences()
           self.remove_unknown_sequences()
           self.check()
           self.print_statistics()



       def __getitem__(				self     ,UpperCAmelCase_      )  ->							Optional[int]:
           return (self.token_ids[index], self.lengths[index])



       def __len__(				self      )  ->							Union[str, Any]:
           return len(self.lengths      )



       def 							_a      (				self      )  ->							int:
           assert len(self.token_ids      ) == len(self.lengths      )
           assert all(self.lengths[i] == len(self.token_ids[i]      ) for i in range(len(self.lengths      )      )      )



       def 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						self.params.max_model_input_size
           lowercase__		       =						self.lengths > max_len
           logger.info(F'''Splitting {sum(UpperCAmelCase_      )} too long sequences.'''      )

           def divide_chunks(UpperCAmelCase_     ,UpperCAmelCase_      ):
               return [l[i : i + n] for i in range(0     ,len(UpperCAmelCase_      )     ,UpperCAmelCase_      )]

           lowercase__		       =						[]
           lowercase__		       =						[]
           if self.params.mlm:
               lowercase__       ,       lowercase__		       =						self.params.special_tok_ids["cls_token"], self.params.special_tok_ids["sep_token"]
           else:
               lowercase__       ,       lowercase__		       =						self.params.special_tok_ids["bos_token"], self.params.special_tok_ids["eos_token"]

           for seq_, len_ in zip(self.token_ids     ,self.lengths      ):
               assert (seq_[0] == cls_id) and (seq_[-1] == sep_id), seq_
               if len_ <= max_len:
                   new_tok_ids.append(seq_      )
                   new_lengths.append(len_      )
               else:
                   lowercase__		       =						[]
                   for sub_s in divide_chunks(seq_     ,max_len - 2      ):
                       if sub_s[0] != cls_id:
                           lowercase__		       =						np.insert(UpperCAmelCase_     ,0     ,UpperCAmelCase_      )
                       if sub_s[-1] != sep_id:
                           lowercase__		       =						np.insert(UpperCAmelCase_     ,len(UpperCAmelCase_      )     ,UpperCAmelCase_      )
                       assert len(UpperCAmelCase_      ) <= max_len
                       assert (sub_s[0] == cls_id) and (sub_s[-1] == sep_id), sub_s
                       sub_seqs.append(UpperCAmelCase_      )

                   new_tok_ids.extend(UpperCAmelCase_      )
                   new_lengths.extend([len(UpperCAmelCase_      ) for l in sub_seqs]      )

           lowercase__		       =						np.array(UpperCAmelCase_      )
           lowercase__		       =						np.array(UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Any:
           lowercase__		       =						len(self      )
           lowercase__		       =						self.lengths > 11
           lowercase__		       =						self.token_ids[indices]
           lowercase__		       =						self.lengths[indices]
           lowercase__		       =						len(self      )
           logger.info(F'''Remove {init_size - new_size} too short (<=11 tokens) sequences.'''      )



       def 							_a      (				self      )  ->							List[Any]:
           if "unk_token" not in self.params.special_tok_ids:
               return
           else:
               lowercase__		       =						self.params.special_tok_ids["unk_token"]
           lowercase__		       =						len(self      )
           lowercase__		       =						np.array([np.count_nonzero(a == unk_token_id      ) for a in self.token_ids]      )
           lowercase__		       =						(unk_occs / self.lengths) < 0.5
           lowercase__		       =						self.token_ids[indices]
           lowercase__		       =						self.lengths[indices]
           lowercase__		       =						len(self      )
           logger.info(F'''Remove {init_size - new_size} sequences with a high level of unknown tokens (50%).'''      )



       def 							_a      (				self      )  ->							Optional[int]:
           if not self.params.is_master:
               return
           logger.info(F'''{len(self      )} sequences'''      )
           # data_len = sum(self.lengths)
           # nb_unique_tokens = len(Counter(list(chain(*self.token_ids))))
           # logger.info(f'{data_len} tokens ({nb_unique_tokens} unique)')

           # unk_idx = self.params.special_tok_ids['unk_token']
           # nb_unknown = sum([(t==unk_idx).sum() for t in self.token_ids])
           # logger.info(f'{nb_unknown} unknown tokens (covering {100*nb_unknown/data_len:.2f}% of the data)')







       def 							_a      (				self     ,UpperCAmelCase_      )  ->							List[str]:
           lowercase__		       =						[t[0] for t in batch]
           lowercase__		       =						[t[1] for t in batch]
           assert len(UpperCAmelCase_      ) == len(UpperCAmelCase_      )

           # Max for paddings
           lowercase__		       =						max(UpperCAmelCase_      )

           # Pad token ids
           if self.params.mlm:
               lowercase__		       =						self.params.special_tok_ids["pad_token"]
           else:
               lowercase__		       =						self.params.special_tok_ids["unk_token"]
           lowercase__		       =						[list(t.astype(UpperCAmelCase_      )      ) + [pad_idx] * (max_seq_len_ - len(UpperCAmelCase_      )) for t in token_ids]
           assert len(tk_      ) == len(UpperCAmelCase_      )
           assert all(len(UpperCAmelCase_      ) == max_seq_len_ for t in tk_      )

           lowercase__		       =						torch.tensor(tk_      )  # (bs, max_seq_len_)
           lowercase__		       =						torch.tensor(UpperCAmelCase_      )  # (bs)
           return tk_t, lg_t




