"""simple docstring"""







from __future__ import annotations

import unittest

from transformers import AutoTokenizer, GPTJConfig, is_tf_available
from transformers.testing_utils import require_tf, slow, tooslow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin
from ...utils.test_modeling_tf_core import TFCoreModelTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers.models.gptj.modeling_tf_gptj import (
						    TFGPTJForCausalLM,
						    TFGPTJForQuestionAnswering,
						    TFGPTJForSequenceClassification,
						    TFGPTJModel,
						    shape_list,
						)



class 			_lowerCamelCase      :







				def __init__(			self     :							Union[str, Any]      ,       lowerCamelCase_     :							Tuple    ):

					"""simple docstring"""






					_lowercase			:     Optional[Any]      					=				parent
					_lowercase			:     Dict      					=				1_3
					_lowercase			:     List[str]      					=				7
					_lowercase			:     List[str]      					=				True
					_lowercase			:     Tuple      					=				True
					_lowercase			:     int      					=				True
					_lowercase			:     Any      					=				True
					_lowercase			:     Tuple      					=				True
					_lowercase			:     str      					=				9_9
					_lowercase			:     Tuple      					=				3_2
					_lowercase			:     List[Any]      					=				4
					_lowercase			:     Union[str, Any]      					=				2
					_lowercase			:     Optional[Any]      					=				4
					_lowercase			:     Union[str, Any]      					=				3_7
					_lowercase			:     Union[str, Any]      					=				'''gelu'''
					_lowercase			:     List[Any]      					=				0.1
					_lowercase			:     Dict      					=				0.1
					_lowercase			:     Tuple      					=				5_1_2
					_lowercase			:     int      					=				1_6
					_lowercase			:     Optional[int]      					=				2
					_lowercase			:     Union[str, Any]      					=				0.02
					_lowercase			:     List[Any]      					=				3
					_lowercase			:     List[str]      					=				4
					_lowercase			:     Tuple      					=				None
					_lowercase			:     Dict      					=				self.vocab_size - 1
					_lowercase			:     Union[str, Any]      					=				self.vocab_size - 1
					_lowercase			:     List[str]      					=				self.vocab_size - 1







				def 				__UpperCAmelCase       (			self     :							List[str]    ):

					"""simple docstring"""






					_lowercase			:     Any      					=				ids_tensor([self.batch_size, self.seq_length]      ,       self.vocab_size    )

					_lowercase			:     Dict      					=				None
					if self.use_input_mask:
						_lowercase			:     List[str]      					=				random_attention_mask([self.batch_size, self.seq_length]    )

					_lowercase			:     Union[str, Any]      					=				None
					if self.use_token_type_ids:
						_lowercase			:     str      					=				ids_tensor([self.batch_size, self.seq_length]      ,       self.type_vocab_size    )

					_lowercase			:     Optional[Any]      					=				None
					if self.use_mc_token_ids:
						_lowercase			:     Union[str, Any]      					=				ids_tensor([self.batch_size, self.num_choices]      ,       self.seq_length    )

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

					_lowercase			:     List[Any]      					=				GPTJConfig(
					    vocab_size=self.vocab_size      ,       n_embd=self.hidden_size      ,       n_layer=self.num_hidden_layers      ,       n_head=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      ,       n_positions=self.max_position_embeddings      ,       type_vocab_size=self.type_vocab_size      ,       initializer_range=self.initializer_range      ,       bos_token_id=self.bos_token_id      ,       eos_token_id=self.eos_token_id      ,       pad_token_id=self.pad_token_id      ,       rotary_dim=self.rotary_dim      ,       return_dict=lowerCAmelCase__      ,       )

					_lowercase			:     Dict      					=				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 				__UpperCAmelCase       (			self     :							Optional[int]      ,       lowerCamelCase_     :							Tuple      ,       lowerCamelCase_     :							int      ,       lowerCamelCase_     :							Dict      ,       lowerCamelCase_     :							Union[str, Any]      ,       lowerCamelCase_     :							Union[str, Any]      ,       *lowerCamelCase_     :							Union[str, Any]    ):

					"""simple docstring"""






					_lowercase			:     str      					=				TFGPTJModel(config=lowerCAmelCase__    )
					_lowercase			:     Union[str, Any]      					=				{
					    '''input_ids''': input_ids,
					    '''attention_mask''': input_mask,
					    '''token_type_ids''': token_type_ids,
					}
					_lowercase			:     Any      					=				model(lowerCAmelCase__    )

					_lowercase			:     Union[str, Any]      					=				[input_ids, None, input_mask]  # None is the input for 'past'
					_lowercase			:     Dict      					=				model(lowerCAmelCase__    )

					_lowercase			:     str      					=				model(lowerCAmelCase__    )

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







				def 				__UpperCAmelCase       (			self     :							Dict      ,       lowerCamelCase_     :							List[Any]      ,       lowerCamelCase_     :							Tuple      ,       lowerCamelCase_     :							Union[str, Any]      ,       lowerCamelCase_     :							List[str]      ,       lowerCamelCase_     :							List[str]      ,       *lowerCamelCase_     :							Union[str, Any]    ):

					"""simple docstring"""






					_lowercase			:     Optional[int]      					=				TFGPTJModel(config=lowerCAmelCase__    )

					# first forward pass
					_lowercase			:     int      					=				model(lowerCAmelCase__      ,       token_type_ids=lowerCAmelCase__      ,       use_cache=lowerCAmelCase__    )
					_lowercase			:     Optional[int]      					=				model(lowerCAmelCase__      ,       token_type_ids=lowerCAmelCase__    )
					_lowercase			:     Tuple      					=				model(lowerCAmelCase__      ,       token_type_ids=lowerCAmelCase__      ,       use_cache=lowerCAmelCase__    )

					self.parent.assertTrue(len(lowerCAmelCase__    ) == len(lowerCAmelCase__    )    )
					self.parent.assertTrue(len(lowerCAmelCase__    ) == len(lowerCAmelCase__    ) + 1    )

					_lowercase			:     str      					=				outputs.to_tuple()

					# create hypothetical next token and extent to next_input_ids
					_lowercase			:     Optional[Any]      					=				ids_tensor((self.batch_size, 1)      ,       config.vocab_size    )
					_lowercase			:     List[str]      					=				ids_tensor([self.batch_size, 1]      ,       self.type_vocab_size    )

					# append to next input_ids and token_type_ids
					_lowercase			:     int      					=				tf.concat([input_ids, next_tokens]      ,       axis=-1    )
					_lowercase			:     List[Any]      					=				tf.concat([token_type_ids, next_token_types]      ,       axis=-1    )

					_lowercase			:     int      					=				model(lowerCAmelCase__      ,       token_type_ids=lowerCAmelCase__    )['''last_hidden_state''']
					_lowercase			:     List[str]      					=				model(lowerCAmelCase__      ,       token_type_ids=lowerCAmelCase__      ,       past_key_values=lowerCAmelCase__    )[
					    '''last_hidden_state'''
					]

					# select random slice
					_lowercase			:     Dict      					=				int(ids_tensor((1,)      ,       shape_list(lowerCAmelCase__    )[-1]    )    )
					_lowercase			:     Optional[int]      					=				output_from_no_past[:, -1, random_slice_idx]
					_lowercase			:     int      					=				output_from_past[:, 0, random_slice_idx]

					# test that outputs are equal for slice
					tf.debugging.assert_near(lowerCAmelCase__      ,       lowerCAmelCase__      ,       rtol=1E-6    )







				def 				__UpperCAmelCase       (			self     :							Optional[Any]      ,       lowerCamelCase_     :							int      ,       lowerCamelCase_     :							Dict      ,       lowerCamelCase_     :							Tuple      ,       lowerCamelCase_     :							Tuple      ,       lowerCamelCase_     :							List[str]      ,       *lowerCamelCase_     :							Tuple    ):

					"""simple docstring"""






					_lowercase			:     Optional[Any]      					=				TFGPTJModel(config=lowerCAmelCase__    )

					# create attention mask
					_lowercase			:     Union[str, Any]      					=				self.seq_length // 2
					_lowercase			:     int      					=				tf.ones((self.batch_size, half_seq_length)      ,       dtype=tf.intaa    )
					_lowercase			:     Any      					=				tf.zeros((self.batch_size, self.seq_length - half_seq_length)      ,       dtype=tf.intaa    )
					_lowercase			:     int      					=				tf.concat([attn_mask_begin, attn_mask_end]      ,       axis=1    )

					# first forward pass
					_lowercase			:     int      					=				model(lowerCAmelCase__      ,       attention_mask=lowerCAmelCase__    ).to_tuple()

					# create hypothetical next token and extent to next_input_ids
					_lowercase			:     Dict      					=				ids_tensor((self.batch_size, 1)      ,       config.vocab_size    )

					# change a random masked slice from input_ids
					_lowercase			:     List[Any]      					=				ids_tensor((1,)      ,       lowerCAmelCase__    ).numpy() + 1
					_lowercase			:     Tuple      					=				ids_tensor((self.batch_size, self.seq_length)      ,       config.vocab_size    )
					_lowercase			:     List[str]      					=				tf.range(self.seq_length    ) == (self.seq_length - random_seq_idx_to_change)
					_lowercase			:     Any      					=				tf.transpose(
					    tf.broadcast_to(tf.expand_dims(lowerCAmelCase__      ,       -1    )      ,       (self.seq_length, self.batch_size)    )    )
					_lowercase			:     Any      					=				tf.where(lowerCAmelCase__      ,       lowerCAmelCase__      ,       lowerCAmelCase__    )

					# append to next input_ids and attn_mask
					_lowercase			:     str      					=				tf.concat([input_ids, next_tokens]      ,       axis=-1    )
					_lowercase			:     int      					=				tf.concat([attn_mask, tf.ones((shape_list(lowerCAmelCase__    )[0], 1)      ,       dtype=tf.intaa    )]      ,       axis=1    )

					# get two different outputs
					_lowercase			:     str      					=				model(lowerCAmelCase__      ,       attention_mask=lowerCAmelCase__    )['''last_hidden_state''']
					_lowercase			:     int      					=				model(lowerCAmelCase__      ,       past_key_values=lowerCAmelCase__      ,       attention_mask=lowerCAmelCase__    )[
					    '''last_hidden_state'''
					]

					# select random slice
					_lowercase			:     List[str]      					=				int(ids_tensor((1,)      ,       shape_list(lowerCAmelCase__    )[-1]    )    )
					_lowercase			:     Tuple      					=				output_from_no_past[:, -1, random_slice_idx]
					_lowercase			:     Dict      					=				output_from_past[:, 0, random_slice_idx]

					# test that outputs are equal for slice
					tf.debugging.assert_near(lowerCAmelCase__      ,       lowerCAmelCase__      ,       rtol=1E-12    )







				def 				__UpperCAmelCase       (			self     :							str      ,       lowerCamelCase_     :							List[Any]      ,       lowerCamelCase_     :							List[str]      ,       lowerCamelCase_     :							Optional[int]      ,       lowerCamelCase_     :							Dict      ,       lowerCamelCase_     :							Any      ,       *lowerCamelCase_     :							Dict    ):

					"""simple docstring"""






					_lowercase			:     Tuple      					=				TFGPTJModel(config=lowerCAmelCase__    )

					_lowercase			:     List[str]      					=				input_ids[:1, :]
					_lowercase			:     Any      					=				input_mask[:1, :]
					_lowercase			:     int      					=				token_type_ids[:1, :]
					_lowercase			:     Any      					=				1

					# first forward pass
					_lowercase			:     Optional[Any]      					=				model(lowerCAmelCase__      ,       attention_mask=lowerCAmelCase__      ,       token_type_ids=lowerCAmelCase__      ,       use_cache=lowerCAmelCase__    )

					_lowercase			:     Tuple      					=				outputs.to_tuple()

					# create hypothetical next token and extent to next_input_ids
					_lowercase			:     Tuple      					=				ids_tensor((self.batch_size, 3)      ,       config.vocab_size    )
					_lowercase			:     int      					=				ids_tensor((self.batch_size, 3)      ,       2    )
					_lowercase			:     str      					=				ids_tensor((self.batch_size, 3)      ,       self.type_vocab_size    )

					# append to next input_ids and token_type_ids
					_lowercase			:     str      					=				tf.concat([input_ids, next_tokens]      ,       axis=-1    )
					_lowercase			:     int      					=				tf.concat([input_mask, next_attn_mask]      ,       axis=-1    )
					_lowercase			:     Tuple      					=				tf.concat([token_type_ids, next_token_types]      ,       axis=-1    )

					_lowercase			:     Union[str, Any]      					=				model(
					    lowerCAmelCase__      ,       token_type_ids=lowerCAmelCase__      ,       attention_mask=lowerCAmelCase__    )['''last_hidden_state''']
					_lowercase			:     Optional[Any]      					=				model(
					    lowerCAmelCase__      ,       token_type_ids=lowerCAmelCase__      ,       attention_mask=lowerCAmelCase__      ,       past_key_values=lowerCAmelCase__      ,       )['''last_hidden_state''']
					self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]    )

					# select random slice
					_lowercase			:     Tuple      					=				int(ids_tensor((1,)      ,       shape_list(lowerCAmelCase__    )[-1]    )    )
					_lowercase			:     int      					=				output_from_no_past[:, -3:, random_slice_idx]
					_lowercase			:     List[Any]      					=				output_from_past[:, :, random_slice_idx]

					# test that outputs are equal for slice
					tf.debugging.assert_near(lowerCAmelCase__      ,       lowerCAmelCase__      ,       rtol=1E-3    )







				def 				__UpperCAmelCase       (			self     :							str      ,       lowerCamelCase_     :							Any      ,       lowerCamelCase_     :							Optional[Any]      ,       lowerCamelCase_     :							int      ,       lowerCamelCase_     :							Any      ,       lowerCamelCase_     :							Union[str, Any]      ,       *lowerCamelCase_     :							Dict    ):

					"""simple docstring"""






					_lowercase			:     List[Any]      					=				TFGPTJForCausalLM(config=lowerCAmelCase__    )
					_lowercase			:     Optional[Any]      					=				{
					    '''input_ids''': input_ids,
					    '''attention_mask''': input_mask,
					    '''token_type_ids''': token_type_ids,
					}
					_lowercase			:     Any      					=				model(lowerCAmelCase__    )
					self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.seq_length, self.vocab_size)    )







				def 				__UpperCAmelCase       (			self     :							List[str]    ):

					"""simple docstring"""






					_lowercase			:     int      					=				self.prepare_config_and_inputs()

					(
					    _lowercase				
					)			:     List[str]      					=				config_and_inputs

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



