



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




	"""simple docstring"""






	def __init__(     self				,UpperCAmelCase_   ):
			_lowercase				: Dict     =   parent
			_lowercase				: List[Any]     =   13
			_lowercase				: Any     =   7
			_lowercase				: List[str]     =   True
			_lowercase				: Dict     =   True
			_lowercase				: Optional[Any]     =   True
			_lowercase				: Union[str, Any]     =   True
			_lowercase				: List[Any]     =   True
			_lowercase				: List[str]     =   99
			_lowercase				: Any     =   32
			_lowercase				: List[Any]     =   4
			_lowercase				: Optional[Any]     =   2
			_lowercase				: List[Any]     =   4
			_lowercase				: Dict     =   37
			_lowercase				: Dict     =   """gelu"""
			_lowercase				: Tuple     =   0.1
			_lowercase				: str     =   0.1
			_lowercase				: int     =   5_12
			_lowercase				: int     =   16
			_lowercase				: List[Any]     =   2
			_lowercase				: List[Any]     =   0.02
			_lowercase				: List[Any]     =   3
			_lowercase				: Optional[Any]     =   4
			_lowercase				: Union[str, Any]     =   None
			_lowercase				: List[str]     =   self.vocab_size - 1
			_lowercase				: List[Any]     =   self.vocab_size - 1
			_lowercase				: int     =   self.vocab_size - 1






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Optional[Any]     =   ids_tensor([self.batch_size, self.seq_length]				,self.vocab_size   )

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

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

			_lowercase				: List[str]     =   None
			if self.use_mc_token_ids:
					_lowercase				: Tuple     =   ids_tensor([self.batch_size, self.num_choices]				,self.seq_length   )

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

			_lowercase				: Optional[int]     =   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=UpperCAmelCase_				,)

			_lowercase				: Optional[Any]     =   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    lowerCamelCase__					(     self				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,*UpperCAmelCase_   ):
			_lowercase				: Union[str, Any]     =   TFGPTJModel(config=UpperCAmelCase_   )
			_lowercase				: str     =   {
			    """input_ids""": input_ids,
			    """attention_mask""": input_mask,
			    """token_type_ids""": token_type_ids,
			}
			_lowercase				: Optional[int]     =   model(UpperCAmelCase_   )

			_lowercase				: Tuple     =   [input_ids, None, input_mask]  # None is the input for 'past'
			_lowercase				: List[str]     =   model(UpperCAmelCase_   )

			_lowercase				: Tuple     =   model(UpperCAmelCase_   )

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






	def    lowerCamelCase__					(     self				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,*UpperCAmelCase_   ):
			_lowercase				: Optional[int]     =   TFGPTJModel(config=UpperCAmelCase_   )

			# first forward pass
			_lowercase				: Optional[int]     =   model(UpperCAmelCase_				,token_type_ids=UpperCAmelCase_				,use_cache=UpperCAmelCase_   )
			_lowercase				: Optional[Any]     =   model(UpperCAmelCase_				,token_type_ids=UpperCAmelCase_   )
			_lowercase				: Union[str, Any]     =   model(UpperCAmelCase_				,token_type_ids=UpperCAmelCase_				,use_cache=UpperCAmelCase_   )

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

			_lowercase     ,  _lowercase				: List[Any]     =   outputs.to_tuple()

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

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

			_lowercase				: int     =   model(UpperCAmelCase_				,token_type_ids=UpperCAmelCase_   )["""last_hidden_state"""]
			_lowercase				: Any     =   model(UpperCAmelCase_				,token_type_ids=UpperCAmelCase_				,past_key_values=UpperCAmelCase_   )[
			    """last_hidden_state"""
			]

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

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






	def    lowerCamelCase__					(     self				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,*UpperCAmelCase_   ):
			_lowercase				: int     =   TFGPTJModel(config=UpperCAmelCase_   )

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

			# first forward pass
			_lowercase     ,  _lowercase				: Optional[Any]     =   model(UpperCAmelCase_				,attention_mask=UpperCAmelCase_   ).to_tuple()

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

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

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

			# get two different outputs
			_lowercase				: Optional[Any]     =   model(UpperCAmelCase_				,attention_mask=UpperCAmelCase_   )["""last_hidden_state"""]
			_lowercase				: Optional[Any]     =   model(UpperCAmelCase_				,past_key_values=UpperCAmelCase_				,attention_mask=UpperCAmelCase_   )[
			    """last_hidden_state"""
			]

			# select random slice
			_lowercase				: int     =   int(ids_tensor((1,)				,shape_list(UpperCAmelCase_   )[-1]   )   )
			_lowercase				: List[str]     =   output_from_no_past[:, -1, random_slice_idx]
			_lowercase				: Optional[Any]     =   output_from_past[:, 0, random_slice_idx]

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






	def    lowerCamelCase__					(     self				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,*UpperCAmelCase_   ):
			_lowercase				: Union[str, Any]     =   TFGPTJModel(config=UpperCAmelCase_   )

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

			# first forward pass
			_lowercase				: Optional[int]     =   model(UpperCAmelCase_				,attention_mask=UpperCAmelCase_				,token_type_ids=UpperCAmelCase_				,use_cache=UpperCAmelCase_   )

			_lowercase     ,  _lowercase				: List[Any]     =   outputs.to_tuple()

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

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

			_lowercase				: Union[str, Any]     =   model(
			    UpperCAmelCase_				,token_type_ids=UpperCAmelCase_				,attention_mask=UpperCAmelCase_   )["""last_hidden_state"""]
			_lowercase				: List[str]     =   model(
			    UpperCAmelCase_				,token_type_ids=UpperCAmelCase_				,attention_mask=UpperCAmelCase_				,past_key_values=UpperCAmelCase_				,)["""last_hidden_state"""]
			self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]   )

			# select random slice
			_lowercase				: int     =   int(ids_tensor((1,)				,shape_list(UpperCAmelCase_   )[-1]   )   )
			_lowercase				: Union[str, Any]     =   output_from_no_past[:, -3:, random_slice_idx]
			_lowercase				: Dict     =   output_from_past[:, :, random_slice_idx]

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






	def    lowerCamelCase__					(     self				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,*UpperCAmelCase_   ):
			_lowercase				: Tuple     =   TFGPTJForCausalLM(config=UpperCAmelCase_   )
			_lowercase				: Tuple     =   {
			    """input_ids""": input_ids,
			    """attention_mask""": input_mask,
			    """token_type_ids""": token_type_ids,
			}
			_lowercase				: List[str]     =   model(UpperCAmelCase_   )
			self.parent.assertEqual(result.logits.shape				,(self.batch_size, self.seq_length, self.vocab_size)   )






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Optional[int]     =   self.prepare_config_and_inputs()

			(
			    (
			    _lowercase
			)     ,  (
			    _lowercase
			)     ,  (
			    _lowercase
			)     ,  (
			    _lowercase
			)     ,  (
			    _lowercase
			)     ,  (
			    _lowercase
			)     ,  (
			    _lowercase
			)     ,  (
			    _lowercase
			)     ,  (
			    _lowercase
			)     ,  
			)				: Optional[int]     =   config_and_inputs

			_lowercase				: Union[str, Any]     =   {
			    """input_ids""": input_ids,
			    """token_type_ids""": token_type_ids,
			    """attention_mask""": input_mask,
			}
			return config, inputs_dict







