code: 
  
  
  
  
  
  '''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   _UpperCAmelCase	:
  
  
  
  
  					def __init__(						self    ,      __lowerCAmelCase       )						->  List[str]:
  							UpperCamelCase__         =	parent
  							UpperCamelCase__         =	13
  							UpperCamelCase__         =	7
  							UpperCamelCase__         =	True
  							UpperCamelCase__         =	True
  							UpperCamelCase__         =	True
  							UpperCamelCase__         =	True
  							UpperCamelCase__         =	True
  							UpperCamelCase__         =	99
  							UpperCamelCase__         =	32
  							UpperCamelCase__         =	4
  							UpperCamelCase__         =	2
  							UpperCamelCase__         =	4
  							UpperCamelCase__         =	37
  							UpperCamelCase__         =	"""gelu"""
  							UpperCamelCase__         =	0.1
  							UpperCamelCase__         =	0.1
  							UpperCamelCase__         =	512
  							UpperCamelCase__         =	16
  							UpperCamelCase__         =	2
  							UpperCamelCase__         =	0.0_2
  							UpperCamelCase__         =	3
  							UpperCamelCase__         =	4
  							UpperCamelCase__         =	None
  							UpperCamelCase__         =	self.vocab_size - 1
  							UpperCamelCase__         =	self.vocab_size - 1
  							UpperCamelCase__         =	self.vocab_size - 1
  
  
  
  
  					def 	snake_case_		(						self       )						->  Any:
  							UpperCamelCase__         =	ids_tensor([self.batch_size, self.seq_length]    ,      self.vocab_size       )
  
  							UpperCamelCase__         =	None
  							if self.use_input_mask:
  									UpperCamelCase__         =	random_attention_mask([self.batch_size, self.seq_length]       )
  
  							UpperCamelCase__         =	None
  							if self.use_token_type_ids:
  									UpperCamelCase__         =	ids_tensor([self.batch_size, self.seq_length]    ,      self.type_vocab_size       )
  
  							UpperCamelCase__         =	None
  							if self.use_mc_token_ids:
  									UpperCamelCase__         =	ids_tensor([self.batch_size, self.num_choices]    ,      self.seq_length       )
  
  							UpperCamelCase__         =	None
  							UpperCamelCase__         =	None
  							UpperCamelCase__         =	None
  							if self.use_labels:
  									UpperCamelCase__         =	ids_tensor([self.batch_size]    ,      self.type_sequence_label_size       )
  									UpperCamelCase__         =	ids_tensor([self.batch_size, self.seq_length]    ,      self.num_labels       )
  									UpperCamelCase__         =	ids_tensor([self.batch_size]    ,      self.num_choices       )
  
  							UpperCamelCase__         =	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    ,      )
  
  							UpperCamelCase__         =	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 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      *__lowerCAmelCase       )						->  Optional[int]:
  							UpperCamelCase__         =	TFGPTJModel(config=__lowerCAmelCase       )
  							UpperCamelCase__         =	{
  							    """input_ids""": input_ids,
  							    """attention_mask""": input_mask,
  							    """token_type_ids""": token_type_ids,
  							}
  							UpperCamelCase__         =	model(__lowerCAmelCase       )
  
  							UpperCamelCase__         =	[input_ids, None, input_mask]  # None is the input for 'past'
  							UpperCamelCase__         =	model(__lowerCAmelCase       )
  
  							UpperCamelCase__         =	model(__lowerCAmelCase       )
  
  							self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, self.seq_length, self.hidden_size)       )
  
  
  
  
  					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      *__lowerCAmelCase       )						->  Tuple:
  							UpperCamelCase__         =	TFGPTJModel(config=__lowerCAmelCase       )
  
  							# first forward pass
  							UpperCamelCase__         =	model(__lowerCAmelCase    ,      token_type_ids=__lowerCAmelCase    ,      use_cache=__lowerCAmelCase       )
  							UpperCamelCase__         =	model(__lowerCAmelCase    ,      token_type_ids=__lowerCAmelCase       )
  							UpperCamelCase__         =	model(__lowerCAmelCase    ,      token_type_ids=__lowerCAmelCase    ,      use_cache=__lowerCAmelCase       )
  
  							self.parent.assertTrue(len(__lowerCAmelCase       ) == len(__lowerCAmelCase       )       )
  							self.parent.assertTrue(len(__lowerCAmelCase       ) == len(__lowerCAmelCase       ) + 1       )
  
  							UpperCamelCase__  ,			UpperCamelCase__         =	outputs.to_tuple()
  
  							# create hypothetical next token and extent to next_input_ids
  							UpperCamelCase__         =	ids_tensor((self.batch_size, 1)    ,      config.vocab_size       )
  							UpperCamelCase__         =	ids_tensor([self.batch_size, 1]    ,      self.type_vocab_size       )
  
  							# append to next input_ids and token_type_ids
  							UpperCamelCase__         =	tf.concat([input_ids, next_tokens]    ,      axis=-1       )
  							UpperCamelCase__         =	tf.concat([token_type_ids, next_token_types]    ,      axis=-1       )
  
  							UpperCamelCase__         =	model(__lowerCAmelCase    ,      token_type_ids=__lowerCAmelCase       )["""last_hidden_state"""]
  							UpperCamelCase__         =	model(__lowerCAmelCase    ,      token_type_ids=__lowerCAmelCase    ,      past_key_values=__lowerCAmelCase       )[
  							    """last_hidden_state"""
  							]
  
  							# select random slice
  							UpperCamelCase__         =	int(ids_tensor((1,)    ,      shape_list(__lowerCAmelCase       )[-1]       )       )
  							UpperCamelCase__         =	output_from_no_past[:, -1, random_slice_idx]
  							UpperCamelCase__         =	output_from_past[:, 0, random_slice_idx]
  
  							# test that outputs are equal for slice
  							tf.debugging.assert_near(__lowerCAmelCase    ,      __lowerCAmelCase    ,      rtol=1E-6       )
  
  
  
  
  					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      *__lowerCAmelCase       )						->  List[str]:
  							UpperCamelCase__         =	TFGPTJModel(config=__lowerCAmelCase       )
  
  							# create attention mask
  							UpperCamelCase__         =	self.seq_length // 2
  							UpperCamelCase__         =	tf.ones((self.batch_size, half_seq_length)    ,      dtype=tf.intaa       )
  							UpperCamelCase__         =	tf.zeros((self.batch_size, self.seq_length - half_seq_length)    ,      dtype=tf.intaa       )
  							UpperCamelCase__         =	tf.concat([attn_mask_begin, attn_mask_end]    ,      axis=1       )
  
  							# first forward pass
  							UpperCamelCase__  ,			UpperCamelCase__         =	model(__lowerCAmelCase    ,      attention_mask=__lowerCAmelCase       ).to_tuple()
  
  							# create hypothetical next token and extent to next_input_ids
  							UpperCamelCase__         =	ids_tensor((self.batch_size, 1)    ,      config.vocab_size       )
  
  							# change a random masked slice from input_ids
  							UpperCamelCase__         =	ids_tensor((1,)    ,      __lowerCAmelCase       ).numpy() + 1
  							UpperCamelCase__         =	ids_tensor((self.batch_size, self.seq_length)    ,      config.vocab_size       )
  							UpperCamelCase__         =	tf.range(self.seq_length       ) == (self.seq_length - random_seq_idx_to_change)
  							UpperCamelCase__         =	tf.transpose(
  							    tf.broadcast_to(tf.expand_dims(__lowerCAmelCase    ,      -1       )    ,      (self.seq_length, self.batch_size)       )       )
  							UpperCamelCase__         =	tf.where(__lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase       )
  
  							# append to next input_ids and attn_mask
  							UpperCamelCase__         =	tf.concat([input_ids, next_tokens]    ,      axis=-1       )
  							UpperCamelCase__         =	tf.concat([attn_mask, tf.ones((shape_list(__lowerCAmelCase       )[0], 1)    ,      dtype=tf.intaa       )]    ,      axis=1       )
  
  							# get two different outputs
  							UpperCamelCase__         =	model(__lowerCAmelCase    ,      attention_mask=__lowerCAmelCase       )["""last_hidden_state"""]
  							UpperCamelCase__         =	model(__lowerCAmelCase    ,      past_key_values=__lowerCAmelCase    ,      attention_mask=__lowerCAmelCase       )[
  							    """last_hidden_state"""
  							]
  
  							# select random slice
  							UpperCamelCase__         =	int(ids_tensor((1,)    ,      shape_list(__lowerCAmelCase       )[-1]       )       )
  							UpperCamelCase__         =	output_from_no_past[:, -1, random_slice_idx]
  							UpperCamelCase__         =	output_from_past[:, 0, random_slice_idx]
  
  							# test that outputs are equal for slice
  							tf.debugging.assert_near(__lowerCAmelCase    ,      __lowerCAmelCase    ,      rtol=1E-1_2       )
  
  
  
  
  					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      *__lowerCAmelCase       )						->  Any:
  							UpperCamelCase__         =	TFGPTJModel(config=__lowerCAmelCase       )
  
  							UpperCamelCase__         =	input_ids[:1, :]
  							UpperCamelCase__         =	input_mask[:1, :]
  							UpperCamelCase__         =	token_type_ids[:1, :]
  							UpperCamelCase__         =	1
  
  							# first forward pass
  							UpperCamelCase__         =	model(__lowerCAmelCase    ,      attention_mask=__lowerCAmelCase    ,      token_type_ids=__lowerCAmelCase    ,      use_cache=__lowerCAmelCase       )
  
  							UpperCamelCase__  ,			UpperCamelCase__         =	outputs.to_tuple()
  
  							# create hypothetical next token and extent to next_input_ids
  							UpperCamelCase__         =	ids_tensor((self.batch_size, 3)    ,      config.vocab_size       )
  							UpperCamelCase__         =	ids_tensor((self.batch_size, 3)    ,      2       )
  							UpperCamelCase__         =	ids_tensor((self.batch_size, 3)    ,      self.type_vocab_size       )
  
  							# append to next input_ids and token_type_ids
  							UpperCamelCase__         =	tf.concat([input_ids, next_tokens]    ,      axis=-1       )
  							UpperCamelCase__         =	tf.concat([input_mask, next_attn_mask]    ,      axis=-1       )
  							UpperCamelCase__         =	tf.concat([token_type_ids, next_token_types]    ,      axis=-1       )
  
  							UpperCamelCase__         =	model(
  							    __lowerCAmelCase    ,      token_type_ids=__lowerCAmelCase    ,      attention_mask=__lowerCAmelCase       )["""last_hidden_state"""]
  							UpperCamelCase__         =	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
  							UpperCamelCase__         =	int(ids_tensor((1,)    ,      shape_list(__lowerCAmelCase       )[-1]       )       )
  							UpperCamelCase__         =	output_from_no_past[:, -3:, random_slice_idx]
  							UpperCamelCase__         =	output_from_past[:, :, random_slice_idx]
  
  							# test that outputs are equal for slice
  							tf.debugging.assert_near(__lowerCAmelCase    ,      __lowerCAmelCase    ,      rtol=1E-3       )
  
  
  
  
  					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      *__lowerCAmelCase       )						->  int:
  							UpperCamelCase__         =	TFGPTJForCausalLM(config=__lowerCAmelCase       )
  							UpperCamelCase__         =	{
  							    """input_ids""": input_ids,
  							    """attention_mask""": input_mask,
  							    """token_type_ids""": token_type_ids,
  							}
  							UpperCamelCase__         =	model(__lowerCAmelCase       )
  							self.parent.assertEqual(result.logits.shape    ,      (self.batch_size, self.seq_length, self.vocab_size)       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  Any:
  							UpperCamelCase__         =	self.prepare_config_and_inputs()
  
  							(
  							    (
  							    UpperCamelCase__			
  							)  ,			(
  							    UpperCamelCase__			
  							)  ,			(
  							    UpperCamelCase__			
  							)  ,			(
  							    UpperCamelCase__			
  							)  ,			(
  							    UpperCamelCase__			
  							)  ,			(
  							    UpperCamelCase__			
  							)  ,			(
  							    UpperCamelCase__			
  							)  ,			(
  							    UpperCamelCase__			
  							)  ,			(
  							    UpperCamelCase__			
  							)  ,						
  							)         =	config_and_inputs
  
  							UpperCamelCase__         =	{
  							    """input_ids""": input_ids,
  							    """token_type_ids""": token_type_ids,
  							    """attention_mask""": input_mask,
  							}
  							return config, inputs_dict
  
  
  
  
  @require_tf
  class   _UpperCAmelCase	(   _lowercase ,       _lowercase ,       _lowercase ,       unittest.TestCase					):
  					__A          =   (
  					    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
  					    if is_tf_available()
  					    else ()
  					)
  
  					__A          =   (TFGPTJForCausalLM,) if is_tf_available() else ()
  					__A          =   (
  					    {
  					        "feature-extraction": TFGPTJModel,
  					        "question-answering": TFGPTJForQuestionAnswering,
  					        "text-classification": TFGPTJForSequenceClassification,
  					        "text-generation": TFGPTJForCausalLM,
  					        "zero-shot": TFGPTJForSequenceClassification,
  					    }
  					    if is_tf_available()
  					    else {}
  					)
  					__A          =   False
  					__A          =   False
  					__A          =   False
  					__A          =   False
  
  
  
  
  					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase       )						->  Dict:
  							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 	snake_case_		(						self       )						->  str:
  							UpperCamelCase__         =	TFGPTJModelTester(self       )
  							UpperCamelCase__         =	ConfigTester(self    ,      config_class=__lowerCAmelCase    ,      n_embd=37       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  Tuple:
  							self.config_tester.run_common_tests()
  
  
  
  
  					def 	snake_case_		(						self       )						->  Any:
  							UpperCamelCase__         =	self.model_tester.prepare_config_and_inputs()
  							self.model_tester.create_and_check_gptj_model(*__lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  Dict:
  							UpperCamelCase__         =	self.model_tester.prepare_config_and_inputs()
  							self.model_tester.create_and_check_gptj_model_past(*__lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  List[str]:
  							UpperCamelCase__         =	self.model_tester.prepare_config_and_inputs()
  							self.model_tester.create_and_check_gptj_model_attention_mask_past(*__lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  Union[str, Any]:
  							UpperCamelCase__         =	self.model_tester.prepare_config_and_inputs()
  							self.model_tester.create_and_check_gptj_model_past_large_inputs(*__lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  Tuple:
  							UpperCamelCase__         =	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 	snake_case_		(						self       )						->  Dict:
  							UpperCamelCase__         =	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 	snake_case_		(						self       )						->  Union[str, Any]:
  							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   _UpperCAmelCase	(   unittest.TestCase					):
  
  
  
  
  					def 	snake_case_		(						self       )						->  Optional[int]:
  							UpperCamelCase__         =	TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""    ,      from_pt=__lowerCAmelCase       )
  							UpperCamelCase__         =	tf.convert_to_tensor([[464, 3_290]]    ,      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
  							UpperCamelCase__         =	[464, 3_290, 318, 257, 582, 338, 1_266, 1_545, 13, 632, 318, 257, 9_112, 15_185, 11, 290, 340, 318, 257, 1_545]
  							# fmt: on
  							UpperCamelCase__         =	model.generate(__lowerCAmelCase    ,      do_sample=__lowerCAmelCase       )
  							self.assertListEqual(output_ids[0].numpy().tolist()    ,      __lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  Union[str, Any]:
  							UpperCamelCase__         =	AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""    ,      revision="""float16"""       )
  							UpperCamelCase__         =	TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""    ,      revision="""float16"""    ,      from_pt=__lowerCAmelCase       )
  
  							UpperCamelCase__         =	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"""       ):
  									UpperCamelCase__         =	model.generate(**__lowerCAmelCase    ,      do_sample=__lowerCAmelCase    ,      seed=[42, 0]       )
  							UpperCamelCase__         =	tokenizer.decode(output_ids[0]    ,      skip_special_tokens=__lowerCAmelCase       )
  
  							UpperCamelCase__         =	"""Today is a nice day and I’m going to go for a walk. I’"""
  							self.assertEqual(__lowerCAmelCase    ,      __lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  str:
  							UpperCamelCase__         =	TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""    ,      revision="""float16"""    ,      from_pt=__lowerCAmelCase       )
  							UpperCamelCase__         =	AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""    ,      revision="""float16"""       )
  
  							UpperCamelCase__         =	"""left"""
  
  							# Define PAD Token = EOS Token = 50256
  							UpperCamelCase__         =	tokenizer.eos_token
  							UpperCamelCase__         =	model.config.eos_token_id
  
  							# use different length sentences to test batching
  							UpperCamelCase__         =	[
  							    """Hello, my dog is a little""",
  							    """Today, I""",
  							]
  							UpperCamelCase__         =	[
  							    """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 	snake_case_		(						self       )						->  List[str]:
  							# Confirms that we get the expected results with left-padded beam search
  							UpperCamelCase__  ,			UpperCamelCase__  ,			UpperCamelCase__  ,			UpperCamelCase__         =	self._get_beam_search_test_objects()
  
  							UpperCamelCase__         =	tokenizer(__lowerCAmelCase    ,      return_tensors="""tf"""    ,      padding=__lowerCAmelCase       )
  							UpperCamelCase__         =	model.generate(**__lowerCAmelCase    ,      do_sample=__lowerCAmelCase    ,      num_beams=2       )
  							UpperCamelCase__         =	tokenizer.batch_decode(__lowerCAmelCase    ,      skip_special_tokens=__lowerCAmelCase       )
  							self.assertListEqual(__lowerCAmelCase    ,      __lowerCAmelCase       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  str:
  							# Confirms that left-padding is working properly
  							UpperCamelCase__  ,			UpperCamelCase__  ,			UpperCamelCase__  ,			UpperCamelCase__         =	self._get_beam_search_test_objects()
  
  							UpperCamelCase__         =	tokenizer(__lowerCAmelCase    ,      return_tensors="""tf"""    ,      padding=__lowerCAmelCase       )
  							UpperCamelCase__         =	tokenizer(sentences[0]    ,      return_tensors="""tf"""       )
  							UpperCamelCase__         =	model.generate(**__lowerCAmelCase    ,      do_sample=__lowerCAmelCase    ,      num_beams=2       )
  							UpperCamelCase__         =	(
  							    shape_list(inputs_non_padded["""input_ids"""]       )[-1]
  							    - tf.reduce_sum(tf.cast(inputs["""attention_mask"""][-1]    ,      tf.intaa       )       ).numpy()
  							)
  							UpperCamelCase__         =	tokenizer(sentences[1]    ,      return_tensors="""tf"""       )
  							UpperCamelCase__         =	model.generate(
  							    **__lowerCAmelCase    ,      do_sample=__lowerCAmelCase    ,      num_beams=2    ,      max_length=model.config.max_length - num_paddings       )
  							UpperCamelCase__         =	tokenizer.decode(output_non_padded[0]    ,      skip_special_tokens=__lowerCAmelCase       )
  							UpperCamelCase__         =	tokenizer.decode(output_padded[0]    ,      skip_special_tokens=__lowerCAmelCase       )
  							self.assertListEqual(__lowerCAmelCase    ,      [non_padded_sentence, padded_sentence]       )
  
  
  
  
  					def 	snake_case_		(						self       )						->  Union[str, Any]:
  							# Confirms that XLA is working properly
  							UpperCamelCase__  ,			UpperCamelCase__  ,			UpperCamelCase__  ,			UpperCamelCase__         =	self._get_beam_search_test_objects()
  
  							UpperCamelCase__         =	tokenizer(__lowerCAmelCase    ,      return_tensors="""tf"""    ,      padding=__lowerCAmelCase       )
  							UpperCamelCase__         =	tf.function(model.generate    ,      jit_compile=__lowerCAmelCase       )
  							UpperCamelCase__         =	xla_generate(**__lowerCAmelCase    ,      do_sample=__lowerCAmelCase    ,      num_beams=2       )
  							UpperCamelCase__         =	tokenizer.batch_decode(__lowerCAmelCase    ,      skip_special_tokens=__lowerCAmelCase       )
  							self.assertListEqual(__lowerCAmelCase    ,      __lowerCAmelCase       )
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  import argparse
  import os
  
  import align
  import numpy as np
  import requests
  import tensorflow as tf
  import torch
  from PIL import Image
  from tokenizer import Tokenizer
  
  from transformers import (
      AlignConfig,
      AlignModel,
      AlignProcessor,
      BertConfig,
      BertTokenizer,
      EfficientNetConfig,
      EfficientNetImageProcessor,
  )
  from transformers.utils import logging
  
  
  logging.set_verbosity_info()
  __lowerCAmelCase     =			logging.get_logger(__name__)
  
  def     a						(      _lowercase   :     Union[str, Any]       )     ->str:
  		UpperCamelCase__         =	tf.image.resize(_lowercase ,						(3_4_6, 3_4_6)       )
  		UpperCamelCase__         =	tf.image.crop_to_bounding_box(_lowercase ,						(3_4_6 - 2_8_9) // 2 ,						(3_4_6 - 2_8_9) // 2 ,						2_8_9 ,						2_8_9       )
  		return image
  
  def     a						(      )     ->Optional[int]:
  		UpperCamelCase__         =	EfficientNetConfig.from_pretrained("""google/efficientnet-b7"""       )
  		UpperCamelCase__         =	2_8_9
  		UpperCamelCase__         =	6_4_0
  		UpperCamelCase__         =	{"""0""": """LABEL_0""", """1""": """LABEL_1"""}
  		UpperCamelCase__         =	{"""LABEL_0""": 0, """LABEL_1""": 1}
  		UpperCamelCase__         =	[]
  
  		UpperCamelCase__         =	BertConfig()
  		UpperCamelCase__         =	AlignConfig.from_text_vision_configs(
  		    text_config=_lowercase ,						vision_config=_lowercase ,						projection_dim=6_4_0       )
  		return config
  
  def     a						(      )     ->Optional[int]:
  		UpperCamelCase__         =	"""http://images.cocodataset.org/val2017/000000039769.jpg"""
  		UpperCamelCase__         =	Image.open(requests.get(_lowercase ,						stream=_lowercase       ).raw       )
  		return im
  
  def     a						(      )     ->Optional[int]:
  		UpperCamelCase__         =	EfficientNetImageProcessor(
  		    do_center_crop=_lowercase ,						rescale_factor=1 / 127.5 ,						rescale_offset=_lowercase ,						do_normalize=_lowercase ,						include_top=_lowercase ,						resample=Image.BILINEAR ,						)
  		UpperCamelCase__         =	BertTokenizer.from_pretrained("""bert-base-uncased"""       )
  		UpperCamelCase__         =	6_4
  		UpperCamelCase__         =	AlignProcessor(image_processor=_lowercase ,						tokenizer=_lowercase       )
  		return processor
  
  def     a						(      _lowercase   :     str       )     ->int:
  		# EfficientNet image encoder
  		UpperCamelCase__         =	[v.split("""_"""       )[0].split("""block"""       )[1] for v in original_param_names if v.startswith("""block"""       )]
  		UpperCamelCase__         =	list(set(_lowercase       )       )
  		UpperCamelCase__         =	sorted(_lowercase       )
  		UpperCamelCase__         =	len(_lowercase       )
  		UpperCamelCase__         =	{b: str(_lowercase       ) for b, i in zip(_lowercase ,						range(_lowercase       )       )}
  
  		UpperCamelCase__         =	[]
  		rename_keys.append(("""stem_conv/kernel:0""", """embeddings.convolution.weight""")       )
  		rename_keys.append(("""stem_bn/gamma:0""", """embeddings.batchnorm.weight""")       )
  		rename_keys.append(("""stem_bn/beta:0""", """embeddings.batchnorm.bias""")       )
  		rename_keys.append(("""stem_bn/moving_mean:0""", """embeddings.batchnorm.running_mean""")       )
  		rename_keys.append(("""stem_bn/moving_variance:0""", """embeddings.batchnorm.running_var""")       )
  
  		for b in block_names:
  				UpperCamelCase__         =	block_name_mapping[b]
  				rename_keys.append((F"""block{b}_expand_conv/kernel:0""", F"""encoder.blocks.{hf_b}.expansion.expand_conv.weight""")       )
  				rename_keys.append((F"""block{b}_expand_bn/gamma:0""", F"""encoder.blocks.{hf_b}.expansion.expand_bn.weight""")       )
  				rename_keys.append((F"""block{b}_expand_bn/beta:0""", F"""encoder.blocks.{hf_b}.expansion.expand_bn.bias""")       )
  				rename_keys.append(
  				    (F"""block{b}_expand_bn/moving_mean:0""", F"""encoder.blocks.{hf_b}.expansion.expand_bn.running_mean""")       )
  				rename_keys.append(
  				    (F"""block{b}_expand_bn/moving_variance:0""", F"""encoder.blocks.{hf_b}.expansion.expand_bn.running_var""")       )
  				rename_keys.append(
  				    (F"""block{b}_dwconv/depthwise_kernel:0""", F"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight""")       )
  				rename_keys.append((F"""block{b}_bn/gamma:0""", F"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight""")       )
  				rename_keys.append((F"""block{b}_bn/beta:0""", F"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias""")       )
  				rename_keys.append(
  				    (F"""block{b}_bn/moving_mean:0""", F"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean""")       )
  				rename_keys.append(
  				    (F"""block{b}_bn/moving_variance:0""", F"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var""")       )
  
  				rename_keys.append((F"""block{b}_se_reduce/kernel:0""", F"""encoder.blocks.{hf_b}.squeeze_excite.reduce.weight""")       )
  				rename_keys.append((F"""block{b}_se_reduce/bias:0""", F"""encoder.blocks.{hf_b}.squeeze_excite.reduce.bias""")       )
  				rename_keys.append((F"""block{b}_se_expand/kernel:0""", F"""encoder.blocks.{hf_b}.squeeze_excite.expand.weight""")       )
  				rename_keys.append((F"""block{b}_se_expand/bias:0""", F"""encoder.blocks.{hf_b}.squeeze_excite.expand.bias""")       )
  				rename_keys.append(
  				    (F"""block{b}_project_conv/kernel:0""", F"""encoder.blocks.{hf_b}.projection.project_conv.weight""")       )
  				rename_keys.append((F"""block{b}_project_bn/gamma:0""", F"""encoder.blocks.{hf_b}.projection.project_bn.weight""")       )
  				rename_keys.append((F"""block{b}_project_bn/beta:0""", F"""encoder.blocks.{hf_b}.projection.project_bn.bias""")       )
  				rename_keys.append(
  				    (F"""block{b}_project_bn/moving_mean:0""", F"""encoder.blocks.{hf_b}.projection.project_bn.running_mean""")       )
  				rename_keys.append(
  				    (F"""block{b}_project_bn/moving_variance:0""", F"""encoder.blocks.{hf_b}.projection.project_bn.running_var""")       )
  
  		UpperCamelCase__         =	{}
  		for item in rename_keys:
  				if item[0] in original_param_names:
  						UpperCamelCase__         =	"""vision_model.""" + item[1]
  
      # BERT text encoder
  		UpperCamelCase__         =	[]
  		UpperCamelCase__         =	"""tf_bert_model/bert"""
  		UpperCamelCase__         =	"""text_model"""
  		for i in range(1_2       ):
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/attention/self/query/kernel:0""",
  				        F"""{new}.encoder.layer.{i}.attention.self.query.weight""",
  				    )       )
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/attention/self/query/bias:0""",
  				        F"""{new}.encoder.layer.{i}.attention.self.query.bias""",
  				    )       )
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/attention/self/key/kernel:0""",
  				        F"""{new}.encoder.layer.{i}.attention.self.key.weight""",
  				    )       )
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/attention/self/key/bias:0""",
  				        F"""{new}.encoder.layer.{i}.attention.self.key.bias""",
  				    )       )
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/attention/self/value/kernel:0""",
  				        F"""{new}.encoder.layer.{i}.attention.self.value.weight""",
  				    )       )
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/attention/self/value/bias:0""",
  				        F"""{new}.encoder.layer.{i}.attention.self.value.bias""",
  				    )       )
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/attention/output/dense/kernel:0""",
  				        F"""{new}.encoder.layer.{i}.attention.output.dense.weight""",
  				    )       )
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/attention/output/dense/bias:0""",
  				        F"""{new}.encoder.layer.{i}.attention.output.dense.bias""",
  				    )       )
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/attention/output/LayerNorm/gamma:0""",
  				        F"""{new}.encoder.layer.{i}.attention.output.LayerNorm.weight""",
  				    )       )
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/attention/output/LayerNorm/beta:0""",
  				        F"""{new}.encoder.layer.{i}.attention.output.LayerNorm.bias""",
  				    )       )
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/intermediate/dense/kernel:0""",
  				        F"""{new}.encoder.layer.{i}.intermediate.dense.weight""",
  				    )       )
  				rename_keys.append(
  				    (
  				        F"""{old}/encoder/layer_._{i}/intermediate/dense/bias:0""",
  				        F"""{new}.encoder.layer.{i}.intermediate.dense.bias""",
  				    )       )
  				rename_keys.append(
  				    (F"""{old}/encoder/layer_._{i}/output/dense/kernel:0""", F"""{new}.encoder.layer.{i}.output.dense.weight""")       )
  				rename_keys.append(
  				    (F"""{old}/encoder/layer_._{i}/output/dense/bias:0""", F"""{new}.encoder.layer.{i}.output.dense.bias""")       )
  				rename_keys.append(
  				    (F"""{old}/encoder/layer_._{i}/output/LayerNorm/gamma:0""", F"""{new}.encoder.layer.{i}.output.LayerNorm.weight""")       )
  				rename_keys.append(
  				    (F"""{old}/encoder/layer_._{i}/output/LayerNorm/beta:0""", F"""{new}.encoder.layer.{i}.output.LayerNorm.bias""")       )
  
  		rename_keys.append((F"""{old}/embeddings/word_embeddings/weight:0""", F"""{new}.embeddings.word_embeddings.weight""")       )
  		rename_keys.append(
  		    (F"""{old}/embeddings/position_embeddings/embeddings:0""", F"""{new}.embeddings.position_embeddings.weight""")       )
  		rename_keys.append(
  		    (F"""{old}/embeddings/token_type_embeddings/embeddings:0""", F"""{new}.embeddings.token_type_embeddings.weight""")       )
  		rename_keys.append((F"""{old}/embeddings/LayerNorm/gamma:0""", F"""{new}.embeddings.LayerNorm.weight""")       )
  		rename_keys.append((F"""{old}/embeddings/LayerNorm/beta:0""", F"""{new}.embeddings.LayerNorm.bias""")       )
  
  		rename_keys.append((F"""{old}/pooler/dense/kernel:0""", F"""{new}.pooler.dense.weight""")       )
  		rename_keys.append((F"""{old}/pooler/dense/bias:0""", F"""{new}.pooler.dense.bias""")       )
  		rename_keys.append(("""dense/kernel:0""", """text_projection.weight""")       )
  		rename_keys.append(("""dense/bias:0""", """text_projection.bias""")       )
  		rename_keys.append(("""dense/bias:0""", """text_projection.bias""")       )
  		rename_keys.append(("""temperature:0""", """temperature""")       )
  
  		for item in rename_keys:
  				if item[0] in original_param_names:
  						UpperCamelCase__         =	item[1]
  		return key_mapping
  
  def     a						(      _lowercase   :     int ,						_lowercase   :     str ,						_lowercase   :     int       )     ->Optional[Any]:
  		list(hf_params.keys()       )
  
  		for key, value in tf_params.items():
  				if key not in key_mapping:
  						continue
  
  				UpperCamelCase__         =	key_mapping[key]
  				if "_conv" in key and "kernel" in key:
  						UpperCamelCase__         =	torch.from_numpy(_lowercase       ).permute(3 ,						2 ,						0 ,						1       )
  				elif "embeddings" in key:
  						UpperCamelCase__         =	torch.from_numpy(_lowercase       )
  				elif "depthwise_kernel" in key:
  						UpperCamelCase__         =	torch.from_numpy(_lowercase       ).permute(2 ,						3 ,						0 ,						1       )
  				elif "kernel" in key:
  						UpperCamelCase__         =	torch.from_numpy(np.transpose(_lowercase       )       )
  				elif "temperature" in key:
  						UpperCamelCase__         =	value
  				elif "bn/gamma" or "bn/beta" in key:
  						UpperCamelCase__         =	torch.from_numpy(np.transpose(_lowercase       )       ).squeeze()
  				else:
  						UpperCamelCase__         =	torch.from_numpy(_lowercase       )
  
  				# Replace HF parameters with original TF model parameters
  				hf_params[hf_key].copy_(_lowercase       )
  
  @torch.no_grad()
  def     a						(      _lowercase   :     List[str] ,						_lowercase   :     Any ,						_lowercase   :     Dict ,						_lowercase   :     Optional[Any]       )     ->str:
  		UpperCamelCase__         =	6_4
  		UpperCamelCase__         =	Tokenizer(_lowercase       )
  		UpperCamelCase__         =	align.Align("""efficientnet-b7""" ,						"""bert-base""" ,						6_4_0 ,						_lowercase ,						tok.get_vocab_size()       )
  		original_model.compile()
  		original_model.load_weights(_lowercase       )
  
  		UpperCamelCase__         =	original_model.trainable_variables
  		UpperCamelCase__         =	original_model.non_trainable_variables
  		UpperCamelCase__         =	{param.name: param.numpy() for param in tf_params}
  		for param in tf_non_train_params:
  				UpperCamelCase__         =	param.numpy()
  		UpperCamelCase__         =	list(tf_params.keys()       )
  
  		# Load HuggingFace model
  		UpperCamelCase__         =	get_align_config()
  		UpperCamelCase__         =	AlignModel(_lowercase       ).eval()
  		UpperCamelCase__         =	hf_model.state_dict()
  
  		# Create src-to-dst parameter name mapping dictionary
  		print("""Converting parameters..."""       )
  		UpperCamelCase__         =	rename_keys(_lowercase       )
  		replace_params(_lowercase ,						_lowercase ,						_lowercase       )
  
  		# Initialize processor
  		UpperCamelCase__         =	get_processor()
  		UpperCamelCase__         =	processor(
  		    images=prepare_img() ,						text="""A picture of a cat""" ,						padding="""max_length""" ,						max_length=6_4 ,						return_tensors="""pt"""       )
  
  		# HF model inference
  		hf_model.eval()
  		with torch.no_grad():
  				UpperCamelCase__         =	hf_model(**_lowercase       )
  
  		UpperCamelCase__         =	outputs.image_embeds.detach().numpy()
  		UpperCamelCase__         =	outputs.text_embeds.detach().numpy()
  
  		# Original model inference
  		UpperCamelCase__         =	False
  		UpperCamelCase__         =	EfficientNetImageProcessor(
  		    do_center_crop=_lowercase ,						do_rescale=_lowercase ,						do_normalize=_lowercase ,						include_top=_lowercase ,						resample=Image.BILINEAR ,						)
  		UpperCamelCase__         =	tf_image_processor(images=prepare_img() ,						return_tensors="""tf""" ,						data_format="""channels_last"""       )["""pixel_values"""]
  		UpperCamelCase__         =	tok(tf.constant(["""A picture of a cat"""]       )       )
  
  		UpperCamelCase__         =	original_model.image_encoder(_lowercase ,						training=_lowercase       )
  		UpperCamelCase__         =	original_model.text_encoder(_lowercase ,						training=_lowercase       )
  
  		UpperCamelCase__         =	tf.nn.la_normalize(_lowercase ,						axis=-1       )
  		UpperCamelCase__         =	tf.nn.la_normalize(_lowercase ,						axis=-1       )
  
  		# Check whether original and HF model outputs match  -> np.allclose
  		if not np.allclose(_lowercase ,						_lowercase ,						atol=1e-3       ):
  				raise ValueError("""The predicted image features are not the same."""       )
  		if not np.allclose(_lowercase ,						_lowercase ,						atol=1e-3       ):
  				raise ValueError("""The predicted text features are not the same."""       )
  		print("""Model outputs match!"""       )
  
  		if save_model:
  				# Create folder to save model
  				if not os.path.isdir(_lowercase       ):
  						os.mkdir(_lowercase       )
  				# Save converted model and image processor
  				hf_model.save_pretrained(_lowercase       )
  				processor.save_pretrained(_lowercase       )
  
  		if push_to_hub:
  				# Push model and image processor to hub
  				print("""Pushing converted ALIGN to the hub..."""       )
  				processor.push_to_hub("""align-base"""       )
  				hf_model.push_to_hub("""align-base"""       )
  
  
  if __name__ == "__main__":
  				__lowerCAmelCase     =			argparse.ArgumentParser()
  				# Required parameters
  				parser.add_argument(
  				    "--checkpoint_path",
  				    default="./weights/model-weights",
  				    type=str,
  				    help="Path to the pretrained TF ALIGN checkpoint.",
  				)
  				parser.add_argument(
  				    "--pytorch_dump_folder_path",
  				    default="hf_model",
  				    type=str,
  				    help="Path to the output PyTorch model directory.",
  				)
  				parser.add_argument("--save_model", action="store_true", help="Save model to local")
  				parser.add_argument("--push_to_hub", action="store_true", help="Push model and image processor to the hub")
  
  				__lowerCAmelCase     =			parser.parse_args()
  				convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)
  
  
  
  
  
  