





from __future__ import annotations

import os
import tempfile
import unittest

from transformers import ConvBertConfig, 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 import (
        TFConvBertForMaskedLM,
        TFConvBertForMultipleChoice,
        TFConvBertForQuestionAnswering,
        TFConvBertForSequenceClassification,
        TFConvBertForTokenClassification,
        TFConvBertModel,
    )



class 					A__   :





      def __init__( self       ,							A_       ,							A_=13       ,							A_=7       ,							A_=True       ,							A_=True       ,							A_=True       ,							A_=True       ,							A_=99       ,							A_=32       ,							A_=2       ,							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       ,							):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   parent
       UpperCamelCase	:					List[Any]       	=   13
       UpperCamelCase	:					str       	=   7
       UpperCamelCase	:					List[str]       	=   True
       UpperCamelCase	:					Optional[Any]       	=   True
       UpperCamelCase	:					str       	=   True
       UpperCamelCase	:					List[str]       	=   True
       UpperCamelCase	:					str       	=   99
       UpperCamelCase	:					Optional[Any]       	=   384
       UpperCamelCase	:					str       	=   2
       UpperCamelCase	:					List[str]       	=   4
       UpperCamelCase	:					Optional[Any]       	=   37
       UpperCamelCase	:					Optional[Any]       	=   """gelu"""
       UpperCamelCase	:					List[str]       	=   0.1
       UpperCamelCase	:					Dict       	=   0.1
       UpperCamelCase	:					List[str]       	=   512
       UpperCamelCase	:					List[str]       	=   16
       UpperCamelCase	:					Union[str, Any]       	=   2
       UpperCamelCase	:					Optional[int]       	=   0.02
       UpperCamelCase	:					int       	=   3
       UpperCamelCase	:					Any       	=   4
       UpperCamelCase	:					Optional[int]       	=   128
       UpperCamelCase	:					Any       	=   2
       UpperCamelCase	:					Optional[int]       	=   9
       UpperCamelCase	:					int       	=   1
       UpperCamelCase	:					List[Any]       	=   None





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




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

       UpperCamelCase	:					Tuple       	=   None
       if self.use_input_mask:
        UpperCamelCase	:					str       	=   random_attention_mask([self.batch_size, self.seq_length]   )

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

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

       UpperCamelCase	:					Optional[int]       	=   ConvBertConfig(
           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       ,							return_dict=A_       ,							)

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





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[Any]       	=   TFConvBertModel(config=A_   )
       UpperCamelCase	:					Dict       	=   {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}

       UpperCamelCase	:					Dict       	=   [input_ids, input_mask]
       UpperCamelCase	:					Optional[Any]       	=   model(A_   )

       UpperCamelCase	:					Optional[int]       	=   model(A_   )

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





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[int]       	=   TFConvBertForMaskedLM(config=A_   )
       UpperCamelCase	:					Optional[Any]       	=   {
           """input_ids""": input_ids,
           """attention_mask""": input_mask,
           """token_type_ids""": token_type_ids,
       }
       UpperCamelCase	:					List[str]       	=   model(A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.vocab_size)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					str       	=   self.num_labels
       UpperCamelCase	:					Dict       	=   TFConvBertForSequenceClassification(config=A_   )
       UpperCamelCase	:					List[Any]       	=   {
           """input_ids""": input_ids,
           """attention_mask""": input_mask,
           """token_type_ids""": token_type_ids,
       }

       UpperCamelCase	:					Optional[Any]       	=   model(A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.num_labels)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					str       	=   self.num_choices
       UpperCamelCase	:					int       	=   TFConvBertForMultipleChoice(config=A_   )
       UpperCamelCase	:					List[Any]       	=   tf.tile(tf.expand_dims(A_       ,							1   )       ,							(1, self.num_choices, 1)   )
       UpperCamelCase	:					Dict       	=   tf.tile(tf.expand_dims(A_       ,							1   )       ,							(1, self.num_choices, 1)   )
       UpperCamelCase	:					Optional[Any]       	=   tf.tile(tf.expand_dims(A_       ,							1   )       ,							(1, self.num_choices, 1)   )
       UpperCamelCase	:					Dict       	=   {
           """input_ids""": multiple_choice_inputs_ids,
           """attention_mask""": multiple_choice_input_mask,
           """token_type_ids""": multiple_choice_token_type_ids,
       }
       UpperCamelCase	:					Tuple       	=   model(A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.num_choices)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[int]       	=   self.num_labels
       UpperCamelCase	:					List[str]       	=   TFConvBertForTokenClassification(config=A_   )
       UpperCamelCase	:					Optional[int]       	=   {
           """input_ids""": input_ids,
           """attention_mask""": input_mask,
           """token_type_ids""": token_type_ids,
       }
       UpperCamelCase	:					List[str]       	=   model(A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.num_labels)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[Any]       	=   TFConvBertForQuestionAnswering(config=A_   )
       UpperCamelCase	:					Dict       	=   {
           """input_ids""": input_ids,
           """attention_mask""": input_mask,
           """token_type_ids""": token_type_ids,
       }

       UpperCamelCase	:					List[Any]       	=   model(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        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   self.prepare_config_and_inputs()
       (
           UpperCamelCase    
       )	:					Dict       	=   config_and_inputs
       UpperCamelCase	:					List[Any]       	=   {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
       return config, inputs_dict



@require_tf
class 					A__   (       UpperCAmelCase__       ,   UpperCAmelCase__       ,   unittest.TestCase			):
      _UpperCAmelCase	:str						    =   (
          (
              TFConvBertModel,
              TFConvBertForMaskedLM,
              TFConvBertForQuestionAnswering,
              TFConvBertForSequenceClassification,
              TFConvBertForTokenClassification,
              TFConvBertForMultipleChoice,
          )
          if is_tf_available()
          else ()
      )
      _UpperCAmelCase	:Union[str, Any]						    =   (
          {
              "feature-extraction": TFConvBertModel,
              "fill-mask": TFConvBertForMaskedLM,
              "question-answering": TFConvBertForQuestionAnswering,
              "text-classification": TFConvBertForSequenceClassification,
              "token-classification": TFConvBertForTokenClassification,
              "zero-shot": TFConvBertForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      _UpperCAmelCase	:Any						    =   False
      _UpperCAmelCase	:List[Any]						    =   False
      _UpperCAmelCase	:str						    =   False





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					List[Any]       	=   TFConvBertModelTester(self   )
       UpperCamelCase	:					Any       	=   ConfigTester(self       ,							config_class=A_       ,							hidden_size=37   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Union[str, Any]       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					str       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Union[str, Any]       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Dict       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Union[str, Any]       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*A_   )





      @slow
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   self.model_tester.prepare_config_and_inputs_for_common()
       UpperCamelCase	:					Tuple       	=   True
       UpperCamelCase	:					List[Any]       	=   True

       if hasattr(A_       ,							"use_cache"   ):
        UpperCamelCase	:					Dict       	=   True

       UpperCamelCase	:					int       	=   getattr(self.model_tester       ,							"encoder_seq_length"       ,							self.model_tester.seq_length   )
       UpperCamelCase	:					Any       	=   getattr(self.model_tester       ,							"key_length"       ,							A_   )

       for model_class in self.all_model_classes:
        UpperCamelCase	:					List[str]       	=   self._prepare_for_class(A_       ,							A_   )
        UpperCamelCase	:					List[Any]       	=   model_class(A_   )
        UpperCamelCase	:					Any       	=   len(model(A_   )   )

        with tempfile.TemporaryDirectory() as tmpdirname:
         model.save_pretrained(A_       ,							saved_model=A_   )
         UpperCamelCase	:					int       	=   os.path.join(A_       ,							"saved_model"       ,							"1"   )
         UpperCamelCase	:					Union[str, Any]       	=   tf.keras.models.load_model(A_   )
         UpperCamelCase	:					Any       	=   model(A_   )

         if self.is_encoder_decoder:
          UpperCamelCase	:					Optional[Any]       	=   outputs["""encoder_hidden_states"""]
          UpperCamelCase	:					Dict       	=   outputs["""encoder_attentions"""]
         else:
          UpperCamelCase	:					Tuple       	=   outputs["""hidden_states"""]
          UpperCamelCase	:					Any       	=   outputs["""attentions"""]

         self.assertEqual(len(A_   )       ,							A_   )

         UpperCamelCase	:					int       	=   getattr(
             self.model_tester       ,							"expected_num_hidden_layers"       ,							self.model_tester.num_hidden_layers + 1   )

         self.assertEqual(len(A_   )       ,							A_   )
         self.assertListEqual(
             list(output_hidden_states[0].shape[-2:]   )       ,							[self.model_tester.seq_length, self.model_tester.hidden_size]       ,							)

         self.assertEqual(len(A_   )       ,							self.model_tester.num_hidden_layers   )
         self.assertListEqual(
             list(output_attentions[0].shape[-3:]   )       ,							[self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]       ,							)





      @slow
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   TFConvBertModel.from_pretrained("YituTech/conv-bert-base"   )
       self.assertIsNotNone(A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   self.model_tester.prepare_config_and_inputs_for_common()
       UpperCamelCase	:					List[Any]       	=   True
       UpperCamelCase	:					Union[str, Any]       	=   getattr(self.model_tester       ,							"decoder_seq_length"       ,							self.model_tester.seq_length   )
       UpperCamelCase	:					List[str]       	=   getattr(self.model_tester       ,							"encoder_seq_length"       ,							self.model_tester.seq_length   )
       UpperCamelCase	:					Optional[int]       	=   getattr(self.model_tester       ,							"key_length"       ,							A_   )
       UpperCamelCase	:					Dict       	=   getattr(self.model_tester       ,							"key_length"       ,							A_   )

       def check_decoder_attentions_output(A_   ):
        UpperCamelCase	:					List[str]       	=   len(A_   )
        self.assertEqual(out_len % 2       ,							0   )
        UpperCamelCase	:					Union[str, Any]       	=   outputs.decoder_attentions
        self.assertEqual(len(A_   )       ,							self.model_tester.num_hidden_layers   )
        self.assertListEqual(
            list(decoder_attentions[0].shape[-3:]   )       ,							[self.model_tester.num_attention_heads / 2, decoder_seq_length, decoder_key_length]       ,							)

       def check_encoder_attentions_output(A_   ):
        UpperCamelCase	:					Tuple       	=   [
            t.numpy() for t in (outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions)
        ]
        self.assertEqual(len(A_   )       ,							self.model_tester.num_hidden_layers   )
        self.assertListEqual(
            list(attentions[0].shape[-3:]   )       ,							[self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]       ,							)

       for model_class in self.all_model_classes:
        UpperCamelCase	:					str       	=   True
        UpperCamelCase	:					Tuple       	=   False
        UpperCamelCase	:					List[Any]       	=   model_class(A_   )
        UpperCamelCase	:					Dict       	=   model(self._prepare_for_class(A_       ,							A_   )   )
        UpperCamelCase	:					Union[str, Any]       	=   len(A_   )
        self.assertEqual(config.output_hidden_states       ,							A_   )
        check_encoder_attentions_output(A_   )

        if self.is_encoder_decoder:
         UpperCamelCase	:					Optional[Any]       	=   model_class(A_   )
         UpperCamelCase	:					str       	=   model(self._prepare_for_class(A_       ,							A_   )   )
         self.assertEqual(config.output_hidden_states       ,							A_   )
         check_decoder_attentions_output(A_   )

        # Check that output attentions can also be changed via the config
        del inputs_dict["output_attentions"]
        UpperCamelCase	:					Union[str, Any]       	=   True
        UpperCamelCase	:					Optional[int]       	=   model_class(A_   )
        UpperCamelCase	:					Union[str, Any]       	=   model(self._prepare_for_class(A_       ,							A_   )   )
        self.assertEqual(config.output_hidden_states       ,							A_   )
        check_encoder_attentions_output(A_   )

        # Check attention is always last and order is fine
        UpperCamelCase	:					Dict       	=   True
        UpperCamelCase	:					Optional[Any]       	=   True
        UpperCamelCase	:					List[Any]       	=   model_class(A_   )
        UpperCamelCase	:					Optional[Any]       	=   model(self._prepare_for_class(A_       ,							A_   )   )

        self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)       ,							len(A_   )   )
        self.assertEqual(model.config.output_hidden_states       ,							A_   )
        check_encoder_attentions_output(A_   )



@require_tf
class 					A__   (       unittest.TestCase			):





      @slow
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   TFConvBertModel.from_pretrained("YituTech/conv-bert-base"   )
       UpperCamelCase	:					Tuple       	=   tf.constant([[0, 1, 2, 3, 4, 5]]   )
       UpperCamelCase	:					Dict       	=   model(A_   )[0]

       UpperCamelCase	:					List[str]       	=   [1, 6, 768]
       self.assertEqual(output.shape       ,							A_   )

       UpperCamelCase	:					Tuple       	=   tf.constant(
           [
               [
                   [-0.03_47_54_93, -0.4_68_60_34, -0.30_63_88_32],
                   [0.22_63_72_48, -0.26_98_86_46, -0.7_42_34_24],
                   [0.10_32_48_68, -0.45_01_35_08, -0.58_28_07_84],
               ]
           ]   )
       tf.debugging.assert_near(output[:, :3, :3]       ,							A_       ,							atol=1e-4   )




'''simple docstring'''







import pytest

from datasets.utils.sharding import _distribute_shards, _number_of_shards_in_gen_kwargs, _split_gen_kwargs


@pytest.mark.parametrize(
    """kwargs, expected"""    ,[
        ({"""num_shards""": 0, """max_num_jobs""": 1}, []),
        ({"""num_shards""": 10, """max_num_jobs""": 1}, [range(10  )]),
        ({"""num_shards""": 10, """max_num_jobs""": 10}, [range(_UpperCAmelCase    ,i + 1  ) for i in range(10  )]),
        ({"""num_shards""": 1, """max_num_jobs""": 10}, [range(1  )]),
        ({"""num_shards""": 10, """max_num_jobs""": 3}, [range(0    ,4  ), range(4    ,7  ), range(7    ,10  )]),
        ({"""num_shards""": 3, """max_num_jobs""": 10}, [range(0    ,1  ), range(1    ,2  ), range(2    ,3  )]),
    ]    ,)
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Union[str, Any]    ,_UpperCAmelCase	:	Dict  )  ->			Optional[Any]:
       _a    : Tuple     					=_distribute_shards(**_UpperCAmelCase  )
       assert out == expected


@pytest.mark.parametrize(
    """gen_kwargs, max_num_jobs, expected"""    ,[
        ({"""foo""": 0}, 10, [{"""foo""": 0}]),
        ({"""shards""": [0, 1, 2, 3]}, 1, [{"""shards""": [0, 1, 2, 3]}]),
        ({"""shards""": [0, 1, 2, 3]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}, {"""shards""": [2]}, {"""shards""": [3]}]),
        ({"""shards""": [0, 1]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}]),
        ({"""shards""": [0, 1, 2, 3]}, 2, [{"""shards""": [0, 1]}, {"""shards""": [2, 3]}]),
    ]    ,)
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Tuple    ,_UpperCAmelCase	:	List[Any]    ,_UpperCAmelCase	:	Union[str, Any]  )  ->			List[str]:
       _a    : List[str]     					=_split_gen_kwargs(_UpperCAmelCase    ,_UpperCAmelCase  )
       assert out == expected


@pytest.mark.parametrize(
    """gen_kwargs, expected"""    ,[
        ({"""foo""": 0}, 1),
        ({"""shards""": [0]}, 1),
        ({"""shards""": [0, 1, 2, 3]}, 4),
        ({"""shards""": [0, 1, 2, 3], """foo""": 0}, 4),
        ({"""shards""": [0, 1, 2, 3], """other""": (0, 1)}, 4),
        ({"""shards""": [0, 1, 2, 3], """shards2""": [0, 1]}, RuntimeError),
    ]    ,)
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Optional[Any]    ,_UpperCAmelCase	:	List[Any]  )  ->			Union[str, Any]:
       if expected is RuntimeError:
              with pytest.raises(_UpperCAmelCase  ):
                     _number_of_shards_in_gen_kwargs(_UpperCAmelCase  )
       else:
              _a    : Dict     					=_number_of_shards_in_gen_kwargs(_UpperCAmelCase  )
              assert out == expected


