


from __future__ import annotations

import unittest

from transformers import 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 (
			    TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TFXLMForMultipleChoice,
			    TFXLMForQuestionAnsweringSimple,
			    TFXLMForSequenceClassification,
			    TFXLMForTokenClassification,
			    TFXLMModel,
			    TFXLMWithLMHeadModel,
			    XLMConfig,
			)




class      SCREAMING_SNAKE_CASE__  :





	'''simple docstring'''



	def __init__(						self			:					Tuple    ,						lowercase			:					Optional[Any]    ,						):


				'''simple docstring'''


				_snake_case						=    parent
				_snake_case						=    13
				_snake_case						=    7
				_snake_case						=    True
				_snake_case						=    True
				_snake_case						=    True
				_snake_case						=    True
				_snake_case						=    True
				_snake_case						=    False
				_snake_case						=    False
				_snake_case						=    False
				_snake_case						=    2
				_snake_case						=    99
				_snake_case						=    0
				_snake_case						=    32
				_snake_case						=    2
				_snake_case						=    4
				_snake_case						=    0.1
				_snake_case						=    0.1
				_snake_case						=    512
				_snake_case						=    16
				_snake_case						=    2
				_snake_case						=    0.02
				_snake_case						=    3
				_snake_case						=    4
				_snake_case						=    'last'
				_snake_case						=    True
				_snake_case						=    None
				_snake_case						=    0



	def 				A							(						self			:					Dict			):


				'''simple docstring'''


				_snake_case						=    ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size			)
				_snake_case						=    random_attention_mask([self.batch_size, self.seq_length]    ,						dtype=tf.floataa			)

				_snake_case						=    None
				if self.use_input_lengths:
							_snake_case						=    (
							    ids_tensor([self.batch_size]    ,						vocab_size=2			) + self.seq_length - 2
							)  # small variation of seq_length

				_snake_case						=    None
				if self.use_token_type_ids:
							_snake_case						=    ids_tensor([self.batch_size, self.seq_length]    ,						self.n_langs			)

				_snake_case						=    None
				_snake_case						=    None
				_snake_case						=    None
				if self.use_labels:
							_snake_case						=    ids_tensor([self.batch_size]    ,						self.type_sequence_label_size			)
							_snake_case						=    ids_tensor([self.batch_size, self.seq_length]    ,						self.num_labels			)
							_snake_case						=    ids_tensor([self.batch_size]    ,						2    ,						dtype=tf.floataa			)
							_snake_case						=    ids_tensor([self.batch_size]    ,						self.num_choices			)

				_snake_case						=    XLMConfig(
				    vocab_size=self.vocab_size    ,						n_special=self.n_special    ,						emb_dim=self.hidden_size    ,						n_layers=self.num_hidden_layers    ,						n_heads=self.num_attention_heads    ,						dropout=self.hidden_dropout_prob    ,						attention_dropout=self.attention_probs_dropout_prob    ,						gelu_activation=self.gelu_activation    ,						sinusoidal_embeddings=self.sinusoidal_embeddings    ,						asm=self.asm    ,						causal=self.causal    ,						n_langs=self.n_langs    ,						max_position_embeddings=self.max_position_embeddings    ,						initializer_range=self.initializer_range    ,						summary_type=self.summary_type    ,						use_proj=self.use_proj    ,						bos_token_id=self.bos_token_id    ,						)

				return (
				    config,
				    input_ids,
				    token_type_ids,
				    input_lengths,
				    sequence_labels,
				    token_labels,
				    is_impossible_labels,
				    choice_labels,
				    input_mask,
				)



	def 				A							(						self			:					List[str]    ,						lowercase			:					List[str]    ,						lowercase			:					Optional[Any]    ,						lowercase			:					Optional[int]    ,						lowercase			:					int    ,						lowercase			:					Dict    ,						lowercase			:					Union[str, Any]    ,						lowercase			:					Any    ,						lowercase			:					Any    ,						lowercase			:					Any    ,						):


				'''simple docstring'''


				_snake_case						=    TFXLMModel(config=lowercase			)
				_snake_case						=    {'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
				_snake_case						=    model(lowercase			)

				_snake_case						=    [input_ids, input_mask]
				_snake_case						=    model(lowercase			)
				self.parent.assertEqual(result.last_hidden_state.shape    ,						(self.batch_size, self.seq_length, self.hidden_size)			)



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


				'''simple docstring'''


				_snake_case						=    TFXLMWithLMHeadModel(lowercase			)

				_snake_case						=    {'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
				_snake_case						=    model(lowercase			)

				_snake_case						=    outputs

				self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)			)



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


				'''simple docstring'''


				_snake_case						=    TFXLMForQuestionAnsweringSimple(lowercase			)

				_snake_case						=    {'input_ids': input_ids, 'lengths': input_lengths}

				_snake_case						=    model(lowercase			)

				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			:					Any    ,						lowercase			:					Optional[int]    ,						lowercase			:					Tuple    ,						lowercase			:					List[Any]    ,						lowercase			:					Tuple    ,						lowercase			:					Tuple    ,						lowercase			:					Any    ,						lowercase			:					Union[str, Any]    ,						lowercase			:					Tuple    ,						lowercase			:					List[str]    ,						):


				'''simple docstring'''


				_snake_case						=    TFXLMForSequenceClassification(lowercase			)

				_snake_case						=    {'input_ids': input_ids, 'lengths': input_lengths}

				_snake_case						=    model(lowercase			)

				self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.type_sequence_label_size)			)



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


				'''simple docstring'''


				_snake_case						=    self.num_labels
				_snake_case						=    TFXLMForTokenClassification(config=lowercase			)
				_snake_case						=    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
				_snake_case						=    model(lowercase			)
				self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.num_labels)			)



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


				'''simple docstring'''


				_snake_case						=    self.num_choices
				_snake_case						=    TFXLMForMultipleChoice(config=lowercase			)
				_snake_case						=    tf.tile(tf.expand_dims(lowercase    ,						1			)    ,						(1, self.num_choices, 1)			)
				_snake_case						=    tf.tile(tf.expand_dims(lowercase    ,						1			)    ,						(1, self.num_choices, 1)			)
				_snake_case						=    tf.tile(tf.expand_dims(lowercase    ,						1			)    ,						(1, self.num_choices, 1)			)
				_snake_case						=    {
				    'input_ids': multiple_choice_inputs_ids,
				    'attention_mask': multiple_choice_input_mask,
				    'token_type_ids': multiple_choice_token_type_ids,
				}
				_snake_case						=    model(lowercase			)
				self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_choices)			)





	def 				A							(						self			:					List[Any]			):


				'''simple docstring'''


				_snake_case						=    self.prepare_config_and_inputs()
				(
				    (
				    _snake_case      
				)			,      (
				    _snake_case      
				)			,      (
				    _snake_case      
				)			,      (
				    _snake_case      
				)			,      (
				    _snake_case      
				)			,      (
				    _snake_case      
				)			,      (
				    _snake_case      
				)			,      (
				    _snake_case      
				)			,      (
				    _snake_case      
				)			,            
				)						=    config_and_inputs
				_snake_case						=    {
				    'input_ids': input_ids,
				    'token_type_ids': token_type_ids,
				    'langs': token_type_ids,
				    'lengths': input_lengths,
				}
				return config, inputs_dict




