
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				:






     """simple docstring"""







     def __init__(							self    :    List[str]      ,		_A    :    Dict				):







       """simple docstring"""




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



     def 				UpperCAmelCase__    (							self    :    str				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Dict			=						ids_tensor([self.batch_size, self.seq_length]      ,		self.vocab_size				)

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

       __SCREAMING_SNAKE_CASE     :					str			=						None
       if self.use_token_type_ids:
         __SCREAMING_SNAKE_CASE     :					Tuple			=						ids_tensor([self.batch_size, self.seq_length]      ,		self.type_vocab_size				)

       __SCREAMING_SNAKE_CASE     :					Any			=						None
       if self.use_mc_token_ids:
         __SCREAMING_SNAKE_CASE     :					int			=						ids_tensor([self.batch_size, self.num_choices]      ,		self.seq_length				)

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

       __SCREAMING_SNAKE_CASE     :					str			=						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=snake_case_      ,		)

       __SCREAMING_SNAKE_CASE     :					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 				UpperCAmelCase__    (							self    :    Any      ,		_A    :    List[Any]      ,		_A    :    Dict      ,		_A    :    str      ,		_A    :    Union[str, Any]      ,		_A    :    Optional[Any]      ,		*_A    :    str				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					int			=						TFGPTJModel(config=snake_case_				)
       __SCREAMING_SNAKE_CASE     :					Any			=						{
           """input_ids""": input_ids,
           """attention_mask""": input_mask,
           """token_type_ids""": token_type_ids,
       }
       __SCREAMING_SNAKE_CASE     :					int			=						model(snake_case_				)

       __SCREAMING_SNAKE_CASE     :					Tuple			=						[input_ids, None, input_mask]  # None is the input for 'past'
       __SCREAMING_SNAKE_CASE     :					Tuple			=						model(snake_case_				)

       __SCREAMING_SNAKE_CASE     :					List[Any]			=						model(snake_case_				)

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



     def 				UpperCAmelCase__    (							self    :    Optional[int]      ,		_A    :    Tuple      ,		_A    :    Union[str, Any]      ,		_A    :    Dict      ,		_A    :    str      ,		_A    :    List[str]      ,		*_A    :    Dict				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Any			=						TFGPTJModel(config=snake_case_				)

       # first forward pass
       __SCREAMING_SNAKE_CASE     :					int			=						model(snake_case_      ,		token_type_ids=snake_case_      ,		use_cache=snake_case_				)
       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						model(snake_case_      ,		token_type_ids=snake_case_				)
       __SCREAMING_SNAKE_CASE     :					List[str]			=						model(snake_case_      ,		token_type_ids=snake_case_      ,		use_cache=snake_case_				)

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

       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						outputs.to_tuple()

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

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

       __SCREAMING_SNAKE_CASE     :					int			=						model(snake_case_      ,		token_type_ids=snake_case_				)["""last_hidden_state"""]
       __SCREAMING_SNAKE_CASE     :					List[Any]			=						model(snake_case_      ,		token_type_ids=snake_case_      ,		past_key_values=snake_case_				)[
           """last_hidden_state"""
       ]

       # select random slice
       __SCREAMING_SNAKE_CASE     :					Any			=						int(ids_tensor((1,)      ,		shape_list(snake_case_				)[-1]				)				)
       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						output_from_no_past[:, -1, random_slice_idx]
       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						output_from_past[:, 0, random_slice_idx]

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



     def 				UpperCAmelCase__    (							self    :    Union[str, Any]      ,		_A    :    Union[str, Any]      ,		_A    :    str      ,		_A    :    str      ,		_A    :    Optional[int]      ,		_A    :    Tuple      ,		*_A    :    str				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					List[str]			=						TFGPTJModel(config=snake_case_				)

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

       # first forward pass
       __SCREAMING_SNAKE_CASE     :					Tuple			=						model(snake_case_      ,		attention_mask=snake_case_				).to_tuple()

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

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

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

       # get two different outputs
       __SCREAMING_SNAKE_CASE     :					Dict			=						model(snake_case_      ,		attention_mask=snake_case_				)["""last_hidden_state"""]
       __SCREAMING_SNAKE_CASE     :					List[Any]			=						model(snake_case_      ,		past_key_values=snake_case_      ,		attention_mask=snake_case_				)[
           """last_hidden_state"""
       ]

       # select random slice
       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						int(ids_tensor((1,)      ,		shape_list(snake_case_				)[-1]				)				)
       __SCREAMING_SNAKE_CASE     :					List[str]			=						output_from_no_past[:, -1, random_slice_idx]
       __SCREAMING_SNAKE_CASE     :					str			=						output_from_past[:, 0, random_slice_idx]

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



     def 				UpperCAmelCase__    (							self    :    Any      ,		_A    :    Union[str, Any]      ,		_A    :    Optional[Any]      ,		_A    :    Union[str, Any]      ,		_A    :    str      ,		_A    :    Optional[int]      ,		*_A    :    Optional[Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						TFGPTJModel(config=snake_case_				)

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

       # first forward pass
       __SCREAMING_SNAKE_CASE     :					int			=						model(snake_case_      ,		attention_mask=snake_case_      ,		token_type_ids=snake_case_      ,		use_cache=snake_case_				)

       __SCREAMING_SNAKE_CASE     :					List[str]			=						outputs.to_tuple()

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

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

       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						model(
           snake_case_      ,		token_type_ids=snake_case_      ,		attention_mask=snake_case_				)["""last_hidden_state"""]
       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						model(
           snake_case_      ,		token_type_ids=snake_case_      ,		attention_mask=snake_case_      ,		past_key_values=snake_case_      ,		)["""last_hidden_state"""]
       self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]				)

       # select random slice
       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						int(ids_tensor((1,)      ,		shape_list(snake_case_				)[-1]				)				)
       __SCREAMING_SNAKE_CASE     :					str			=						output_from_no_past[:, -3:, random_slice_idx]
       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						output_from_past[:, :, random_slice_idx]

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



     def 				UpperCAmelCase__    (							self    :    Union[str, Any]      ,		_A    :    int      ,		_A    :    Any      ,		_A    :    int      ,		_A    :    List[Any]      ,		_A    :    Tuple      ,		*_A    :    List[str]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					List[str]			=						TFGPTJForCausalLM(config=snake_case_				)
       __SCREAMING_SNAKE_CASE     :					Any			=						{
           """input_ids""": input_ids,
           """attention_mask""": input_mask,
           """token_type_ids""": token_type_ids,
       }
       __SCREAMING_SNAKE_CASE     :					Dict			=						model(snake_case_				)
       self.parent.assertEqual(result.logits.shape      ,		(self.batch_size, self.seq_length, self.vocab_size)				)




     def 				UpperCAmelCase__    (							self    :    Any				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						self.prepare_config_and_inputs()

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

       __SCREAMING_SNAKE_CASE     :					int			=						{
           """input_ids""": input_ids,
           """token_type_ids""": token_type_ids,
           """attention_mask""": input_mask,
       }
       return config, inputs_dict






@require_tf
class   __UpperCamelCase				(    _a				,		_a				,		_a				,		unittest.TestCase   ):






     """simple docstring"""







     lowerCAmelCase_     =	(
         (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
         if is_tf_available()
         else ()
     )

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



     def 				UpperCAmelCase__    (							self    :    str      ,		_A    :    List[str]      ,		_A    :    Optional[int]      ,		_A    :    Dict      ,		_A    :    Tuple      ,		_A    :    Optional[Any]				):







       """simple docstring"""




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

       return False



     def 				UpperCAmelCase__    (							self    :    Any				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Tuple			=						TFGPTJModelTester(self				)
       __SCREAMING_SNAKE_CASE     :					Dict			=						ConfigTester(self      ,		config_class=snake_case_      ,		n_embd=37				)



     def 				UpperCAmelCase__    (							self    :    Optional[Any]				):







       """simple docstring"""




       self.config_tester.run_common_tests()



     def 				UpperCAmelCase__    (							self    :    Optional[Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_gptj_model(*snake_case_				)



     def 				UpperCAmelCase__    (							self    :    List[Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_gptj_model_past(*snake_case_				)



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







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Dict			=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_gptj_model_attention_mask_past(*snake_case_				)



     def 				UpperCAmelCase__    (							self    :    List[Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Dict			=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_gptj_model_past_large_inputs(*snake_case_				)



     def 				UpperCAmelCase__    (							self    :    Any				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					str			=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_gptj_lm_head_model(*snake_case_				)



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







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					str			=						TFGPTJModel.from_pretrained('''EleutherAI/gpt-j-6B'''      ,		from_pt=snake_case_				)
       self.assertIsNotNone(snake_case_				)




     @unittest.skip(reason='''Currently, model embeddings are going to undergo a major refactor.'''				)
     def 				UpperCAmelCase__    (							self    :    List[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   __UpperCamelCase				(    unittest.TestCase   ):






     """simple docstring"""







     def 				UpperCAmelCase__    (							self    :    Optional[int]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					List[Any]			=						TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''      ,		from_pt=snake_case_				)
       __SCREAMING_SNAKE_CASE     :					str			=						tf.convert_to_tensor([[464, 3290]]      ,		dtype=tf.intaa				)  # The dog
       # fmt: off
       # The dog is a man's best friend. It is a loyal companion, and it is a friend
       __SCREAMING_SNAKE_CASE     :					int			=						[464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 1_5185, 11, 290, 340, 318, 257, 1545]
       # fmt: on
       __SCREAMING_SNAKE_CASE     :					str			=						model.generate(snake_case_      ,		do_sample=snake_case_				)
       self.assertListEqual(output_ids[0].numpy().tolist()      ,		snake_case_				)



     def 				UpperCAmelCase__    (							self    :    Tuple				):







       """simple docstring"""




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

       __SCREAMING_SNAKE_CASE     :					Any			=						tokenizer('''Today is a nice day and'''      ,		return_tensors='''tf'''				)
       # forces the generation to happen on CPU, to avoid GPU-related quirks
       with tf.device(''':/CPU:0'''				):
         __SCREAMING_SNAKE_CASE     :					Tuple			=						model.generate(**snake_case_      ,		do_sample=snake_case_      ,		seed=[42, 0]				)
       __SCREAMING_SNAKE_CASE     :					Tuple			=						tokenizer.decode(output_ids[0]      ,		skip_special_tokens=snake_case_				)

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



     def 				UpperCAmelCase__    (							self    :    Dict				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''      ,		revision='''float16'''      ,		from_pt=snake_case_				)
       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						AutoTokenizer.from_pretrained('''EleutherAI/gpt-j-6B'''      ,		revision='''float16'''				)

       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						"""left"""

       # Define PAD Token = EOS Token = 50256
       __SCREAMING_SNAKE_CASE     :					Dict			=						tokenizer.eos_token
       __SCREAMING_SNAKE_CASE     :					List[Any]			=						model.config.eos_token_id

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



     def 				UpperCAmelCase__    (							self    :    List[str]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Dict			=						self._get_beam_search_test_objects()

       __SCREAMING_SNAKE_CASE     :					Dict			=						tokenizer(snake_case_      ,		return_tensors='''tf'''      ,		padding=snake_case_				)
       __SCREAMING_SNAKE_CASE     :					Any			=						model.generate(**snake_case_      ,		do_sample=snake_case_      ,		num_beams=2				)
       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						tokenizer.batch_decode(snake_case_      ,		skip_special_tokens=snake_case_				)
       self.assertListEqual(snake_case_      ,		snake_case_				)



     def 				UpperCAmelCase__    (							self    :    Tuple				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Tuple			=						self._get_beam_search_test_objects()

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




     def 				UpperCAmelCase__    (							self    :    List[Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Dict			=						self._get_beam_search_test_objects()

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


import unittest

from transformers import load_tool
from transformers.tools.agent_types import AGENT_TYPE_MAPPING

from .test_tools_common import ToolTesterMixin, output_types





class   __UpperCamelCase				(    unittest.TestCase				,		lowerCAmelCase__   ):






     """simple docstring"""







     def 				UpperCAmelCase__    (							self    :    Optional[Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					str			=						load_tool('''translation'''				)
       self.tool.setup()
       __SCREAMING_SNAKE_CASE     :					List[Any]			=						load_tool('''translation'''      ,		remote=_A				)



     def 				UpperCAmelCase__    (							self    :    List[Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Any			=						self.tool('''Hey, what\'s up?'''      ,		src_lang='''English'''      ,		tgt_lang='''French'''				)
       self.assertEqual(_A      ,		'''- Hé, comment ça va?'''				)



     def 				UpperCAmelCase__    (							self    :    List[str]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					List[str]			=						self.remote_tool('''Hey, what\'s up?'''      ,		src_lang='''English'''      ,		tgt_lang='''French'''				)
       self.assertEqual(_A      ,		'''- Hé, comment ça va?'''				)



     def 				UpperCAmelCase__    (							self    :    Optional[Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					int			=						self.tool(text='''Hey, what\'s up?'''      ,		src_lang='''English'''      ,		tgt_lang='''French'''				)
       self.assertEqual(_A      ,		'''- Hé, comment ça va?'''				)



     def 				UpperCAmelCase__    (							self    :    int				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					List[Any]			=						self.remote_tool(text='''Hey, what\'s up?'''      ,		src_lang='''English'''      ,		tgt_lang='''French'''				)
       self.assertEqual(_A      ,		'''- Hé, comment ça va?'''				)



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







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						['''Hey, what\'s up?''', '''English''', '''Spanish''']
       __SCREAMING_SNAKE_CASE     :					List[Any]			=						self.tool(*_A				)

       # There is a single output
       if len(self.tool.outputs				) == 1:
         __SCREAMING_SNAKE_CASE     :					Any			=						[outputs]

       self.assertListEqual(output_types(_A				)      ,		self.tool.outputs				)



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







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					int			=						['''Hey, what\'s up?''', '''English''', '''Spanish''']
       __SCREAMING_SNAKE_CASE     :					List[str]			=						self.tool(*_A				)

       if not isinstance(_A      ,		_A				):
         __SCREAMING_SNAKE_CASE     :					Dict			=						[outputs]

       self.assertEqual(len(_A				)      ,		len(self.tool.outputs				)				)

       for output, output_type in zip(_A      ,		self.tool.outputs				):
         __SCREAMING_SNAKE_CASE     :					List[Any]			=						AGENT_TYPE_MAPPING[output_type]
         self.assertTrue(isinstance(_A      ,		_A				)				)




     def 				UpperCAmelCase__    (							self    :    List[str]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Dict			=						['''Hey, what\'s up?''', '''English''', '''Spanish''']

       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						[]

       for _input, input_type in zip(_A      ,		self.tool.inputs				):
         if isinstance(_A      ,		_A				):
           _inputs.append([AGENT_TYPE_MAPPING[_input_type](_input				) for _input_type in input_type]				)
         else:
           _inputs.append(AGENT_TYPE_MAPPING[input_type](_input				)				)

        # Should not raise an error
       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						self.tool(*_A				)

       if not isinstance(_A      ,		_A				):
         __SCREAMING_SNAKE_CASE     :					Dict			=						[outputs]

       self.assertEqual(len(_A				)      ,		len(self.tool.outputs				)				)
