





"""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__				)




from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


_SCREAMING_SNAKE_CASE:    List[Any]     						=							logging.get_logger(__name__)

class 						_SCREAMING_SNAKE_CASE				(  _lowercase				):






		"""simple docstring"""







		A_    			=   ['''pixel_values''']



		def __init__(      self     :					Tuple	,       __SCREAMING_SNAKE_CASE     :					bool = True	,       __SCREAMING_SNAKE_CASE     :					Optional[Dict[str, int]] = None	,       __SCREAMING_SNAKE_CASE     :					PILImageResampling = PILImageResampling.BILINEAR	,       __SCREAMING_SNAKE_CASE     :					bool = True	,       __SCREAMING_SNAKE_CASE     :					Union[int, float] = 1 / 255	,       __SCREAMING_SNAKE_CASE     :					bool = True	,       __SCREAMING_SNAKE_CASE     :					Optional[Union[float, List[float]]] = None	,       __SCREAMING_SNAKE_CASE     :					Optional[Union[float, List[float]]] = None	,       **__SCREAMING_SNAKE_CASE     :					str	,       ):






					"""simple docstring"""
					super().__init__(**__SCREAMING_SNAKE_CASE )
					_snake_case														=   size if size is not None else {"height": 224, "width": 224}
					_snake_case														=   get_size_dict(__SCREAMING_SNAKE_CASE )
					_snake_case														=   do_resize
					_snake_case														=   do_rescale
					_snake_case														=   do_normalize
					_snake_case														=   size
					_snake_case														=   resample
					_snake_case														=   rescale_factor
					_snake_case														=   image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
					_snake_case														=   image_std if image_std is not None else IMAGENET_STANDARD_STD



		def 					_a					(      self     :					int	,       __SCREAMING_SNAKE_CASE     :					np.ndarray	,       __SCREAMING_SNAKE_CASE     :					Dict[str, int]	,       __SCREAMING_SNAKE_CASE     :					PILImageResampling = PILImageResampling.BILINEAR	,       __SCREAMING_SNAKE_CASE     :					Optional[Union[str, ChannelDimension]] = None	,       **__SCREAMING_SNAKE_CASE     :					Tuple	,       ):






					"""simple docstring"""
					_snake_case														=   get_size_dict(__SCREAMING_SNAKE_CASE )
					if "height" not in size or "width" not in size:
								raise ValueError(F'The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}' )
					return resize(
					    __SCREAMING_SNAKE_CASE	,       size=(size["height"], size["width"])	,       resample=__SCREAMING_SNAKE_CASE	,       data_format=__SCREAMING_SNAKE_CASE	,       **__SCREAMING_SNAKE_CASE )



		def 					_a					(      self     :					Tuple	,       __SCREAMING_SNAKE_CASE     :					np.ndarray	,       __SCREAMING_SNAKE_CASE     :					float	,       __SCREAMING_SNAKE_CASE     :					Optional[Union[str, ChannelDimension]] = None	,       **__SCREAMING_SNAKE_CASE     :					List[str] ):






					"""simple docstring"""
					return rescale(__SCREAMING_SNAKE_CASE	,       scale=__SCREAMING_SNAKE_CASE	,       data_format=__SCREAMING_SNAKE_CASE	,       **__SCREAMING_SNAKE_CASE )



		def 					_a					(      self     :					Optional[Any]	,       __SCREAMING_SNAKE_CASE     :					np.ndarray	,       __SCREAMING_SNAKE_CASE     :					Union[float, List[float]]	,       __SCREAMING_SNAKE_CASE     :					Union[float, List[float]]	,       __SCREAMING_SNAKE_CASE     :					Optional[Union[str, ChannelDimension]] = None	,       **__SCREAMING_SNAKE_CASE     :					Optional[Any]	,       ):






					"""simple docstring"""
					return normalize(__SCREAMING_SNAKE_CASE	,       mean=__SCREAMING_SNAKE_CASE	,       std=__SCREAMING_SNAKE_CASE	,       data_format=__SCREAMING_SNAKE_CASE	,       **__SCREAMING_SNAKE_CASE )







		def 					_a					(      self     :					int	,       __SCREAMING_SNAKE_CASE     :					ImageInput	,       __SCREAMING_SNAKE_CASE     :					Optional[bool] = None	,       __SCREAMING_SNAKE_CASE     :					Dict[str, int] = None	,       __SCREAMING_SNAKE_CASE     :					PILImageResampling = None	,       __SCREAMING_SNAKE_CASE     :					Optional[bool] = None	,       __SCREAMING_SNAKE_CASE     :					Optional[float] = None	,       __SCREAMING_SNAKE_CASE     :					Optional[bool] = None	,       __SCREAMING_SNAKE_CASE     :					Optional[Union[float, List[float]]] = None	,       __SCREAMING_SNAKE_CASE     :					Optional[Union[float, List[float]]] = None	,       __SCREAMING_SNAKE_CASE     :					Optional[Union[str, TensorType]] = None	,       __SCREAMING_SNAKE_CASE     :					Union[str, ChannelDimension] = ChannelDimension.FIRST	,       **__SCREAMING_SNAKE_CASE     :					List[Any]	,       ):






					"""simple docstring"""
					_snake_case														=   do_resize if do_resize is not None else self.do_resize
					_snake_case														=   do_rescale if do_rescale is not None else self.do_rescale
					_snake_case														=   do_normalize if do_normalize is not None else self.do_normalize
					_snake_case														=   resample if resample is not None else self.resample
					_snake_case														=   rescale_factor if rescale_factor is not None else self.rescale_factor
					_snake_case														=   image_mean if image_mean is not None else self.image_mean
					_snake_case														=   image_std if image_std is not None else self.image_std

					_snake_case														=   size if size is not None else self.size
					_snake_case														=   get_size_dict(__SCREAMING_SNAKE_CASE )

					_snake_case														=   make_list_of_images(__SCREAMING_SNAKE_CASE )

					if not valid_images(__SCREAMING_SNAKE_CASE ):
								raise ValueError(
								    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
								    "torch.Tensor, tf.Tensor or jax.ndarray." )

					if do_resize and size is None:
								raise ValueError("Size must be specified if do_resize is True." )

					if do_rescale and rescale_factor is None:
								raise ValueError("Rescale factor must be specified if do_rescale is True." )

					# All transformations expect numpy arrays.
					_snake_case														=   [to_numpy_array(__SCREAMING_SNAKE_CASE ) for image in images]

					if do_resize:
								_snake_case														=   [self.resize(image=__SCREAMING_SNAKE_CASE	,       size=__SCREAMING_SNAKE_CASE	,       resample=__SCREAMING_SNAKE_CASE ) for image in images]

					if do_rescale:
								_snake_case														=   [self.rescale(image=__SCREAMING_SNAKE_CASE	,       scale=__SCREAMING_SNAKE_CASE ) for image in images]

					if do_normalize:
								_snake_case														=   [self.normalize(image=__SCREAMING_SNAKE_CASE	,       mean=__SCREAMING_SNAKE_CASE	,       std=__SCREAMING_SNAKE_CASE ) for image in images]

					_snake_case														=   [to_channel_dimension_format(__SCREAMING_SNAKE_CASE	,       __SCREAMING_SNAKE_CASE ) for image in images]

					_snake_case														=   {"pixel_values": images}
					return BatchFeature(data=__SCREAMING_SNAKE_CASE	,       tensor_type=__SCREAMING_SNAKE_CASE )

