



"""simple docstring"""







from __future__ import annotations

import unittest

from transformers import CTRLConfig, 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.ctrl.modeling_tf_ctrl import (
     TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
     TFCTRLForSequenceClassification,
     TFCTRLLMHeadModel,
     TFCTRLModel,
 )

class       UpperCamelCase  (					a				):





     """simple docstring"""






     def __init__(					self :							Any	,      _lowerCamelCase :							Union[str, Any]	,      ):
          A__				=   parent
          A__				=   1_3
          A__				=   7
          A__				=   True
          A__				=   True
          A__				=   True
          A__				=   True
          A__				=   True
          A__				=   9_9
          A__				=   3_2
          A__				=   2
          A__				=   4
          A__				=   3_7
          A__				=   '''gelu'''
          A__				=   0.1
          A__				=   0.1
          A__				=   5_1_2
          A__				=   1_6
          A__				=   2
          A__				=   0.02
          A__				=   3
          A__				=   4
          A__				=   None
          A__				=   self.vocab_size - 1





     def 						A__  (					self :							str ):
          A__				=   ids_tensor([self.batch_size, self.seq_length]	,      self.vocab_size )

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

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

          A__				=   None
          if self.use_mc_token_ids:
               A__				=   ids_tensor([self.batch_size, self.num_choices]	,      self.seq_length )

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

          A__				=   CTRLConfig(
              vocab_size=self.vocab_size	,      n_embd=self.hidden_size	,      n_layer=self.num_hidden_layers	,      n_head=self.num_attention_heads	,      n_positions=self.max_position_embeddings	,      pad_token_id=self.pad_token_id	,      )

          A__				=   ids_tensor([self.num_hidden_layers, self.num_attention_heads]	,      2 )

          return (
              config,
              input_ids,
              input_mask,
              head_mask,
              token_type_ids,
              mc_token_ids,
              sequence_labels,
              token_labels,
              choice_labels,
          )





     def 						A__  (					self :							Any	,      _lowerCamelCase :							str	,      _lowerCamelCase :							int	,      _lowerCamelCase :							List[Any]	,      _lowerCamelCase :							Dict	,      _lowerCamelCase :							List[str]	,      *_lowerCamelCase :							Dict ):
          A__				=   TFCTRLModel(config=_lowerCamelCase )
          A__				=   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
          A__				=   model(_lowerCamelCase )

          A__				=   [input_ids, None, input_mask]  # None is the input for 'past'
          A__				=   model(_lowerCamelCase )

          A__				=   model(_lowerCamelCase )

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





     def 						A__  (					self :							Tuple	,      _lowerCamelCase :							Tuple	,      _lowerCamelCase :							Optional[int]	,      _lowerCamelCase :							Any	,      _lowerCamelCase :							List[str]	,      _lowerCamelCase :							int	,      *_lowerCamelCase :							Optional[Any] ):
          A__				=   TFCTRLLMHeadModel(config=_lowerCamelCase )
          A__				=   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
          A__				=   model(_lowerCamelCase )
          self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.seq_length, self.vocab_size) )





     def 						A__  (					self :							Union[str, Any]	,      _lowerCamelCase :							Optional[int]	,      _lowerCamelCase :							Dict	,      _lowerCamelCase :							Any	,      _lowerCamelCase :							int	,      _lowerCamelCase :							str	,      *_lowerCamelCase :							Optional[int] ):
          A__				=   self.num_labels
          A__				=   ids_tensor([self.batch_size]	,      self.type_sequence_label_size )
          A__				=   {
              '''input_ids''': input_ids,
              '''token_type_ids''': token_type_ids,
              '''labels''': sequence_labels,
          }
          A__				=   TFCTRLForSequenceClassification(_lowerCamelCase )
          A__				=   model(_lowerCamelCase )
          self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.num_labels) )





     def 						A__  (					self :							Optional[Any] ):
          A__				=   self.prepare_config_and_inputs()

          (
              (
              A__
          )	,     (
              A__
          )	,     (
              A__
          )	,     (
              A__
          )	,     (
              A__
          )	,     (
              A__
          )	,     (
              A__
          )	,     (
              A__
          )	,     (
              A__
          )	,     
          )				=   config_and_inputs

          A__				=   {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
          return config, inputs_dict


@require_tf
class       UpperCamelCase  (					a      , a      , unittest.TestCase				):





     """simple docstring"""






     _lowerCamelCase					: Union[str, Any]				   =(TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
     _lowerCamelCase					: Tuple				   =(TFCTRLLMHeadModel,) if is_tf_available() else ()
     _lowerCamelCase					: List[str]				   =(
         {
             "feature-extraction": TFCTRLModel,
             "text-classification": TFCTRLForSequenceClassification,
             "text-generation": TFCTRLLMHeadModel,
             "zero-shot": TFCTRLForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     _lowerCamelCase					: int				   =False
     _lowerCamelCase					: Union[str, Any]				   =False





     def 						A__  (					self :							Dict	,      _lowerCamelCase :							Tuple	,      _lowerCamelCase :							Tuple	,      _lowerCamelCase :							Optional[Any]	,      _lowerCamelCase :							Optional[int]	,      _lowerCamelCase :							int ):
          if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
               # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
               # `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
               # config could not be created.
               return True

          return False





     def 						A__  (					self :							List[Any] ):
          A__				=   TFCTRLModelTester(self )
          A__				=   ConfigTester(self	,      config_class=_lowerCamelCase	,      n_embd=3_7 )





     def 						A__  (					self :							str ):
          self.config_tester.run_common_tests()





     def 						A__  (					self :							str ):
          A__				=   self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_ctrl_model(*_lowerCamelCase )





     def 						A__  (					self :							Optional[Any] ):
          A__				=   self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_ctrl_lm_head(*_lowerCamelCase )





     def 						A__  (					self :							Optional[Any] ):
          A__				=   self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_ctrl_for_sequence_classification(*_lowerCamelCase )





     def 						A__  (					self :							str ):
          A__	,     A__				=   self.model_tester.prepare_config_and_inputs_for_common()
          A__				=   [TFCTRLLMHeadModel]
          A__				=   [TFCTRLForSequenceClassification]

          for model_class in self.all_model_classes:
               A__				=   model_class(_lowerCamelCase )
               model.build()  # may be needed for the get_bias() call below
               assert isinstance(model.get_input_embeddings()	,      tf.keras.layers.Layer )

               if model_class in list_lm_models:
                    A__				=   model.get_output_embeddings()
                    assert isinstance(_lowerCamelCase	,      tf.keras.layers.Layer )
                    A__				=   model.get_bias()
                    assert isinstance(_lowerCamelCase	,      _lowerCamelCase )
                    for k, v in name.items():
                         assert isinstance(_lowerCamelCase	,      tf.Variable )
               elif model_class in list_other_models_with_output_ebd:
                    A__				=   model.get_output_embeddings()
                    assert isinstance(_lowerCamelCase	,      tf.keras.layers.Layer )
                    A__				=   model.get_bias()
                    assert name is None
               else:
                    A__				=   model.get_output_embeddings()
                    assert x is None
                    A__				=   model.get_bias()
                    assert name is None





     @slow
     def 						A__  (					self :							List[Any] ):
          for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               A__				=   TFCTRLModel.from_pretrained(_lowerCamelCase )
               self.assertIsNotNone(_lowerCamelCase )


@require_tf
class       UpperCamelCase  (					unittest.TestCase				):





     """simple docstring"""






     @slow
     def 						A__  (					self :							Any ):
          A__				=   TFCTRLLMHeadModel.from_pretrained('''ctrl''' )
          A__				=   tf.convert_to_tensor([[1_1_8_5_9, 0, 1_6_1_1, 8]]	,      dtype=tf.intaa )  # Legal the president is
          A__				=   [
              1_1_8_5_9,
              0,
              1_6_1_1,
              8,
              5,
              1_5_0,
              2_6_4_4_9,
              2,
              1_9,
              3_4_8,
              4_6_9,
              3,
              2_5_9_5,
              4_8,
              2_0_7_4_0,
              2_4_6_5_3_3,
              2_4_6_5_3_3,
              1_9,
              3_0,
              5,
          ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

          A__				=   model.generate(_lowerCamelCase	,      do_sample=_lowerCamelCase )
          self.assertListEqual(output_ids[0].numpy().tolist()	,      _lowerCamelCase )








"""simple docstring"""







import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


__snake_case	:				Optional[Any]  					=			500_000

__snake_case					,  __snake_case	:				Any  					=			os.path.split(__file__)
__snake_case	:				Optional[Any]  					=			os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))



@get_duration
def   a_	(			__a      ):
     A__				=   dataset.select(range(0		,							len(__a      )		,							2      )      )



@get_duration
def   a_	(			__a      ):
     A__				=   dataset.sort('''numbers'''      )



@get_duration
def   a_	(			__a      ):
     A__				=   dataset.shuffle()



@get_duration
def   a_	(			__a      ):
     A__				=   dataset.train_test_split(0.1      )



@get_duration
def   a_	(			__a		,							__a=10      ):
     for shard_id in range(__a      ):
          A__				=   dataset.shard(__a		,							__a      )



def   a_	(			):
     A__				=   {'''num examples''': SPEED_TEST_N_EXAMPLES}
     A__				=   (select, sort, shuffle, train_test_split, shard)
     with tempfile.TemporaryDirectory() as tmp_dir:
          print('''generating dataset'''      )
          A__				=   datasets.Features({'''text''': datasets.Value('''string'''      ), '''numbers''': datasets.Value('''float32'''      )}      )
          A__				=   generate_example_dataset(
              os.path.join(__a		,							'''dataset.arrow'''      )		,							__a		,							num_examples=__a      )
          print('''Functions'''      )
          for func in functions:
               print(func.__name__      )
               A__				=   func(__a      )

     with open(__a		,							'''wb'''      ) as f:
          f.write(json.dumps(__a      ).encode('''utf-8'''      )      )


if __name__ == "__main__":  # useful to run the profiler
 benchmark_indices_mapping()



