





"""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 					_a		:






				def __init__(					self						:       Union[str, Any],					lowerCAmelCase__						:       Tuple				)      -> List[Any]:




										'''simple docstring'''



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






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




										'''simple docstring'''



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

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

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

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

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

										_UpperCamelCase     :   List[Any]      =       GPTJConfig(
										    vocab_size=self.vocab_size,					n_embd=self.hidden_size,					n_layer=self.num_hidden_layers,					n_head=self.num_attention_heads,					intermediate_size=self.intermediate_size,					hidden_act=self.hidden_act,					hidden_dropout_prob=self.hidden_dropout_prob,					attention_probs_dropout_prob=self.attention_probs_dropout_prob,					n_positions=self.max_position_embeddings,					type_vocab_size=self.type_vocab_size,					initializer_range=self.initializer_range,					bos_token_id=self.bos_token_id,					eos_token_id=self.eos_token_id,					pad_token_id=self.pad_token_id,					rotary_dim=self.rotary_dim,					return_dict=lowerCAmelCase__,					)

										_UpperCamelCase     :   Dict      =       ids_tensor([self.num_hidden_layers, self.num_attention_heads],					2				)

										return (
										    config,
										    input_ids,
										    input_mask,
										    head_mask,
										    token_type_ids,
										    mc_token_ids,
										    sequence_labels,
										    token_labels,
										    choice_labels,
										)






				def 							snake_case      (					self						:       Optional[int],					lowerCAmelCase__						:       Tuple,					lowerCAmelCase__						:       int,					lowerCAmelCase__						:       Dict,					lowerCAmelCase__						:       Union[str, Any],					lowerCAmelCase__						:       Union[str, Any],					*lowerCAmelCase__						:       Union[str, Any]				)      -> int:




										'''simple docstring'''



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

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

										_UpperCamelCase     :   str      =       model(lowerCAmelCase__				)

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






				def 							snake_case      (					self						:       Dict,					lowerCAmelCase__						:       List[Any],					lowerCAmelCase__						:       Tuple,					lowerCAmelCase__						:       Union[str, Any],					lowerCAmelCase__						:       List[str],					lowerCAmelCase__						:       List[str],					*lowerCAmelCase__						:       Union[str, Any]				)      -> Dict:




										'''simple docstring'''



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

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

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

										_UpperCamelCase  ,				_UpperCamelCase     :   str      =       outputs.to_tuple()

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

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

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

										# select random slice
										_UpperCamelCase     :   Dict      =       int(ids_tensor((1,),					shape_list(lowerCAmelCase__				)[-1]				)				)
										_UpperCamelCase     :   Optional[int]      =       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-6				)






				def 							snake_case      (					self						:       Optional[Any],					lowerCAmelCase__						:       int,					lowerCAmelCase__						:       Dict,					lowerCAmelCase__						:       Tuple,					lowerCAmelCase__						:       Tuple,					lowerCAmelCase__						:       List[str],					*lowerCAmelCase__						:       Tuple				)      -> Tuple:




										'''simple docstring'''



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

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

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

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

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

										# append to next input_ids and attn_mask
										_UpperCamelCase     :   str      =       tf.concat([input_ids, next_tokens],					axis=-1				)
										_UpperCamelCase     :   int      =       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     :   List[str]      =       int(ids_tensor((1,),					shape_list(lowerCAmelCase__				)[-1]				)				)
										_UpperCamelCase     :   Tuple      =       output_from_no_past[:, -1, random_slice_idx]
										_UpperCamelCase     :   Dict      =       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						:       str,					lowerCAmelCase__						:       List[Any],					lowerCAmelCase__						:       List[str],					lowerCAmelCase__						:       Optional[int],					lowerCAmelCase__						:       Dict,					lowerCAmelCase__						:       Any,					*lowerCAmelCase__						:       Dict				)      -> str:




										'''simple docstring'''



										_UpperCamelCase     :   Tuple      =       TFGPTJModel(config=lowerCAmelCase__				)

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

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

										_UpperCamelCase  ,				_UpperCamelCase     :   Tuple      =       outputs.to_tuple()

										# create hypothetical next token and extent to next_input_ids
										_UpperCamelCase     :   Tuple      =       ids_tensor((self.batch_size, 3),					config.vocab_size				)
										_UpperCamelCase     :   int      =       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     :   str      =       tf.concat([input_ids, next_tokens],					axis=-1				)
										_UpperCamelCase     :   int      =       tf.concat([input_mask, next_attn_mask],					axis=-1				)
										_UpperCamelCase     :   Tuple      =       tf.concat([token_type_ids, next_token_types],					axis=-1				)

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

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

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






				def 							snake_case      (					self						:       str,					lowerCAmelCase__						:       Any,					lowerCAmelCase__						:       Optional[Any],					lowerCAmelCase__						:       int,					lowerCAmelCase__						:       Any,					lowerCAmelCase__						:       Union[str, Any],					*lowerCAmelCase__						:       Dict				)      -> Union[str, Any]:




										'''simple docstring'''



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






				def 							snake_case      (					self						:       List[str]				)      -> Union[str, Any]:




										'''simple docstring'''



										_UpperCamelCase     :   int      =       self.prepare_config_and_inputs()

										(
										    (
										    _UpperCamelCase
										)  ,				(
										    _UpperCamelCase
										)  ,				(
										    _UpperCamelCase
										)  ,				(
										    _UpperCamelCase
										)  ,				(
										    _UpperCamelCase
										)  ,				(
										    _UpperCamelCase
										)  ,				(
										    _UpperCamelCase
										)  ,				(
										    _UpperCamelCase
										)  ,				(
										    _UpperCamelCase
										)  ,				
										)     :   List[str]      =       config_and_inputs

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




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

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






				def 							snake_case      (					self						:       Any,					lowerCAmelCase__						:       int,					lowerCAmelCase__						:       str,					lowerCAmelCase__						:       Optional[int],					lowerCAmelCase__						:       List[str],					lowerCAmelCase__						:       List[str]				)      -> int:




										'''simple docstring'''



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

										return False






				def 							snake_case      (					self						:       Tuple				)      -> Any:




										'''simple docstring'''



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






				def 							snake_case      (					self						:       Optional[Any]				)      -> str:




										'''simple docstring'''



										self.config_tester.run_common_tests()






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




										'''simple docstring'''



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






				def 							snake_case      (					self						:       Optional[Any]				)      -> Tuple:




										'''simple docstring'''



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






				def 							snake_case      (					self						:       Optional[int]				)      -> Dict:




										'''simple docstring'''



										_UpperCamelCase     :   str      =       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]				)      -> Any:




										'''simple docstring'''



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






				def 							snake_case      (					self						:       Dict				)      -> List[str]:




										'''simple docstring'''



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






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




										'''simple docstring'''



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






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




										'''simple docstring'''



										super().test_resize_token_embeddings()




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






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




										'''simple docstring'''



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






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




										'''simple docstring'''



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

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






				def 							snake_case      (					self						:       int				)      -> List[Any]:




										'''simple docstring'''



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

										_UpperCamelCase     :   Tuple      =       '''left'''

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

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






				def 							snake_case      (					self						:       Any				)      -> Optional[Any]:




										'''simple docstring'''



										_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase     :   Optional[Any]      =       self._get_beam_search_test_objects()

										_UpperCamelCase     :   List[str]      =       tokenizer(lowerCAmelCase__,					return_tensors='''tf''',					padding=lowerCAmelCase__				)
										_UpperCamelCase     :   Optional[Any]      =       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 							snake_case      (					self						:       Union[str, Any]				)      -> List[Any]:




										'''simple docstring'''



										_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase     :   Any      =       self._get_beam_search_test_objects()

										_UpperCamelCase     :   Optional[Any]      =       tokenizer(lowerCAmelCase__,					return_tensors='''tf''',					padding=lowerCAmelCase__				)
										_UpperCamelCase     :   Dict      =       tokenizer(sentences[0],					return_tensors='''tf'''				)
										_UpperCamelCase     :   Any      =       model.generate(**lowerCAmelCase__,					do_sample=lowerCAmelCase__,					num_beams=2				)
										_UpperCamelCase     :   Union[str, Any]      =       (
										    shape_list(inputs_non_padded['''input_ids''']				)[-1]
										    - tf.reduce_sum(tf.cast(inputs['''attention_mask'''][-1],					tf.intaa				)				).numpy()
										)
										_UpperCamelCase     :   Any      =       tokenizer(sentences[1],					return_tensors='''tf'''				)
										_UpperCamelCase     :   Tuple      =       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[str]      =       tokenizer.decode(output_padded[0],					skip_special_tokens=lowerCAmelCase__				)
										self.assertListEqual(lowerCAmelCase__,					[non_padded_sentence, padded_sentence]				)






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




										'''simple docstring'''



										_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase  ,				_UpperCamelCase     :   str      =       self._get_beam_search_test_objects()

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










'''simple docstring'''






from .hash_table import HashTable
from .number_theory.prime_numbers import is_prime, next_prime



class  lowerCamelCase__				(						snake_case_						):


	"""simple docstring"""





	def __init__(					self   ,     *UpperCAmelCase__   ,     **UpperCAmelCase__			)       ->				Optional[int]:
					super().__init__(*UpperCAmelCase__   ,     **UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__			)       ->				Optional[int]:
					_A :				Union[str, Any]				=  (
					    next_prime(value % self.size_table			)
					    if not is_prime(value % self.size_table			)
					    else value % self.size_table
					)  # gt = bigger than
					return next_prime_gt - (data % next_prime_gt)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__			)       ->				Union[str, Any]:
					return (increment * self.__hash_function_a(UpperCAmelCase__   ,     UpperCAmelCase__			)) % self.size_table




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__=None			)       ->				List[Any]:
					_A :				Optional[int]				=  1
					_A :				Optional[int]				=  self.hash_function(UpperCAmelCase__			)

					while self.values[new_key] is not None and self.values[new_key] != key:
									_A :				int				=  (
									    self.__hash_double_function(UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__			)
									    if self.balanced_factor() >= self.lim_charge
									    else None
									)
									if new_key is None:
													break
									else:
													i += 1

					return new_key
