


from __future__ import annotations

import math
from itertools import permutations

def       _a	(					SCREAMING_SNAKE_CASE						:					int  ):







 """simple docstring"""



 if 1 < number < 4:
  # 2 and 3 are primes
  return True
 elif number < 2 or number % 2 == 0 or number % 3 == 0:
  # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
  return False

 # All primes number are in format of 6k +/- 1
 for i in range(5     ,			int(math.sqrt(SCREAMING_SNAKE_CASE  ) + 1  )     ,			6  ):
  if number % i == 0 or number % (i + 2) == 0:
   return False
 return True

def       _a	(					SCREAMING_SNAKE_CASE						:					int = 7  ):







 """simple docstring"""



 UpperCamelCase__       :     Optional[Any] 							= ''''''.join(str(SCREAMING_SNAKE_CASE  ) for i in range(1     ,			n + 1  )  )
 UpperCamelCase__       :     List[str] 							= [int(''''''.join(SCREAMING_SNAKE_CASE  )  ) for i in permutations(SCREAMING_SNAKE_CASE     ,			SCREAMING_SNAKE_CASE  )]
 UpperCamelCase__       :     Any 							= [num for num in perm_list if is_prime(SCREAMING_SNAKE_CASE  )]
 return max(SCREAMING_SNAKE_CASE  ) if pandigitals else 0


