

"""simple docstring"""




from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
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.models.distilbert.modeling_tf_distilbert import (
          TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFDistilBertForMaskedLM,
          TFDistilBertForMultipleChoice,
          TFDistilBertForQuestionAnswering,
          TFDistilBertForSequenceClassification,
          TFDistilBertForTokenClassification,
          TFDistilBertModel,
      )







class   __lowercase      :

  def __init__(	self		:      Optional[int]       ,A		:      Optional[Any]       ,):






   '''simple docstring'''
   UpperCAmelCase__       :	List[Any]        =       parent
   UpperCAmelCase__       :	Any        =       13
   UpperCAmelCase__       :	Union[str, Any]        =       7
   UpperCAmelCase__       :	str        =       True
   UpperCAmelCase__       :	Tuple        =       True
   UpperCAmelCase__       :	Optional[int]        =       False
   UpperCAmelCase__       :	Optional[int]        =       True
   UpperCAmelCase__       :	Union[str, Any]        =       99
   UpperCAmelCase__       :	Union[str, Any]        =       32
   UpperCAmelCase__       :	Dict        =       2
   UpperCAmelCase__       :	List[str]        =       4
   UpperCAmelCase__       :	str        =       37
   UpperCAmelCase__       :	Any        =       """gelu"""
   UpperCAmelCase__       :	int        =       0.1
   UpperCAmelCase__       :	str        =       0.1
   UpperCAmelCase__       :	Optional[Any]        =       512
   UpperCAmelCase__       :	Tuple        =       16
   UpperCAmelCase__       :	Union[str, Any]        =       2
   UpperCAmelCase__       :	int        =       0.0_2
   UpperCAmelCase__       :	Union[str, Any]        =       3
   UpperCAmelCase__       :	int        =       4
   UpperCAmelCase__       :	Optional[Any]        =       None

  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Tuple        =       ids_tensor([self.batch_size, self.seq_length]       ,self.vocab_size     )

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

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

   UpperCAmelCase__       :	Optional[Any]        =       DistilBertConfig(
       vocab_size=self.vocab_size       ,dim=self.hidden_size       ,n_layers=self.num_hidden_layers       ,n_heads=self.num_attention_heads       ,hidden_dim=self.intermediate_size       ,hidden_act=self.hidden_act       ,dropout=self.hidden_dropout_prob       ,attention_dropout=self.attention_probs_dropout_prob       ,max_position_embeddings=self.max_position_embeddings       ,initializer_range=self.initializer_range       ,)

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

  def  __lowercase							(	self		:      Any       ,A		:      Tuple       ,A		:      Any       ,A		:      Union[str, Any]       ,A		:      List[Any]       ,A		:      Optional[int]       ,A		:      Any     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[str]        =       TFDistilBertModel(config=_SCREAMING_SNAKE_CASE     )
   UpperCAmelCase__       :	str        =       {"""input_ids""": input_ids, """attention_mask""": input_mask}

   UpperCAmelCase__       :	Optional[int]        =       model(_SCREAMING_SNAKE_CASE     )

   UpperCAmelCase__       :	Optional[Any]        =       [input_ids, input_mask]

   UpperCAmelCase__       :	int        =       model(_SCREAMING_SNAKE_CASE     )

   self.parent.assertEqual(result.last_hidden_state.shape       ,(self.batch_size, self.seq_length, self.hidden_size)     )

  def  __lowercase							(	self		:      List[str]       ,A		:      str       ,A		:      Optional[int]       ,A		:      Union[str, Any]       ,A		:      Tuple       ,A		:      int       ,A		:      Union[str, Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Tuple        =       TFDistilBertForMaskedLM(config=_SCREAMING_SNAKE_CASE     )
   UpperCAmelCase__       :	Optional[int]        =       {"""input_ids""": input_ids, """attention_mask""": input_mask}
   UpperCAmelCase__       :	Optional[Any]        =       model(_SCREAMING_SNAKE_CASE     )
   self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.seq_length, self.vocab_size)     )

  def  __lowercase							(	self		:      Any       ,A		:      Optional[Any]       ,A		:      Tuple       ,A		:      Optional[Any]       ,A		:      Tuple       ,A		:      Optional[Any]       ,A		:      Optional[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[Any]        =       TFDistilBertForQuestionAnswering(config=_SCREAMING_SNAKE_CASE     )
   UpperCAmelCase__       :	List[str]        =       {
       """input_ids""": input_ids,
       """attention_mask""": input_mask,
   }
   UpperCAmelCase__       :	List[Any]        =       model(_SCREAMING_SNAKE_CASE     )
   self.parent.assertEqual(result.start_logits.shape       ,(self.batch_size, self.seq_length)     )
   self.parent.assertEqual(result.end_logits.shape       ,(self.batch_size, self.seq_length)     )

  def  __lowercase							(	self		:      int       ,A		:      str       ,A		:      str       ,A		:      str       ,A		:      Optional[Any]       ,A		:      int       ,A		:      List[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Dict        =       self.num_labels
   UpperCAmelCase__       :	List[Any]        =       TFDistilBertForSequenceClassification(_SCREAMING_SNAKE_CASE     )
   UpperCAmelCase__       :	Optional[int]        =       {"""input_ids""": input_ids, """attention_mask""": input_mask}
   UpperCAmelCase__       :	List[Any]        =       model(_SCREAMING_SNAKE_CASE     )
   self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.num_labels)     )

  def  __lowercase							(	self		:      Union[str, Any]       ,A		:      int       ,A		:      Dict       ,A		:      int       ,A		:      Tuple       ,A		:      str       ,A		:      Optional[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[Any]        =       self.num_choices
   UpperCAmelCase__       :	Union[str, Any]        =       TFDistilBertForMultipleChoice(_SCREAMING_SNAKE_CASE     )
   UpperCAmelCase__       :	int        =       tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE       ,1     )       ,(1, self.num_choices, 1)     )
   UpperCAmelCase__       :	str        =       tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE       ,1     )       ,(1, self.num_choices, 1)     )
   UpperCAmelCase__       :	List[Any]        =       {
       """input_ids""": multiple_choice_inputs_ids,
       """attention_mask""": multiple_choice_input_mask,
   }
   UpperCAmelCase__       :	int        =       model(_SCREAMING_SNAKE_CASE     )
   self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.num_choices)     )

  def  __lowercase							(	self		:      List[str]       ,A		:      Optional[Any]       ,A		:      Optional[Any]       ,A		:      Tuple       ,A		:      Optional[int]       ,A		:      str       ,A		:      str     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[int]        =       self.num_labels
   UpperCAmelCase__       :	str        =       TFDistilBertForTokenClassification(_SCREAMING_SNAKE_CASE     )
   UpperCAmelCase__       :	Union[str, Any]        =       {"""input_ids""": input_ids, """attention_mask""": input_mask}
   UpperCAmelCase__       :	str        =       model(_SCREAMING_SNAKE_CASE     )
   self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.seq_length, self.num_labels)     )




  def  __lowercase							(	self		:      Optional[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[str]        =       self.prepare_config_and_inputs()
   ((UpperCAmelCase__) ,   (UpperCAmelCase__) ,   (UpperCAmelCase__) ,   (UpperCAmelCase__) ,   (UpperCAmelCase__) ,   (UpperCAmelCase__))       :	Optional[int]        =       config_and_inputs
   UpperCAmelCase__       :	Any        =       {"""input_ids""": input_ids, """attention_mask""": input_mask}
   return config, inputs_dict







@require_tf
class   __lowercase      (     __snake_case						,      __snake_case						,      unittest.TestCase       ):
  snake_case_												=							(
      (
          TFDistilBertModel,
          TFDistilBertForMaskedLM,
          TFDistilBertForQuestionAnswering,
          TFDistilBertForSequenceClassification,
          TFDistilBertForTokenClassification,
          TFDistilBertForMultipleChoice,
      )
      if is_tf_available()
      else None
  )
  snake_case_												=							(
      {
          """feature-extraction""": TFDistilBertModel,
          """fill-mask""": TFDistilBertForMaskedLM,
          """question-answering""": TFDistilBertForQuestionAnswering,
          """text-classification""": TFDistilBertForSequenceClassification,
          """token-classification""": TFDistilBertForTokenClassification,
          """zero-shot""": TFDistilBertForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  snake_case_												=							False
  snake_case_												=							False

  def  __lowercase							(	self		:      Tuple     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[Any]        =       TFDistilBertModelTester(self     )
   UpperCAmelCase__       :	Optional[int]        =       ConfigTester(self       ,config_class=_SCREAMING_SNAKE_CASE       ,dim=37     )

  def  __lowercase							(	self		:      Union[str, Any]     ):






   '''simple docstring'''
   self.config_tester.run_common_tests()

  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   UpperCAmelCase__       :	int        =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_model(*_SCREAMING_SNAKE_CASE     )

  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Tuple        =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_masked_lm(*_SCREAMING_SNAKE_CASE     )

  def  __lowercase							(	self		:      Tuple     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[Any]        =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_question_answering(*_SCREAMING_SNAKE_CASE     )

  def  __lowercase							(	self		:      Dict     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[Any]        =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_sequence_classification(*_SCREAMING_SNAKE_CASE     )

  def  __lowercase							(	self		:      int     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Union[str, Any]        =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_multiple_choice(*_SCREAMING_SNAKE_CASE     )

  def  __lowercase							(	self		:      Optional[Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Dict        =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_token_classification(*_SCREAMING_SNAKE_CASE     )




  @slow
  def  __lowercase							(	self		:      List[str]     ):






   '''simple docstring'''




   for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]     ):
    UpperCAmelCase__       :	Any        =       TFDistilBertModel.from_pretrained(_SCREAMING_SNAKE_CASE     )
    self.assertIsNotNone(_SCREAMING_SNAKE_CASE     )







@require_tf
class   __lowercase      (     unittest.TestCase       ):

  @slow
  def  __lowercase							(	self		:      Union[str, Any]     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[str]        =       TFDistilBertModel.from_pretrained("""distilbert-base-uncased"""     )
   UpperCAmelCase__       :	Dict        =       tf.constant([[0, 1, 2, 3, 4, 5]]     )
   UpperCAmelCase__       :	Optional[int]        =       model(_SCREAMING_SNAKE_CASE     )[0]

   UpperCAmelCase__       :	str        =       [1, 6, 768]
   self.assertEqual(output.shape       ,_SCREAMING_SNAKE_CASE     )

   UpperCAmelCase__       :	int        =       tf.constant(
       [
           [
               [0.1_9_2_6_1_8_8_5, -0.1_3_7_3_2_9_5_5, 0.4_1_1_9_7_9_9],
               [0.2_2_1_5_0_1_5_6, -0.0_7_4_2_2_6_6_1, 0.3_9_0_3_7_2_0_4],
               [0.2_2_7_5_6_0_1_8, -0.0_8_9_6_4_1_4, 0.3_7_0_1_4_6_7],
           ]
       ]     )
   tf.debugging.assert_near(output[:, :3, :3]       ,_SCREAMING_SNAKE_CASE       ,atol=1e-4     )













from transformers import BertTokenizer, EncoderDecoderModel, SeqaSeqTrainer, SeqaSeqTrainingArguments
from transformers.testing_utils import TestCasePlus, require_torch, slow
from transformers.utils import is_datasets_available


if is_datasets_available():
       import datasets





class 					A__     (     __snake_case ):

     '''simple docstring'''






     @slow
     @require_torch
     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           UpperCamelCase					       =  EncoderDecoderModel.from_encoder_decoder_pretrained('prajjwal1/bert-tiny'   ,						'prajjwal1/bert-tiny' )
           UpperCamelCase					       =  BertTokenizer.from_pretrained('bert-base-uncased' )

           UpperCamelCase					       =  bertabert.config.encoder.vocab_size
           UpperCamelCase					       =  tokenizer.sep_token_id
           UpperCamelCase					       =  tokenizer.cls_token_id
           UpperCamelCase					       =  128

           UpperCamelCase					       =  datasets.load_dataset('cnn_dailymail'   ,						'3.0.0'   ,						split='train[:1%]' )
           UpperCamelCase					       =  datasets.load_dataset('cnn_dailymail'   ,						'3.0.0'   ,						split='validation[:1%]' )

           UpperCamelCase					       =  train_dataset.select(range(32 ) )
           UpperCamelCase					       =  val_dataset.select(range(16 ) )

           UpperCamelCase					       =  4

           def _map_to_encoder_decoder_inputs(_SCREAMING_SNAKE_CASE      :    Tuple ):
                 # Tokenizer will automatically set [BOS] <text> [EOS]
                 UpperCamelCase					       =  tokenizer(batch['article']   ,						padding='max_length'   ,						truncation=_SCREAMING_SNAKE_CASE   ,						max_length=512 )
                 UpperCamelCase					       =  tokenizer(batch['highlights']   ,						padding='max_length'   ,						truncation=_SCREAMING_SNAKE_CASE   ,						max_length=128 )
                 UpperCamelCase					       =  inputs.input_ids
                 UpperCamelCase					       =  inputs.attention_mask

                 UpperCamelCase					       =  outputs.input_ids
                 UpperCamelCase					       =  outputs.input_ids.copy()
                 UpperCamelCase					       =  [
                     [-100 if token == tokenizer.pad_token_id else token for token in labels] for labels in batch['labels']
                 ]
                 UpperCamelCase					       =  outputs.attention_mask

                 assert all(len(_SCREAMING_SNAKE_CASE ) == 512 for x in inputs.input_ids )
                 assert all(len(_SCREAMING_SNAKE_CASE ) == 128 for x in outputs.input_ids )

                 return batch

           def _compute_metrics(_SCREAMING_SNAKE_CASE      :    str ):
                 UpperCamelCase					       =  pred.label_ids
                 UpperCamelCase					       =  pred.predictions

                 # all unnecessary tokens are removed
                 UpperCamelCase					       =  tokenizer.batch_decode(_SCREAMING_SNAKE_CASE   ,						skip_special_tokens=_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  tokenizer.batch_decode(_SCREAMING_SNAKE_CASE   ,						skip_special_tokens=_SCREAMING_SNAKE_CASE )

                 UpperCamelCase					       =  sum([int(pred_str[i] == label_str[i] ) for i in range(len(_SCREAMING_SNAKE_CASE ) )] ) / len(_SCREAMING_SNAKE_CASE )

                 return {"accuracy": accuracy}

           # map train dataset
           UpperCamelCase					       =  train_dataset.map(
               _map_to_encoder_decoder_inputs   ,						batched=_SCREAMING_SNAKE_CASE   ,						batch_size=_SCREAMING_SNAKE_CASE   ,						remove_columns=['article', 'highlights']   ,						)
           train_dataset.set_format(
               type='torch'   ,						columns=['input_ids', 'attention_mask', 'decoder_input_ids', 'decoder_attention_mask', 'labels']   ,						)

           # same for validation dataset
           UpperCamelCase					       =  val_dataset.map(
               _map_to_encoder_decoder_inputs   ,						batched=_SCREAMING_SNAKE_CASE   ,						batch_size=_SCREAMING_SNAKE_CASE   ,						remove_columns=['article', 'highlights']   ,						)
           val_dataset.set_format(
               type='torch'   ,						columns=['input_ids', 'attention_mask', 'decoder_input_ids', 'decoder_attention_mask', 'labels']   ,						)

           UpperCamelCase					       =  self.get_auto_remove_tmp_dir()

           UpperCamelCase					       =  SeqaSeqTrainingArguments(
               output_dir=_SCREAMING_SNAKE_CASE   ,						per_device_train_batch_size=_SCREAMING_SNAKE_CASE   ,						per_device_eval_batch_size=_SCREAMING_SNAKE_CASE   ,						predict_with_generate=_SCREAMING_SNAKE_CASE   ,						evaluation_strategy='steps'   ,						do_train=_SCREAMING_SNAKE_CASE   ,						do_eval=_SCREAMING_SNAKE_CASE   ,						warmup_steps=0   ,						eval_steps=2   ,						logging_steps=2   ,						)

           # instantiate trainer
           UpperCamelCase					       =  SeqaSeqTrainer(
               model=_SCREAMING_SNAKE_CASE   ,						args=_SCREAMING_SNAKE_CASE   ,						compute_metrics=_compute_metrics   ,						train_dataset=_SCREAMING_SNAKE_CASE   ,						eval_dataset=_SCREAMING_SNAKE_CASE   ,						tokenizer=_SCREAMING_SNAKE_CASE   ,						)

           # start training
           trainer.train()



