
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       :List[Any]		,						lowerCAmelCase__       :Any		,						)     ->Optional[int]:
										lowercase						    =			parent
										lowercase						    =			13
										lowercase						    =			7
										lowercase						    =			True
										lowercase						    =			True
										lowercase						    =			True
										lowercase						    =			True
										lowercase						    =			True
										lowercase						    =			99
										lowercase						    =			32
										lowercase						    =			2
										lowercase						    =			4
										lowercase						    =			37
										lowercase						    =			"gelu"
										lowercase						    =			0.1
										lowercase						    =			0.1
										lowercase						    =			512
										lowercase						    =			16
										lowercase						    =			2
										lowercase						    =			0.02
										lowercase						    =			3
										lowercase						    =			4
										lowercase						    =			None
										lowercase						    =			self.vocab_size - 1





						def 					SCREAMING_SNAKE_CASE( self       :str    )     ->List[str]:
										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
										if self.use_mc_token_ids:
														lowercase						    =			ids_tensor([self.batch_size, self.num_choices]		,						self.seq_length    )

										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						    =			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		,						)

										lowercase						    =			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 					SCREAMING_SNAKE_CASE( self       :str		,						lowerCAmelCase__       :Union[str, Any]		,						lowerCAmelCase__       :int		,						lowerCAmelCase__       :int		,						lowerCAmelCase__       :Optional[int]		,						lowerCAmelCase__       :List[Any]		,						*lowerCAmelCase__       :str    )     ->Union[str, Any]:
										lowercase						    =			TFCTRLModel(config=lowerCAmelCase__    )
										lowercase						    =			{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
										lowercase						    =			model(lowerCAmelCase__    )

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

										lowercase						    =			model(lowerCAmelCase__    )

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





						def 					SCREAMING_SNAKE_CASE( self       :Any		,						lowerCAmelCase__       :Any		,						lowerCAmelCase__       :Optional[int]		,						lowerCAmelCase__       :Optional[int]		,						lowerCAmelCase__       :str		,						lowerCAmelCase__       :int		,						*lowerCAmelCase__       :Union[str, Any]    )     ->Tuple:
										lowercase						    =			TFCTRLLMHeadModel(config=lowerCAmelCase__    )
										lowercase						    =			{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
										lowercase						    =			model(lowerCAmelCase__    )
										self.parent.assertEqual(result.logits.shape		,						(self.batch_size, self.seq_length, self.vocab_size)    )





						def 					SCREAMING_SNAKE_CASE( self       :int		,						lowerCAmelCase__       :Optional[Any]		,						lowerCAmelCase__       :Dict		,						lowerCAmelCase__       :List[str]		,						lowerCAmelCase__       :str		,						lowerCAmelCase__       :int		,						*lowerCAmelCase__       :List[str]    )     ->List[str]:
										lowercase						    =			self.num_labels
										lowercase						    =			ids_tensor([self.batch_size]		,						self.type_sequence_label_size    )
										lowercase						    =			{
										    "input_ids": input_ids,
										    "token_type_ids": token_type_ids,
										    "labels": sequence_labels,
										}
										lowercase						    =			TFCTRLForSequenceClassification(lowerCAmelCase__    )
										lowercase						    =			model(lowerCAmelCase__    )
										self.parent.assertEqual(result.logits.shape		,						(self.batch_size, self.num_labels)    )





						def 					SCREAMING_SNAKE_CASE( self       :Dict    )     ->Tuple:
										lowercase						    =			self.prepare_config_and_inputs()

										(
										    (
										    lowercase
										)    ,  (
										    lowercase
										)    ,  (
										    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





@require_tf
class      UpperCamelCase_			(			__a		,						__a		,						unittest.TestCase		):







						'''simple docstring'''



						UpperCamelCase			:     Dict							      =				(TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
						UpperCamelCase			:     str							      =				(TFCTRLLMHeadModel,) if is_tf_available() else ()
						UpperCamelCase			:     Any							      =				(
						    {
						        '''feature-extraction''': TFCTRLModel,
						        '''text-classification''': TFCTRLForSequenceClassification,
						        '''text-generation''': TFCTRLLMHeadModel,
						        '''zero-shot''': TFCTRLForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						UpperCamelCase			:     Optional[Any]							      =				False
						UpperCamelCase			:     str							      =				False





						def 					SCREAMING_SNAKE_CASE( self       :Any		,						lowerCAmelCase__       :Optional[Any]		,						lowerCAmelCase__       :List[Any]		,						lowerCAmelCase__       :Tuple		,						lowerCAmelCase__       :Any		,						lowerCAmelCase__       :Optional[int]    )     ->Dict:
										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 					SCREAMING_SNAKE_CASE( self       :List[str]    )     ->Any:
										lowercase						    =			TFCTRLModelTester(self    )
										lowercase						    =			ConfigTester(self		,						config_class=lowerCAmelCase__		,						n_embd=37    )





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





						def 					SCREAMING_SNAKE_CASE( self       :Any    )     ->str:
										lowercase						    =			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_ctrl_model(*lowerCAmelCase__    )





						def 					SCREAMING_SNAKE_CASE( self       :Optional[int]    )     ->Optional[int]:
										lowercase						    =			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_ctrl_lm_head(*lowerCAmelCase__    )





						def 					SCREAMING_SNAKE_CASE( self       :int    )     ->Dict:
										lowercase						    =			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_ctrl_for_sequence_classification(*lowerCAmelCase__    )





						def 					SCREAMING_SNAKE_CASE( self       :Dict    )     ->Optional[int]:
										lowercase    ,  lowercase						    =			self.model_tester.prepare_config_and_inputs_for_common()
										lowercase						    =			[TFCTRLLMHeadModel]
										lowercase						    =			[TFCTRLForSequenceClassification]

										for model_class in self.all_model_classes:
														lowercase						    =			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:
																		lowercase						    =			model.get_output_embeddings()
																		assert isinstance(lowerCAmelCase__		,						tf.keras.layers.Layer    )
																		lowercase						    =			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:
																		lowercase						    =			model.get_output_embeddings()
																		assert isinstance(lowerCAmelCase__		,						tf.keras.layers.Layer    )
																		lowercase						    =			model.get_bias()
																		assert name is None
														else:
																		lowercase						    =			model.get_output_embeddings()
																		assert x is None
																		lowercase						    =			model.get_bias()
																		assert name is None





						@slow
						def 					SCREAMING_SNAKE_CASE( self       :Optional[int]    )     ->int:
										for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														lowercase						    =			TFCTRLModel.from_pretrained(lowerCAmelCase__    )
														self.assertIsNotNone(lowerCAmelCase__    )





@require_tf
class      UpperCamelCase_			(			unittest.TestCase		):







						'''simple docstring'''



						@slow
						def 					SCREAMING_SNAKE_CASE( self       :List[Any]    )     ->List[Any]:
										lowercase						    =			TFCTRLLMHeadModel.from_pretrained("ctrl"    )
										lowercase						    =			tf.convert_to_tensor([[11859, 0, 1611, 8]]		,						dtype=tf.intaa    )  # Legal the president is
										lowercase						    =			[
										    11859,
										    0,
										    1611,
										    8,
										    5,
										    150,
										    26449,
										    2,
										    19,
										    348,
										    469,
										    3,
										    2595,
										    48,
										    20740,
										    246533,
										    246533,
										    19,
										    30,
										    5,
										]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

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


# This test is meant to be run in on an instance with TPUs like this:
#
#   python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
#
# Replace 8 with the number of TPU cores you have.
#

import sys
from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.utils import logging


_snake_case			:       Optional[Any]      =						logging.get_logger(__name__)


if is_torch_available():
						import torch
						from torch import nn
						from torch.utils.data import Dataset

						from transformers import Trainer





						class      UpperCamelCase_			(			__a		):







												'''simple docstring'''



												def __init__( self       :Optional[int]		,						lowerCAmelCase__       :int = 101    )     ->Any:
																lowercase						    =			length





												def __len__( self       :Tuple    )     ->Any:
																return self.length





												def __getitem__( self       :Any		,						lowerCAmelCase__       :Tuple    )     ->int:
																return i





						class      UpperCamelCase_			:







												'''simple docstring'''



												def __call__( self       :Optional[Any]		,						lowerCAmelCase__       :List[Any]    )     ->int:
																return {"input_ids": torch.tensor(lowerCAmelCase__    ), "labels": torch.tensor(lowerCAmelCase__    )}





						class      UpperCamelCase_			(			nn.Module		):







												'''simple docstring'''



												def __init__( self       :Any    )     ->str:
																super().__init__()
																# Add some (unused) params otherwise DDP will complain.
																lowercase						    =			nn.Linear(120		,						80    )





												def 					SCREAMING_SNAKE_CASE( self       :int		,						lowerCAmelCase__       :Optional[int]		,						lowerCAmelCase__       :Union[str, Any]=None    )     ->List[Any]:
																if labels is not None:
																				return torch.tensor(0.0		,						device=input_ids.device    ), input_ids
																else:
																				return input_ids

def    __snake_case       (  ):







				'''simple docstring'''


				lowercase						    =			HfArgumentParser((TrainingArguments,)	)
				sys.argv += ["--output_dir", "./examples"]
				lowercase						    =			parser.parse_args_into_dataclasses()[0]

				logger.warning(
				    F'''Process rank: {training_args.local_rank}, device: {training_args.device}, '''
				    F'''tpu_num_cores: {training_args.tpu_num_cores}'''		,						)

				# Essentially, what we want to verify in the distributed case is
				# that we get all samples back, in the right order.
				# (this is crucial for prediction for instance)
				for dataset_length in [1001, 256, 15]:
								lowercase						    =			DummyDataset(__magic_name__	)

								def compute_metrics(__magic_name__	) -> Dict:
												lowercase						    =			list(range(len(__magic_name__	)	)	)
												lowercase						    =			p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
												return {"success": success}

								lowercase						    =			Trainer(
								    model=DummyModel()		,						args=__magic_name__		,						data_collator=DummyDataCollator()		,						eval_dataset=__magic_name__		,						compute_metrics=__magic_name__		,						)
								lowercase						    =			trainer.evaluate()
								logger.info(__magic_name__	)
								if metrics["eval_success"] is not True:
												logger.error(__magic_name__	)
												exit(1	)

								lowercase						    =			trainer.predict(__magic_name__	)
								logger.info(p.metrics	)
								if p.metrics["test_success"] is not True:
												logger.error(p.metrics	)
												exit(1	)

								lowercase						    =			2

								lowercase						    =			trainer.evaluate()
								logger.info(__magic_name__	)
								if metrics["eval_success"] is not True:
												logger.error(__magic_name__	)
												exit(1	)

								lowercase						    =			trainer.predict(__magic_name__	)
								logger.info(p.metrics	)
								if p.metrics["test_success"] is not True:
												logger.error(p.metrics	)
												exit(1	)

								lowercase						    =			None

				logger.info("🔥 All distributed tests successful"	)

def    __snake_case       (  __magic_name__	):







				'''simple docstring'''


				main()


if __name__ == "__main__":
						main()
