





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       )












__lowerCAmelCase     =			{
    "0": 1,
    "1": 1,
    "2": 2,
    "3": 6,
    "4": 24,
    "5": 120,
    "6": 720,
    "7": 5_040,
    "8": 40_320,
    "9": 362_880,
}

__lowerCAmelCase     =			{145: 145}

__lowerCAmelCase     =			{
    145: 0,
    169: 3,
    36_301: 3,
    1_454: 3,
    871: 2,
    45_361: 2,
    872: 2,
}

def     a						(      _lowercase   :     int       )     ->int:
		if n in CACHE_SUM_DIGIT_FACTORIALS:
				return CACHE_SUM_DIGIT_FACTORIALS[n]
		UpperCamelCase__         =	sum(DIGIT_FACTORIALS[let] for let in str(_lowercase       )       )
		UpperCamelCase__         =	ret
		return ret

def     a						(      _lowercase   :     int ,						_lowercase   :     set | None = None       )     ->int:
		UpperCamelCase__         =	previous or set()
		if n in CHAIN_LENGTH_CACHE:
				return CHAIN_LENGTH_CACHE[n]
		UpperCamelCase__         =	sum_digit_factorials(_lowercase       )
		if next_number in previous:
				UpperCamelCase__         =	0
				return 0
		else:
				previous.add(_lowercase       )
				UpperCamelCase__         =	1 + chain_length(_lowercase ,						_lowercase       )
				UpperCamelCase__         =	ret
				return ret

def     a						(      _lowercase   :     int = 6_0 ,						_lowercase   :     int = 1_0_0_0_0_0_0       )     ->int:
		return sum(1 for i in range(1 ,						_lowercase       ) if chain_length(_lowercase       ) == num_terms       )


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