if __name__ == "__main__":
       print(f"{solution() = }")






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    __magic_name__						:


     def __init__(   self				:  int				,						lowerCamelCase__				:  Union[str, Any]      )	->   Optional[int]:





      '''simple docstring'''






      UpperCamelCase__       :     List[str] 							= parent
      UpperCamelCase__       :     Tuple 							= 13
      UpperCamelCase__       :     List[str] 							= 7
      UpperCamelCase__       :     str 							= True
      UpperCamelCase__       :     Tuple 							= True
      UpperCamelCase__       :     str 							= True
      UpperCamelCase__       :     Dict 							= True
      UpperCamelCase__       :     Any 							= True
      UpperCamelCase__       :     Tuple 							= 99
      UpperCamelCase__       :     Tuple 							= 32
      UpperCamelCase__       :     int 							= 4
      UpperCamelCase__       :     Dict 							= 2
      UpperCamelCase__       :     List[str] 							= 4
      UpperCamelCase__       :     str 							= 37
      UpperCamelCase__       :     Optional[Any] 							= '''gelu'''
      UpperCamelCase__       :     Optional[int] 							= 0.1
      UpperCamelCase__       :     List[str] 							= 0.1
      UpperCamelCase__       :     Tuple 							= 512
      UpperCamelCase__       :     Optional[int] 							= 16
      UpperCamelCase__       :     int 							= 2
      UpperCamelCase__       :     Dict 							= 0.02
      UpperCamelCase__       :     Union[str, Any] 							= 3
      UpperCamelCase__       :     Optional[Any] 							= 4
      UpperCamelCase__       :     int 							= None
      UpperCamelCase__       :     Any 							= self.vocab_size - 1
      UpperCamelCase__       :     Optional[int] 							= self.vocab_size - 1
      UpperCamelCase__       :     Tuple 							= self.vocab_size - 1


     def   UpperCAmelCase__						(   self				:  Optional[int]      )	->   int:





      '''simple docstring'''






      UpperCamelCase__       :     List[str] 							= ids_tensor([self.batch_size, self.seq_length]				,						self.vocab_size      )

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

      UpperCamelCase__       :     List[str] 							= None
      if self.use_token_type_ids:
       UpperCamelCase__       :     int 							= ids_tensor([self.batch_size, self.seq_length]				,						self.type_vocab_size      )

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

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

      UpperCamelCase__       :     Dict 							= 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__       :     List[str] 							= 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[Any]				,						lowerCamelCase__				:  Dict				,						lowerCamelCase__				:  List[Any]				,						lowerCamelCase__				:  Optional[Any]				,						lowerCamelCase__				:  Optional[int]				,						lowerCamelCase__				:  Tuple				,						*lowerCamelCase__				:  Optional[Any]      )	->   List[str]:





      '''simple docstring'''






      UpperCamelCase__       :     Union[str, Any] 							= TFGPTJModel(config=lowerCamelCase__      )
      UpperCamelCase__       :     Optional[int] 							= {
          '''input_ids''': input_ids,
          '''attention_mask''': input_mask,
          '''token_type_ids''': token_type_ids,
      }
      UpperCamelCase__       :     Tuple 							= model(lowerCamelCase__      )

      UpperCamelCase__       :     Dict 							= [input_ids, None, input_mask]  # None is the input for 'past'
      UpperCamelCase__       :     List[str] 							= model(lowerCamelCase__      )

      UpperCamelCase__       :     List[Any] 							= model(lowerCamelCase__      )

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


     def   UpperCAmelCase__						(   self				:  List[str]				,						lowerCamelCase__				:  List[Any]				,						lowerCamelCase__				:  Union[str, Any]				,						lowerCamelCase__				:  Any				,						lowerCamelCase__				:  Optional[int]				,						lowerCamelCase__				:  Optional[Any]				,						*lowerCamelCase__				:  Dict      )	->   Union[str, Any]:





      '''simple docstring'''






      UpperCamelCase__       :     str 							= TFGPTJModel(config=lowerCamelCase__      )

      # first forward pass
      UpperCamelCase__       :     Tuple 							= model(lowerCamelCase__				,						token_type_ids=lowerCamelCase__				,						use_cache=lowerCamelCase__      )
      UpperCamelCase__       :     List[Any] 							= model(lowerCamelCase__				,						token_type_ids=lowerCamelCase__      )
      UpperCamelCase__       :     Union[str, Any] 							= 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__       :     Optional[Any] 							= outputs.to_tuple()

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

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

      UpperCamelCase__       :     int 							= model(lowerCamelCase__				,						token_type_ids=lowerCamelCase__      )['''last_hidden_state''']
      UpperCamelCase__       :     Optional[Any] 							= model(lowerCamelCase__				,						token_type_ids=lowerCamelCase__				,						past_key_values=lowerCamelCase__      )[
          '''last_hidden_state'''
      ]

      # select random slice
      UpperCamelCase__       :     int 							= int(ids_tensor((1,)				,						shape_list(lowerCamelCase__      )[-1]      )      )
      UpperCamelCase__       :     Any 							= output_from_no_past[:, -1, random_slice_idx]
      UpperCamelCase__       :     Optional[Any] 							= 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				:  Union[str, Any]				,						lowerCamelCase__				:  int				,						lowerCamelCase__				:  List[str]				,						lowerCamelCase__				:  Optional[int]				,						lowerCamelCase__				:  Any				,						lowerCamelCase__				:  List[Any]				,						*lowerCamelCase__				:  Tuple      )	->   Optional[Any]:





      '''simple docstring'''






      UpperCamelCase__       :     Optional[Any] 							= TFGPTJModel(config=lowerCamelCase__      )

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

      # first forward pass
      UpperCamelCase__							,					UpperCamelCase__       :     str 							= model(lowerCamelCase__				,						attention_mask=lowerCamelCase__      ).to_tuple()

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

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

      # append to next input_ids and attn_mask
      UpperCamelCase__       :     Dict 							= tf.concat([input_ids, next_tokens]				,						axis=-1      )
      UpperCamelCase__       :     Optional[Any] 							= tf.concat([attn_mask, tf.ones((shape_list(lowerCamelCase__      )[0], 1)				,						dtype=tf.intaa      )]				,						axis=1      )

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

      # select random slice
      UpperCamelCase__       :     Union[str, Any] 							= int(ids_tensor((1,)				,						shape_list(lowerCamelCase__      )[-1]      )      )
      UpperCamelCase__       :     List[str] 							= output_from_no_past[:, -1, random_slice_idx]
      UpperCamelCase__       :     int 							= 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   UpperCAmelCase__						(   self				:  Tuple				,						lowerCamelCase__				:  str				,						lowerCamelCase__				:  List[Any]				,						lowerCamelCase__				:  int				,						lowerCamelCase__				:  Optional[Any]				,						lowerCamelCase__				:  str				,						*lowerCamelCase__				:  Optional[Any]      )	->   Union[str, Any]:





      '''simple docstring'''






      UpperCamelCase__       :     Optional[int] 							= TFGPTJModel(config=lowerCamelCase__      )

      UpperCamelCase__       :     List[str] 							= input_ids[:1, :]
      UpperCamelCase__       :     List[Any] 							= input_mask[:1, :]
      UpperCamelCase__       :     Optional[Any] 							= token_type_ids[:1, :]
      UpperCamelCase__       :     List[Any] 							= 1

      # first forward pass
      UpperCamelCase__       :     Tuple 							= model(lowerCamelCase__				,						attention_mask=lowerCamelCase__				,						token_type_ids=lowerCamelCase__				,						use_cache=lowerCamelCase__      )

      UpperCamelCase__							,					UpperCamelCase__       :     Union[str, Any] 							= outputs.to_tuple()

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

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

      UpperCamelCase__       :     Dict 							= model(
          lowerCamelCase__				,						token_type_ids=lowerCamelCase__				,						attention_mask=lowerCamelCase__      )['''last_hidden_state''']
      UpperCamelCase__       :     str 							= 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__       :     Dict 							= int(ids_tensor((1,)				,						shape_list(lowerCamelCase__      )[-1]      )      )
      UpperCamelCase__       :     Tuple 							= output_from_no_past[:, -3:, random_slice_idx]
      UpperCamelCase__       :     Optional[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				:  Union[str, Any]				,						lowerCamelCase__				:  Optional[int]				,						lowerCamelCase__				:  Dict				,						lowerCamelCase__				:  Dict				,						lowerCamelCase__				:  Union[str, Any]				,						lowerCamelCase__				:  Optional[Any]				,						*lowerCamelCase__				:  Tuple      )	->   List[Any]:





      '''simple docstring'''






      UpperCamelCase__       :     List[str] 							= TFGPTJForCausalLM(config=lowerCamelCase__      )
      UpperCamelCase__       :     int 							= {
          '''input_ids''': input_ids,
          '''attention_mask''': input_mask,
          '''token_type_ids''': token_type_ids,
      }
      UpperCamelCase__       :     Union[str, Any] 							= model(lowerCamelCase__      )
      self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length, self.vocab_size)      )


     def   UpperCAmelCase__						(   self				:  Optional[Any]      )	->   List[Any]:





      '''simple docstring'''






      UpperCamelCase__       :     str 							= self.prepare_config_and_inputs()

      (
          (
          UpperCamelCase__
      )							,					(
          UpperCamelCase__
      )							,					(
          UpperCamelCase__
      )							,					(
          UpperCamelCase__
      )							,					(
          UpperCamelCase__
      )							,					(
          UpperCamelCase__
      )							,					(
          UpperCamelCase__
      )							,					(
          UpperCamelCase__
      )							,					(
          UpperCamelCase__
      )							,					
      )       :     Union[str, Any] 							= config_and_inputs

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