@require_tf
class   UpperCamelCase	(     snake_case					,   snake_case					,   snake_case					,   unittest.TestCase       ):




	"""simple docstring"""
	SCREAMING_SNAKE_CASE_							:       int						       =  (
	    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
	    if is_tf_available()
	    else ()
	)

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






	def    lowerCamelCase__					(     self				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_   ):
			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    lowerCamelCase__					(     self   ):
			_lowercase				: str     =   TFGPTJModelTester(self   )
			_lowercase				: List[str]     =   ConfigTester(self				,config_class=UpperCAmelCase_				,n_embd=37   )






	def    lowerCamelCase__					(     self   ):
			self.config_tester.run_common_tests()






	def    lowerCamelCase__					(     self   ):
			_lowercase				: List[Any]     =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_gptj_model(*UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Dict     =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_gptj_model_past(*UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Optional[Any]     =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_gptj_model_attention_mask_past(*UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Any     =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_gptj_model_past_large_inputs(*UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Dict     =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_gptj_lm_head_model(*UpperCAmelCase_   )






	@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    lowerCamelCase__					(     self   ):
			_lowercase				: str     =   TFGPTJModel.from_pretrained("""EleutherAI/gpt-j-6B"""				,from_pt=UpperCAmelCase_   )
			self.assertIsNotNone(UpperCAmelCase_   )






	@unittest.skip(reason="""Currently, model embeddings are going to undergo a major refactor."""   )
	def    lowerCamelCase__					(     self   ):
			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    lowerCamelCase__					(     self   ):
			_lowercase				: Dict     =   TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""				,from_pt=UpperCAmelCase_   )
			_lowercase				: List[str]     =   tf.convert_to_tensor([[4_64, 32_90]]				,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
			_lowercase				: Union[str, Any]     =   [4_64, 32_90, 3_18, 2_57, 5_82, 3_38, 12_66, 15_45, 13, 6_32, 3_18, 2_57, 91_12, 1_51_85, 11, 2_90, 3_40, 3_18, 2_57, 15_45]
			# fmt: on
			_lowercase				: List[str]     =   model.generate(UpperCAmelCase_				,do_sample=UpperCAmelCase_   )
			self.assertListEqual(output_ids[0].numpy().tolist()				,UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Optional[Any]     =   AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""				,revision="""float16"""   )
			_lowercase				: int     =   TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""				,revision="""float16"""				,from_pt=UpperCAmelCase_   )

			_lowercase				: str     =   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"""   ):
					_lowercase				: List[Any]     =   model.generate(**UpperCAmelCase_				,do_sample=UpperCAmelCase_				,seed=[42, 0]   )
			_lowercase				: Any     =   tokenizer.decode(output_ids[0]				,skip_special_tokens=UpperCAmelCase_   )

			_lowercase				: Union[str, Any]     =   """Today is a nice day and I’m going to go for a walk. I’"""
			self.assertEqual(UpperCAmelCase_				,UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase				: List[Any]     =   TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""				,revision="""float16"""				,from_pt=UpperCAmelCase_   )
			_lowercase				: int     =   AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""				,revision="""float16"""   )

			_lowercase				: List[Any]     =   """left"""

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

			# use different length sentences to test batching
			_lowercase				: List[str]     =   [
			    """Hello, my dog is a little""",
			    """Today, I""",
			]
			_lowercase				: List[str]     =   [
			    """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    lowerCamelCase__					(     self   ):
			# Confirms that we get the expected results with left-padded beam search
			_lowercase     ,  _lowercase     ,  _lowercase     ,  _lowercase				: List[str]     =   self._get_beam_search_test_objects()

			_lowercase				: List[str]     =   tokenizer(UpperCAmelCase_				,return_tensors="""tf"""				,padding=UpperCAmelCase_   )
			_lowercase				: List[str]     =   model.generate(**UpperCAmelCase_				,do_sample=UpperCAmelCase_				,num_beams=2   )
			_lowercase				: List[str]     =   tokenizer.batch_decode(UpperCAmelCase_				,skip_special_tokens=UpperCAmelCase_   )
			self.assertListEqual(UpperCAmelCase_				,UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			# Confirms that left-padding is working properly
			_lowercase     ,  _lowercase     ,  _lowercase     ,  _lowercase				: List[str]     =   self._get_beam_search_test_objects()

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






	def    lowerCamelCase__					(     self   ):
			# Confirms that XLA is working properly
			_lowercase     ,  _lowercase     ,  _lowercase     ,  _lowercase				: List[Any]     =   self._get_beam_search_test_objects()

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








import unittest

import pytest

from transformers import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    MBartaaTokenizerFast,
    MBartConfig,
    MBartForConditionalGeneration,
    TranslationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, require_tf, require_torch, slow

from .test_pipelines_common import ANY


@is_pipeline_test
class    _lowercase (unittest.TestCase   ):
				'''simple docstring'''







				lowercase__     =    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
				lowercase__     =    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

				def    _lowerCamelCase       (			self  ,			snake_case__  ,			snake_case__  ,			snake_case__     ):

									'''simple docstring'''


									if isinstance(model.config  ,			snake_case__     ):
														UpperCamelCase_	,     UpperCamelCase_			= list(tokenizer.lang_code_to_id.keys()     )[:2]
														UpperCamelCase_			= TranslationPipeline(model=snake_case__  ,			tokenizer=snake_case__  ,			src_lang=snake_case__  ,			tgt_lang=snake_case__     )
									else:
														UpperCamelCase_			= TranslationPipeline(model=snake_case__  ,			tokenizer=snake_case__     )
									return translator, ["Some string", "Some other text"]

				def    _lowerCamelCase       (			self  ,			snake_case__  ,			snake_case__     ):

									'''simple docstring'''


									UpperCamelCase_			= translator("Some string"     )
									self.assertEqual(snake_case__  ,			[{"translation_text": ANY(snake_case__     )}]     )

									UpperCamelCase_			= translator(["Some string"]     )
									self.assertEqual(snake_case__  ,			[{"translation_text": ANY(snake_case__     )}]     )

									UpperCamelCase_			= translator(["Some string", "other string"]     )
									self.assertEqual(snake_case__  ,			[{"translation_text": ANY(snake_case__     )}, {"translation_text": ANY(snake_case__     )}]     )

				@require_torch
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= pipeline("translation_en_to_ro"  ,			model="patrickvonplaten/t5-tiny-random"  ,			framework="pt"     )
									UpperCamelCase_			= translator("This is a test string"  ,			max_length=20     )
									self.assertEqual(
									    snake_case__  ,			[
									        {
									            "translation_text": (
									                "Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide"
									                " Beide Beide"
									            )
									        }
									    ]  ,			)

				@require_tf
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= pipeline("translation_en_to_ro"  ,			model="patrickvonplaten/t5-tiny-random"  ,			framework="tf"     )
									UpperCamelCase_			= translator("This is a test string"  ,			max_length=20     )
									self.assertEqual(
									    snake_case__  ,			[
									        {
									            "translation_text": (
									                "Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide"
									                " Beide Beide"
									            )
									        }
									    ]  ,			)

				@require_torch
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= pipeline("translation_en_to_de"  ,			model="patrickvonplaten/t5-tiny-random"  ,			framework="pt"     )
									UpperCamelCase_			= translator("This is a test string"  ,			max_length=20     )
									self.assertEqual(
									    snake_case__  ,			[
									        {
									            "translation_text": (
									                "monoton monoton monoton monoton monoton monoton monoton monoton monoton monoton urine urine"
									                " urine urine urine urine urine urine urine"
									            )
									        }
									    ]  ,			)

				@require_tf
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= pipeline("translation_en_to_de"  ,			model="patrickvonplaten/t5-tiny-random"  ,			framework="tf"     )
									UpperCamelCase_			= translator("This is a test string"  ,			max_length=20     )
									self.assertEqual(
									    snake_case__  ,			[
									        {
									            "translation_text": (
									                "monoton monoton monoton monoton monoton monoton monoton monoton monoton monoton urine urine"
									                " urine urine urine urine urine urine urine"
									            )
									        }
									    ]  ,			)




class    _lowercase (unittest.TestCase   ):
				'''simple docstring'''







				@require_torch
				@slow
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									with self.assertRaises(snake_case__     ):
														pipeline(task="translation_cn_to_ar"     )

									# but we do for this one
									UpperCamelCase_			= pipeline(task="translation_en_to_de"     )
									self.assertEqual(translator._preprocess_params["src_lang"]  ,			"en"     )
									self.assertEqual(translator._preprocess_params["tgt_lang"]  ,			"de"     )

				@require_torch
				@slow
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt"     )
									UpperCamelCase_			= MBartaaTokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt"     )

									UpperCamelCase_			= pipeline(task="translation"  ,			model=snake_case__  ,			tokenizer=snake_case__     )
									# Missing src_lang, tgt_lang
									with self.assertRaises(snake_case__     ):
														translator("This is a test"     )

									UpperCamelCase_			= translator("This is a test"  ,			src_lang="en_XX"  ,			tgt_lang="ar_AR"     )
									self.assertEqual(snake_case__  ,			[{"translation_text": "هذا إختبار"}]     )

									UpperCamelCase_			= translator("This is a test"  ,			src_lang="en_XX"  ,			tgt_lang="hi_IN"     )
									self.assertEqual(snake_case__  ,			[{"translation_text": "यह एक परीक्षण है"}]     )

									# src_lang, tgt_lang can be defined at pipeline call time
									UpperCamelCase_			= pipeline(task="translation"  ,			model=snake_case__  ,			tokenizer=snake_case__  ,			src_lang="en_XX"  ,			tgt_lang="ar_AR"     )
									UpperCamelCase_			= translator("This is a test"     )
									self.assertEqual(snake_case__  ,			[{"translation_text": "هذا إختبار"}]     )

				@require_torch
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= "patrickvonplaten/t5-tiny-random"
									UpperCamelCase_			= pipeline(task="translation_cn_to_ar"  ,			model=snake_case__     )
									self.assertEqual(translator._preprocess_params["src_lang"]  ,			"cn"     )
									self.assertEqual(translator._preprocess_params["tgt_lang"]  ,			"ar"     )

				@require_torch
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= "patrickvonplaten/t5-tiny-random"
									with pytest.warns(snake_case__  ,			match=R".*translation_en_to_de.*"     ):
														UpperCamelCase_			= pipeline(task="translation"  ,			model=snake_case__     )
									self.assertEqual(translator.task  ,			"translation_en_to_de"     )
									self.assertEqual(translator._preprocess_params["src_lang"]  ,			"en"     )
									self.assertEqual(translator._preprocess_params["tgt_lang"]  ,			"de"     )

				@require_torch
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									with self.assertRaises(snake_case__     ):
														pipeline(task="translation"     )