@require_tf
class      SCREAMING_SNAKE_CASE__  (			UpperCAmelCase							,UpperCAmelCase							,unittest.TestCase      ):





	'''simple docstring'''



	_UpperCAmelCase :	Optional[int]					     =      (
	    (
	        TFXLMModel,
	        TFXLMWithLMHeadModel,
	        TFXLMForSequenceClassification,
	        TFXLMForQuestionAnsweringSimple,
	        TFXLMForTokenClassification,
	        TFXLMForMultipleChoice,
	    )
	    if is_tf_available()
	    else ()
	)
	_UpperCAmelCase :	Tuple					     =      (
	    (TFXLMWithLMHeadModel,) if is_tf_available() else ()
	)  # TODO (PVP): Check other models whether language generation is also applicable
	_UpperCAmelCase :	int					     =      (
	    {
	        "feature-extraction": TFXLMModel,
	        "fill-mask": TFXLMWithLMHeadModel,
	        "question-answering": TFXLMForQuestionAnsweringSimple,
	        "text-classification": TFXLMForSequenceClassification,
	        "text-generation": TFXLMWithLMHeadModel,
	        "token-classification": TFXLMForTokenClassification,
	        "zero-shot": TFXLMForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	_UpperCAmelCase :	List[Any]					     =      False
	_UpperCAmelCase :	str					     =      False



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


				'''simple docstring'''


				if (
				    pipeline_test_casse_name == "QAPipelineTests"
				    and tokenizer_name is not None
				    and not tokenizer_name.endswith('Fast'			)
				):
							# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
							# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
							# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
							return True

				return False



	def 				A							(						self			:					Optional[Any]			):


				'''simple docstring'''


				_snake_case						=    TFXLMModelTester(self			)
				_snake_case						=    ConfigTester(self    ,						config_class=lowercase    ,						emb_dim=37			)



	def 				A							(						self			:					str			):


				'''simple docstring'''


				self.config_tester.run_common_tests()



	def 				A							(						self			:					Any			):


				'''simple docstring'''


				_snake_case						=    self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_model(*lowercase			)



	def 				A							(						self			:					Dict			):


				'''simple docstring'''


				_snake_case						=    self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_lm_head(*lowercase			)



	def 				A							(						self			:					Dict			):


				'''simple docstring'''


				_snake_case						=    self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_qa(*lowercase			)



	def 				A							(						self			:					int			):


				'''simple docstring'''


				_snake_case						=    self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_sequence_classif(*lowercase			)



	def 				A							(						self			:					Any			):


				'''simple docstring'''


				_snake_case						=    self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_for_token_classification(*lowercase			)



	def 				A							(						self			:					Optional[Any]			):


				'''simple docstring'''


				_snake_case						=    self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_xlm_for_multiple_choice(*lowercase			)





	@slow
	def 				A							(						self			:					Optional[int]			):


				'''simple docstring'''


				for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							_snake_case						=    TFXLMModel.from_pretrained(lowercase			)
							self.assertIsNotNone(lowercase			)




@require_tf
class      SCREAMING_SNAKE_CASE__  (			unittest.TestCase      ):





	'''simple docstring'''



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


				'''simple docstring'''


				_snake_case						=    TFXLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048'			)
				_snake_case						=    tf.convert_to_tensor([[14, 447]]    ,						dtype=tf.intaa			)  # the president
				_snake_case						=    [
				    14,
				    447,
				    14,
				    447,
				    14,
				    447,
				    14,
				    447,
				    14,
				    447,
				    14,
				    447,
				    14,
				    447,
				    14,
				    447,
				    14,
				    447,
				    14,
				    447,
				]  # the president the president the president the president the president the president the president the president the president the president
				# TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
				_snake_case						=    model.generate(lowercase    ,						do_sample=lowercase			)
				self.assertListEqual(output_ids[0].numpy().tolist()    ,						lowercase			)

import argparse

import torch

from transformers import ImageGPTConfig, ImageGPTForCausalLM, load_tf_weights_in_imagegpt
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()






def  a_	( __lowercase     :			Optional[Any] ,      __lowercase     :			Union[str, Any] ,      __lowercase     :			str				)				->						Any:
			# Construct configuration depending on size
			_snake_case						=    {'small': (512, 8, 24), 'medium': (1_024, 8, 36), 'large': (1_536, 16, 48)}
			_snake_case			,      _snake_case			,      _snake_case						=    MODELS[model_size]  # set model hyperparameters
			_snake_case						=    ImageGPTConfig(n_embd=__lowercase ,      n_layer=__lowercase ,      n_head=__lowercase				)
			_snake_case						=    ImageGPTForCausalLM(__lowercase				)

			# Load weights from numpy
			load_tf_weights_in_imagegpt(__lowercase ,      __lowercase ,      __lowercase				)

			# Save pytorch-model
			_snake_case						=    pytorch_dump_folder_path + '/' + WEIGHTS_NAME
			_snake_case						=    pytorch_dump_folder_path + '/' + CONFIG_NAME
			print(f'''Save PyTorch model to {pytorch_weights_dump_path}'''				)
			torch.save(model.state_dict() ,      __lowercase				)
			print(f'''Save configuration file to {pytorch_config_dump_path}'''				)
			with open(__lowercase ,      'w' ,      encoding='utf-8'				) as f:
						f.write(config.to_json_string()				)


if __name__ == "__main__":
			_lowerCamelCase :       List[Any]     					=    argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '''--imagegpt_checkpoint_path''',
			    default=None,
			    type=str,
			    required=True,
			    help='''Path to the TensorFlow checkpoint path.''',
			)
			parser.add_argument(
			    '''--model_size''',
			    default=None,
			    type=str,
			    required=True,
			    help='''Size of the model (can be either \'small\', \'medium\' or \'large\').''',
			)
			parser.add_argument(
			    '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
			)
			_lowerCamelCase :       Optional[Any]     					=    parser.parse_args()
			convert_imagegpt_checkpoint_to_pytorch(
			    args.imagegpt_checkpoint_path, args.model_size, args.pytorch_dump_folder_path
			)