@require_tf
class    __magic_name__						(  __lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					unittest.TestCase):
     A:		Dict								=   (
         (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
         if is_tf_available()
         else ()
     )

     A:		Tuple								=   (TFGPTJForCausalLM,) if is_tf_available() else ()
     A:		Optional[Any]								=   (
         {
             "feature-extraction": TFGPTJModel,
             "question-answering": TFGPTJForQuestionAnswering,
             "text-classification": TFGPTJForSequenceClassification,
             "text-generation": TFGPTJForCausalLM,
             "zero-shot": TFGPTJForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     A:		List[Any]								=   False
     A:		Optional[Any]								=   False
     A:		str								=   False
     A:		Optional[Any]								=   False


     def   UpperCAmelCase__						(   self				:  int				,						lowerCamelCase__				:  Any				,						lowerCamelCase__				:  Any				,						lowerCamelCase__				:  List[Any]				,						lowerCamelCase__				:  Union[str, Any]				,						lowerCamelCase__				:  Any      )	->   Optional[Any]:





      '''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				:  List[Any]      )	->   str:





      '''simple docstring'''






      UpperCamelCase__       :     Optional[int] 							= TFGPTJModelTester(self      )
      UpperCamelCase__       :     Tuple 							= ConfigTester(self				,						config_class=lowerCamelCase__				,						n_embd=37      )


     def   UpperCAmelCase__						(   self				:  str      )	->   str:





      '''simple docstring'''






      self.config_tester.run_common_tests()


     def   UpperCAmelCase__						(   self				:  Tuple      )	->   int:





      '''simple docstring'''






      UpperCamelCase__       :     Optional[int] 							= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_gptj_model(*lowerCamelCase__      )


     def   UpperCAmelCase__						(   self				:  Optional[Any]      )	->   List[Any]:





      '''simple docstring'''






      UpperCamelCase__       :     str 							= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_gptj_model_past(*lowerCamelCase__      )


     def   UpperCAmelCase__						(   self				:  Optional[int]      )	->   Optional[int]:





      '''simple docstring'''






      UpperCamelCase__       :     Optional[int] 							= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_gptj_model_attention_mask_past(*lowerCamelCase__      )


     def   UpperCAmelCase__						(   self				:  Tuple      )	->   Optional[int]:





      '''simple docstring'''






      UpperCamelCase__       :     str 							= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_gptj_model_past_large_inputs(*lowerCamelCase__      )


     def   UpperCAmelCase__						(   self				:  int      )	->   Dict:





      '''simple docstring'''






      UpperCamelCase__       :     int 							= 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				:  Tuple      )	->   Union[str, Any]:





      '''simple docstring'''






      UpperCamelCase__       :     Dict 							= 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				:  Optional[int]      )	->   str:





      '''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    __magic_name__						(  unittest.TestCase):


     def   UpperCAmelCase__						(   self				:  int      )	->   Union[str, Any]:





      '''simple docstring'''






      UpperCamelCase__       :     str 							= TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''				,						from_pt=lowerCamelCase__      )
      UpperCamelCase__       :     Dict 							= tf.convert_to_tensor([[464, 3290]]				,						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__       :     Dict 							= [464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 15185, 11, 290, 340, 318, 257, 1545]
      # fmt: on
      UpperCamelCase__       :     Optional[int] 							= model.generate(lowerCamelCase__				,						do_sample=lowerCamelCase__      )
      self.assertListEqual(output_ids[0].numpy().tolist()				,						lowerCamelCase__      )


     def   UpperCAmelCase__						(   self				:  str      )	->   str:





      '''simple docstring'''






      UpperCamelCase__       :     Union[str, Any] 							= AutoTokenizer.from_pretrained('''EleutherAI/gpt-j-6B'''				,						revision='''float16'''      )
      UpperCamelCase__       :     List[Any] 							= TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''				,						revision='''float16'''				,						from_pt=lowerCamelCase__      )

      UpperCamelCase__       :     Optional[Any] 							= 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__       :     int 							= model.generate(**lowerCamelCase__				,						do_sample=lowerCamelCase__				,						seed=[42, 0]      )
      UpperCamelCase__       :     Any 							= tokenizer.decode(output_ids[0]				,						skip_special_tokens=lowerCamelCase__      )

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


     def   UpperCAmelCase__						(   self				:  List[Any]      )	->   Dict:





      '''simple docstring'''






      UpperCamelCase__       :     Any 							= TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''				,						revision='''float16'''				,						from_pt=lowerCamelCase__      )
      UpperCamelCase__       :     Union[str, Any] 							= AutoTokenizer.from_pretrained('''EleutherAI/gpt-j-6B'''				,						revision='''float16'''      )

      UpperCamelCase__       :     str 							= '''left'''

      # Define PAD Token = EOS Token = 50256
      UpperCamelCase__       :     Optional[Any] 							= tokenizer.eos_token
      UpperCamelCase__       :     List[str] 							= model.config.eos_token_id

      # use different length sentences to test batching
      UpperCamelCase__       :     Tuple 							= [
          '''Hello, my dog is a little''',
          '''Today, I''',
      ]
      UpperCamelCase__       :     Optional[Any] 							= [
          '''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				:  str      )	->   Optional[int]:





      '''simple docstring'''






      UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__       :     Dict 							= self._get_beam_search_test_objects()

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


     def   UpperCAmelCase__						(   self				:  List[Any]      )	->   str:





      '''simple docstring'''






      UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__       :     Tuple 							= self._get_beam_search_test_objects()

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


     def   UpperCAmelCase__						(   self				:  Optional[int]      )	->   str:





      '''simple docstring'''






      UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__       :     int 							= self._get_beam_search_test_objects()

      UpperCamelCase__       :     Union[str, Any] 							= tokenizer(lowerCamelCase__				,						return_tensors='''tf'''				,						padding=lowerCamelCase__      )
      UpperCamelCase__       :     List[Any] 							= tf.function(model.generate				,						jit_compile=lowerCamelCase__      )
      UpperCamelCase__       :     str 							= xla_generate(**lowerCamelCase__				,						do_sample=lowerCamelCase__				,						num_beams=2      )
      UpperCamelCase__       :     str 							= tokenizer.batch_decode(lowerCamelCase__				,						skip_special_tokens=lowerCamelCase__      )
      self.assertListEqual(lowerCamelCase__				,						lowerCamelCase__      )