@require_tf
class 			_lowerCamelCase      (_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase    ,				unittest.TestCase  ):
				_snake_case  =	(
				    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
				    if is_tf_available()
				    else ()
				)

				_snake_case  =	(TFGPTJForCausalLM,) if is_tf_available() else ()
				_snake_case  =	(
				    {
				        "feature-extraction": TFGPTJModel,
				        "question-answering": TFGPTJForQuestionAnswering,
				        "text-classification": TFGPTJForSequenceClassification,
				        "text-generation": TFGPTJForCausalLM,
				        "zero-shot": TFGPTJForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				_snake_case  =	False
				_snake_case  =	False
				_snake_case  =	False
				_snake_case  =	False







				def 				__UpperCAmelCase       (			self     :							Any      ,       lowerCamelCase_     :							int      ,       lowerCamelCase_     :							str      ,       lowerCamelCase_     :							Optional[int]      ,       lowerCamelCase_     :							List[str]      ,       lowerCamelCase_     :							List[str]    ):

					"""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 				__UpperCAmelCase       (			self     :							Tuple    ):

					"""simple docstring"""






					_lowercase			:     Optional[int]      					=				TFGPTJModelTester(self    )
					_lowercase			:     Optional[Any]      					=				ConfigTester(self      ,       config_class=lowerCAmelCase__      ,       n_embd=3_7    )







				def 				__UpperCAmelCase       (			self     :							Optional[Any]    ):

					"""simple docstring"""






					self.config_tester.run_common_tests()







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

					"""simple docstring"""






					_lowercase			:     List[Any]      					=				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_gptj_model(*lowerCAmelCase__    )







				def 				__UpperCAmelCase       (			self     :							Optional[Any]    ):

					"""simple docstring"""






					_lowercase			:     Any      					=				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_gptj_model_past(*lowerCAmelCase__    )







				def 				__UpperCAmelCase       (			self     :							Optional[int]    ):

					"""simple docstring"""






					_lowercase			:     str      					=				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_gptj_model_attention_mask_past(*lowerCAmelCase__    )







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

					"""simple docstring"""






					_lowercase			:     Dict      					=				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_gptj_model_past_large_inputs(*lowerCAmelCase__    )







				def 				__UpperCAmelCase       (			self     :							Dict    ):

					"""simple docstring"""






					_lowercase			:     Tuple      					=				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_gptj_lm_head_model(*lowerCAmelCase__    )







				@slow
				@unittest.skipIf(
				    not is_tf_available() or len(tf.config.list_physical_devices('GPU'    )    ) > 0      ,       'skip testing on GPU for now to avoid GPU OOM.'      ,       )
				def 				__UpperCAmelCase       (			self     :							Any    ):

					"""simple docstring"""






					_lowercase			:     List[str]      					=				TFGPTJModel.from_pretrained('EleutherAI/gpt-j-6B'      ,       from_pt=lowerCAmelCase__    )
					self.assertIsNotNone(lowerCAmelCase__    )







				@unittest.skip(reason='Currently, model embeddings are going to undergo a major refactor.'    )
				def 				__UpperCAmelCase       (			self     :							Union[str, Any]    ):

					"""simple docstring"""






					super().test_resize_token_embeddings()



@require_tf
@tooslow
# Marked as @tooslow due to GPU OOM -- but still useful to run locally. Requires ~39GB of RAM.
class 			_lowerCamelCase      (unittest.TestCase  ):







				def 				__UpperCAmelCase       (			self     :							Optional[int]    ):

					"""simple docstring"""






					_lowercase			:     List[str]      					=				TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'      ,       from_pt=lowerCAmelCase__    )
					_lowercase			:     List[Any]      					=				tf.convert_to_tensor([[4_6_4, 3_2_9_0]]      ,       dtype=tf.intaa    )  # The dog
					# fmt: off
					# The dog is a man's best friend. It is a loyal companion, and it is a friend
					_lowercase			:     Any      					=				[4_6_4, 3_2_9_0, 3_1_8, 2_5_7, 5_8_2, 3_3_8, 1_2_6_6, 1_5_4_5, 1_3, 6_3_2, 3_1_8, 2_5_7, 9_1_1_2, 1_5_1_8_5, 1_1, 2_9_0, 3_4_0, 3_1_8, 2_5_7, 1_5_4_5]
					# fmt: on
					_lowercase			:     Tuple      					=				model.generate(lowerCAmelCase__      ,       do_sample=lowerCAmelCase__    )
					self.assertListEqual(output_ids[0].numpy().tolist()      ,       lowerCAmelCase__    )







				def 				__UpperCAmelCase       (			self     :							List[Any]    ):

					"""simple docstring"""






					_lowercase			:     Optional[int]      					=				AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B'      ,       revision='float16'    )
					_lowercase			:     List[Any]      					=				TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'      ,       revision='float16'      ,       from_pt=lowerCAmelCase__    )

					_lowercase			:     Dict      					=				tokenizer('Today is a nice day and'      ,       return_tensors='tf'    )
					# forces the generation to happen on CPU, to avoid GPU-related quirks
					with tf.device(':/CPU:0'    ):
						_lowercase			:     Tuple      					=				model.generate(**lowerCAmelCase__      ,       do_sample=lowerCAmelCase__      ,       seed=[4_2, 0]    )
					_lowercase			:     List[str]      					=				tokenizer.decode(output_ids[0]      ,       skip_special_tokens=lowerCAmelCase__    )

					_lowercase			:     List[str]      					=				'''Today is a nice day and I’m going to go for a walk. I’'''
					self.assertEqual(lowerCAmelCase__      ,       lowerCAmelCase__    )







				def 				__UpperCAmelCase       (			self     :							int    ):

					"""simple docstring"""






					_lowercase			:     List[str]      					=				TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'      ,       revision='float16'      ,       from_pt=lowerCAmelCase__    )
					_lowercase			:     Optional[int]      					=				AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B'      ,       revision='float16'    )

					_lowercase			:     Tuple      					=				'''left'''

					# Define PAD Token = EOS Token = 50256
					_lowercase			:     List[str]      					=				tokenizer.eos_token
					_lowercase			:     Union[str, Any]      					=				model.config.eos_token_id

					# use different length sentences to test batching
					_lowercase			:     Union[str, Any]      					=				[
					    '''Hello, my dog is a little''',
					    '''Today, I''',
					]
					_lowercase			:     int      					=				[
					    '''Hello, my dog is a little over a year old and has been diagnosed with hip dysplasia''',
					    '''Today, I’m going to be talking about a topic that’''',
					]
					return model, tokenizer, sentences, expected_output_sentences







				def 				__UpperCAmelCase       (			self     :							Any    ):

					"""simple docstring"""






					_lowercase			:     Optional[Any]      					=				self._get_beam_search_test_objects()

					_lowercase			:     List[str]      					=				tokenizer(lowerCAmelCase__      ,       return_tensors='tf'      ,       padding=lowerCAmelCase__    )
					_lowercase			:     Optional[Any]      					=				model.generate(**lowerCAmelCase__      ,       do_sample=lowerCAmelCase__      ,       num_beams=2    )
					_lowercase			:     Union[str, Any]      					=				tokenizer.batch_decode(lowerCAmelCase__      ,       skip_special_tokens=lowerCAmelCase__    )
					self.assertListEqual(lowerCAmelCase__      ,       lowerCAmelCase__    )







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

					"""simple docstring"""






					_lowercase			:     Any      					=				self._get_beam_search_test_objects()

					_lowercase			:     Optional[Any]      					=				tokenizer(lowerCAmelCase__      ,       return_tensors='tf'      ,       padding=lowerCAmelCase__    )
					_lowercase			:     Dict      					=				tokenizer(sentences[0]      ,       return_tensors='tf'    )
					_lowercase			:     Any      					=				model.generate(**lowerCAmelCase__      ,       do_sample=lowerCAmelCase__      ,       num_beams=2    )
					_lowercase			:     Union[str, Any]      					=				(
					    shape_list(inputs_non_padded['input_ids']    )[-1]
					    - tf.reduce_sum(tf.cast(inputs['attention_mask'][-1]      ,       tf.intaa    )    ).numpy()
					)
					_lowercase			:     Any      					=				tokenizer(sentences[1]      ,       return_tensors='tf'    )
					_lowercase			:     Tuple      					=				model.generate(
					    **lowerCAmelCase__      ,       do_sample=lowerCAmelCase__      ,       num_beams=2      ,       max_length=model.config.max_length - num_paddings    )
					_lowercase			:     Dict      					=				tokenizer.decode(output_non_padded[0]      ,       skip_special_tokens=lowerCAmelCase__    )
					_lowercase			:     List[str]      					=				tokenizer.decode(output_padded[0]      ,       skip_special_tokens=lowerCAmelCase__    )
					self.assertListEqual(lowerCAmelCase__      ,       [non_padded_sentence, padded_sentence]    )







				def 				__UpperCAmelCase       (			self     :							List[Any]    ):

					"""simple docstring"""






					_lowercase			:     str      					=				self._get_beam_search_test_objects()

					_lowercase			:     List[str]      					=				tokenizer(lowerCAmelCase__      ,       return_tensors='tf'      ,       padding=lowerCAmelCase__    )
					_lowercase			:     str      					=				tf.function(model.generate      ,       jit_compile=lowerCAmelCase__    )
					_lowercase			:     str      					=				xla_generate(**lowerCAmelCase__      ,       do_sample=lowerCAmelCase__      ,       num_beams=2    )
					_lowercase			:     Optional[int]      					=				tokenizer.batch_decode(lowerCAmelCase__      ,       skip_special_tokens=lowerCAmelCase__    )
					self.assertListEqual(lowerCAmelCase__      ,       lowerCAmelCase__    )










"""simple docstring"""




import argparse

import torch
from clip import load

from transformers import CLIPConfig, CLIPModel





def 			a_							(				_lowercase			,		_lowercase  ):
						_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase     :   Any      =       pt_attn_layer.in_proj_weight.chunk(3			,		dim=0  )
						_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase     :   Dict      =       pt_attn_layer.in_proj_bias.chunk(3			,		dim=0  )

						_UpperCamelCase     :   int      =       pt_attn_layer.out_proj.weight
						_UpperCamelCase     :   List[Any]      =       pt_attn_layer.out_proj.bias

						_UpperCamelCase     :   Any      =       q_proj
						_UpperCamelCase     :   str      =       q_proj_bias

						_UpperCamelCase     :   str      =       k_proj
						_UpperCamelCase     :   List[str]      =       k_proj_bias

						_UpperCamelCase     :   Dict      =       v_proj
						_UpperCamelCase     :   Dict      =       v_proj_bias

						_UpperCamelCase     :   Dict      =       out_proj_weights
						_UpperCamelCase     :   Optional[Any]      =       out_proj_bias





def 			a_							(				_lowercase			,		_lowercase  ):
						copy_linear(hf_mlp.fca			,		pt_mlp.c_fc  )
						copy_linear(hf_mlp.fca			,		pt_mlp.c_proj  )





def 			a_							(				_lowercase			,		_lowercase  ):
						_UpperCamelCase     :   str      =       pt_linear.weight
						_UpperCamelCase     :   List[str]      =       pt_linear.bias





def 			a_							(				_lowercase			,		_lowercase  ):
						# copy layer norms
						copy_linear(hf_layer.layer_norma			,		pt_layer.ln_a  )
						copy_linear(hf_layer.layer_norma			,		pt_layer.ln_a  )

						# copy MLP
						copy_mlp(hf_layer.mlp			,		pt_layer.mlp  )

						# copy attn
						copy_attn_layer(hf_layer.self_attn			,		pt_layer.attn  )





def 			a_							(				_lowercase			,		_lowercase  ):
						for hf_layer, pt_layer in zip(_lowercase			,		_lowercase  ):
												copy_layer(_lowercase			,		_lowercase  )





def 			a_							(				_lowercase			,		_lowercase  ):
						# copy  embeds
						_UpperCamelCase     :   Tuple      =       pt_model.token_embedding.weight
						_UpperCamelCase     :   List[Any]      =       pt_model.positional_embedding

						# copy layer norm
						copy_linear(hf_encoder.final_layer_norm			,		pt_model.ln_final  )

						# copy hidden layers
						copy_layers(hf_encoder.encoder.layers			,		pt_model.transformer.resblocks  )





def 			a_							(				_lowercase			,		_lowercase  ):
						# copy projection
						_UpperCamelCase     :   str      =       pt_model.text_projection.data.T

						# copy text encoder
						copy_encoder(hf_model.text_model			,		_lowercase  )





def 			a_							(				_lowercase			,		_lowercase  ):
						# copy projection
						_UpperCamelCase     :   Tuple      =       pt_model.visual.proj.data.T

						# copy layer norms
						copy_linear(hf_model.vision_model.pre_layrnorm			,		pt_model.visual.ln_pre  )
						copy_linear(hf_model.vision_model.post_layernorm			,		pt_model.visual.ln_post  )

						# copy embeds
						_UpperCamelCase     :   str      =       pt_model.visual.conva.weight.data
						_UpperCamelCase     :   Dict      =       pt_model.visual.class_embedding
						_UpperCamelCase     :   Union[str, Any]      =       pt_model.visual.positional_embedding.data

						# copy encoder
						copy_layers(hf_model.vision_model.encoder.layers			,		pt_model.visual.transformer.resblocks  )





@torch.no_grad()
def 			a_							(				_lowercase			,		_lowercase			,		_lowercase=None  ):
						if config_path is not None:
												_UpperCamelCase     :   List[Any]      =       CLIPConfig.from_pretrained(_lowercase  )
						else:
												_UpperCamelCase     :   Optional[int]      =       CLIPConfig(projection_dim=512			,		text_config={}			,		vision_config={}  )

						_UpperCamelCase     :   str      =       CLIPModel(_lowercase  ).eval()

						_UpperCamelCase  ,				_UpperCamelCase     :   Optional[int]      =       load(_lowercase			,		device='''cpu'''			,		jit=_lowercase  )
						_UpperCamelCase     :   Union[str, Any]      =       pt_model.eval()

						copy_text_model_and_projection(_lowercase			,		_lowercase  )
						copy_vison_model_and_projection(_lowercase			,		_lowercase  )
						_UpperCamelCase     :   int      =       pt_model.logit_scale

						_UpperCamelCase     :   List[Any]      =       torch.arange(0			,		77  ).unsqueeze(0  )
						_UpperCamelCase     :   Optional[int]      =       torch.randn(1			,		3			,		224			,		224  )

						_UpperCamelCase     :   Tuple      =       hf_model(input_ids=_lowercase			,		pixel_values=_lowercase			,		return_dict=_lowercase  )
						_UpperCamelCase     :   Dict      =       hf_outputs.logits_per_image
						_UpperCamelCase     :   Any      =       hf_outputs.logits_per_text
						_UpperCamelCase  ,				_UpperCamelCase     :   Any      =       pt_model(_lowercase			,		_lowercase  )

						assert torch.allclose(_lowercase			,		_lowercase			,		atol=1E-3  )
						assert torch.allclose(_lowercase			,		_lowercase			,		atol=1E-3  )

						hf_model.save_pretrained(_lowercase  )


if __name__ == "__main__":
						UpperCamelCase_	 =argparse.ArgumentParser()
						parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
						parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to fairseq checkpoint""")
						parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
						UpperCamelCase_	 =parser.parse_args()

						convert_clip_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)



