code: 
  
  
  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        lowercase_				:
  
  
  
  
  
  
  							"""simple docstring"""
  
  
  
  
  
  
  
  							def __init__(   self			,    _UpperCAmelCase		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		parent
  													a_					=		13
  													a_					=		7
  													a_					=		True
  													a_					=		True
  													a_					=		True
  													a_					=		True
  													a_					=		True
  													a_					=		99
  													a_					=		32
  													a_					=		4
  													a_					=		2
  													a_					=		4
  													a_					=		37
  													a_					=		"""gelu"""
  													a_					=		0.1
  													a_					=		0.1
  													a_					=		512
  													a_					=		16
  													a_					=		2
  													a_					=		0.0_2
  													a_					=		3
  													a_					=		4
  													a_					=		None
  													a_					=		self.vocab_size - 1
  													a_					=		self.vocab_size - 1
  													a_					=		self.vocab_size - 1
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		ids_tensor([self.batch_size, self.seq_length]			,    self.vocab_size		)
  
  													a_					=		None
  													if self.use_input_mask:
  																			a_					=		random_attention_mask([self.batch_size, self.seq_length]		)
  
  													a_					=		None
  													if self.use_token_type_ids:
  																			a_					=		ids_tensor([self.batch_size, self.seq_length]			,    self.type_vocab_size		)
  
  													a_					=		None
  													if self.use_mc_token_ids:
  																			a_					=		ids_tensor([self.batch_size, self.num_choices]			,    self.seq_length		)
  
  													a_					=		None
  													a_					=		None
  													a_					=		None
  													if self.use_labels:
  																			a_					=		ids_tensor([self.batch_size]			,    self.type_sequence_label_size		)
  																			a_					=		ids_tensor([self.batch_size, self.seq_length]			,    self.num_labels		)
  																			a_					=		ids_tensor([self.batch_size]			,    self.num_choices		)
  
  													a_					=		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			,    )
  
  													a_					=		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 						lowercase__     (   self			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    *_UpperCAmelCase		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		TFGPTJModel(config=_UpperCAmelCase		)
  													a_					=		{
  													    """input_ids""": input_ids,
  													    """attention_mask""": input_mask,
  													    """token_type_ids""": token_type_ids,
  													}
  													a_					=		model(_UpperCAmelCase		)
  
  													a_					=		[input_ids, None, input_mask]  # None is the input for 'past'
  													a_					=		model(_UpperCAmelCase		)
  
  													a_					=		model(_UpperCAmelCase		)
  
  													self.parent.assertEqual(result.last_hidden_state.shape			,    (self.batch_size, self.seq_length, self.hidden_size)		)
  
  
  							def 						lowercase__     (   self			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    *_UpperCAmelCase		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		TFGPTJModel(config=_UpperCAmelCase		)
  
  													# first forward pass
  													a_					=		model(_UpperCAmelCase			,    token_type_ids=_UpperCAmelCase			,    use_cache=_UpperCAmelCase		)
  													a_					=		model(_UpperCAmelCase			,    token_type_ids=_UpperCAmelCase		)
  													a_					=		model(_UpperCAmelCase			,    token_type_ids=_UpperCAmelCase			,    use_cache=_UpperCAmelCase		)
  
  													self.parent.assertTrue(len(_UpperCAmelCase		) == len(_UpperCAmelCase		)		)
  													self.parent.assertTrue(len(_UpperCAmelCase		) == len(_UpperCAmelCase		) + 1		)
  
  													a_			,   a_					=		outputs.to_tuple()
  
  													# create hypothetical next token and extent to next_input_ids
  													a_					=		ids_tensor((self.batch_size, 1)			,    config.vocab_size		)
  													a_					=		ids_tensor([self.batch_size, 1]			,    self.type_vocab_size		)
  
  													# append to next input_ids and token_type_ids
  													a_					=		tf.concat([input_ids, next_tokens]			,    axis=-1		)
  													a_					=		tf.concat([token_type_ids, next_token_types]			,    axis=-1		)
  
  													a_					=		model(_UpperCAmelCase			,    token_type_ids=_UpperCAmelCase		)["""last_hidden_state"""]
  													a_					=		model(_UpperCAmelCase			,    token_type_ids=_UpperCAmelCase			,    past_key_values=_UpperCAmelCase		)[
  													    """last_hidden_state"""
  													]
  
  													# select random slice
  													a_					=		int(ids_tensor((1,)			,    shape_list(_UpperCAmelCase		)[-1]		)		)
  													a_					=		output_from_no_past[:, -1, random_slice_idx]
  													a_					=		output_from_past[:, 0, random_slice_idx]
  
  													# test that outputs are equal for slice
  													tf.debugging.assert_near(_UpperCAmelCase			,    _UpperCAmelCase			,    rtol=1e-6		)
  
  
  							def 						lowercase__     (   self			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    *_UpperCAmelCase		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		TFGPTJModel(config=_UpperCAmelCase		)
  
  													# create attention mask
  													a_					=		self.seq_length // 2
  													a_					=		tf.ones((self.batch_size, half_seq_length)			,    dtype=tf.intaa		)
  													a_					=		tf.zeros((self.batch_size, self.seq_length - half_seq_length)			,    dtype=tf.intaa		)
  													a_					=		tf.concat([attn_mask_begin, attn_mask_end]			,    axis=1		)
  
  													# first forward pass
  													a_			,   a_					=		model(_UpperCAmelCase			,    attention_mask=_UpperCAmelCase		).to_tuple()
  
  													# create hypothetical next token and extent to next_input_ids
  													a_					=		ids_tensor((self.batch_size, 1)			,    config.vocab_size		)
  
  													# change a random masked slice from input_ids
  													a_					=		ids_tensor((1,)			,    _UpperCAmelCase		).numpy() + 1
  													a_					=		ids_tensor((self.batch_size, self.seq_length)			,    config.vocab_size		)
  													a_					=		tf.range(self.seq_length		) == (self.seq_length - random_seq_idx_to_change)
  													a_					=		tf.transpose(
  													    tf.broadcast_to(tf.expand_dims(_UpperCAmelCase			,    -1		)			,    (self.seq_length, self.batch_size)		)		)
  													a_					=		tf.where(_UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase		)
  
  													# append to next input_ids and attn_mask
  													a_					=		tf.concat([input_ids, next_tokens]			,    axis=-1		)
  													a_					=		tf.concat([attn_mask, tf.ones((shape_list(_UpperCAmelCase		)[0], 1)			,    dtype=tf.intaa		)]			,    axis=1		)
  
  													# get two different outputs
  													a_					=		model(_UpperCAmelCase			,    attention_mask=_UpperCAmelCase		)["""last_hidden_state"""]
  													a_					=		model(_UpperCAmelCase			,    past_key_values=_UpperCAmelCase			,    attention_mask=_UpperCAmelCase		)[
  													    """last_hidden_state"""
  													]
  
  													# select random slice
  													a_					=		int(ids_tensor((1,)			,    shape_list(_UpperCAmelCase		)[-1]		)		)
  													a_					=		output_from_no_past[:, -1, random_slice_idx]
  													a_					=		output_from_past[:, 0, random_slice_idx]
  
  													# test that outputs are equal for slice
  													tf.debugging.assert_near(_UpperCAmelCase			,    _UpperCAmelCase			,    rtol=1e-1_2		)
  
  
  							def 						lowercase__     (   self			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    *_UpperCAmelCase		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		TFGPTJModel(config=_UpperCAmelCase		)
  
  													a_					=		input_ids[:1, :]
  													a_					=		input_mask[:1, :]
  													a_					=		token_type_ids[:1, :]
  													a_					=		1
  
  													# first forward pass
  													a_					=		model(_UpperCAmelCase			,    attention_mask=_UpperCAmelCase			,    token_type_ids=_UpperCAmelCase			,    use_cache=_UpperCAmelCase		)
  
  													a_			,   a_					=		outputs.to_tuple()
  
  													# create hypothetical next token and extent to next_input_ids
  													a_					=		ids_tensor((self.batch_size, 3)			,    config.vocab_size		)
  													a_					=		ids_tensor((self.batch_size, 3)			,    2		)
  													a_					=		ids_tensor((self.batch_size, 3)			,    self.type_vocab_size		)
  
  													# append to next input_ids and token_type_ids
  													a_					=		tf.concat([input_ids, next_tokens]			,    axis=-1		)
  													a_					=		tf.concat([input_mask, next_attn_mask]			,    axis=-1		)
  													a_					=		tf.concat([token_type_ids, next_token_types]			,    axis=-1		)
  
  													a_					=		model(
  													    _UpperCAmelCase			,    token_type_ids=_UpperCAmelCase			,    attention_mask=_UpperCAmelCase		)["""last_hidden_state"""]
  													a_					=		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
  													a_					=		int(ids_tensor((1,)			,    shape_list(_UpperCAmelCase		)[-1]		)		)
  													a_					=		output_from_no_past[:, -3:, random_slice_idx]
  													a_					=		output_from_past[:, :, random_slice_idx]
  
  													# test that outputs are equal for slice
  													tf.debugging.assert_near(_UpperCAmelCase			,    _UpperCAmelCase			,    rtol=1e-3		)
  
  
  							def 						lowercase__     (   self			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    *_UpperCAmelCase		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		TFGPTJForCausalLM(config=_UpperCAmelCase		)
  													a_					=		{
  													    """input_ids""": input_ids,
  													    """attention_mask""": input_mask,
  													    """token_type_ids""": token_type_ids,
  													}
  													a_					=		model(_UpperCAmelCase		)
  													self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.seq_length, self.vocab_size)		)
  
  
  
  
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		self.prepare_config_and_inputs()
  
  													(
  													    (
  													    a_   
  													)			,   (
  													    a_   
  													)			,   (
  													    a_   
  													)			,   (
  													    a_   
  													)			,   (
  													    a_   
  													)			,   (
  													    a_   
  													)			,   (
  													    a_   
  													)			,   (
  													    a_   
  													)			,   (
  													    a_   
  													)			,      
  													)					=		config_and_inputs
  
  													a_					=		{
  													    """input_ids""": input_ids,
  													    """token_type_ids""": token_type_ids,
  													    """attention_mask""": input_mask,
  													}
  													return config, inputs_dict
  
  
  
  
  
  @require_tf
  class        lowercase_				( UpperCamelCase__			,UpperCamelCase__			,UpperCamelCase__			,unittest.TestCase):
  
  
  
  
  
  
  							"""simple docstring"""
  
  
  
  
  
  
  
  							snake_case_        =		(
  							    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
  							    if is_tf_available()
  							    else ()
  							)
  
  							snake_case_        =		(TFGPTJForCausalLM,) if is_tf_available() else ()
  							snake_case_        =		(
  							    {
  							        '''feature-extraction''': TFGPTJModel,
  							        '''question-answering''': TFGPTJForQuestionAnswering,
  							        '''text-classification''': TFGPTJForSequenceClassification,
  							        '''text-generation''': TFGPTJForCausalLM,
  							        '''zero-shot''': TFGPTJForSequenceClassification,
  							    }
  							    if is_tf_available()
  							    else {}
  							)
  							snake_case_        =		False
  							snake_case_        =		False
  							snake_case_        =		False
  							snake_case_        =		False
  
  
  							def 						lowercase__     (   self			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase		):
  
  
  
  
  
  
  
  													"""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 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		TFGPTJModelTester(self		)
  													a_					=		ConfigTester(self			,    config_class=_UpperCAmelCase			,    n_embd=37		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													self.config_tester.run_common_tests()
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		self.model_tester.prepare_config_and_inputs()
  													self.model_tester.create_and_check_gptj_model(*_UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		self.model_tester.prepare_config_and_inputs()
  													self.model_tester.create_and_check_gptj_model_past(*_UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		self.model_tester.prepare_config_and_inputs()
  													self.model_tester.create_and_check_gptj_model_attention_mask_past(*_UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		self.model_tester.prepare_config_and_inputs()
  													self.model_tester.create_and_check_gptj_model_past_large_inputs(*_UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		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 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		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 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""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        lowercase_				( unittest.TestCase):
  
  
  
  
  
  
  							"""simple docstring"""
  
  
  
  
  
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""			,    from_pt=_UpperCAmelCase		)
  													a_					=		tf.convert_to_tensor([[464, 3_290]]			,    dtype=tf.intaa		)  # The dog
  													# fmt: off
  													# The dog is a man's best friend. It is a loyal companion, and it is a friend
  													a_					=		[464, 3_290, 318, 257, 582, 338, 1_266, 1_545, 13, 632, 318, 257, 9_112, 15_185, 11, 290, 340, 318, 257, 1_545]
  													# fmt: on
  													a_					=		model.generate(_UpperCAmelCase			,    do_sample=_UpperCAmelCase		)
  													self.assertListEqual(output_ids[0].numpy().tolist()			,    _UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""			,    revision="""float16"""		)
  													a_					=		TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""			,    revision="""float16"""			,    from_pt=_UpperCAmelCase		)
  
  													a_					=		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"""		):
  																			a_					=		model.generate(**_UpperCAmelCase			,    do_sample=_UpperCAmelCase			,    seed=[42, 0]		)
  													a_					=		tokenizer.decode(output_ids[0]			,    skip_special_tokens=_UpperCAmelCase		)
  
  													a_					=		"""Today is a nice day and I’m going to go for a walk. I’"""
  													self.assertEqual(_UpperCAmelCase			,    _UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""			,    revision="""float16"""			,    from_pt=_UpperCAmelCase		)
  													a_					=		AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""			,    revision="""float16"""		)
  
  													a_					=		"""left"""
  
  													# Define PAD Token = EOS Token = 50256
  													a_					=		tokenizer.eos_token
  													a_					=		model.config.eos_token_id
  
  													# use different length sentences to test batching
  													a_					=		[
  													    """Hello, my dog is a little""",
  													    """Today, I""",
  													]
  													a_					=		[
  													    """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 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_			,   a_			,   a_			,   a_					=		self._get_beam_search_test_objects()
  
  													a_					=		tokenizer(_UpperCAmelCase			,    return_tensors="""tf"""			,    padding=_UpperCAmelCase		)
  													a_					=		model.generate(**_UpperCAmelCase			,    do_sample=_UpperCAmelCase			,    num_beams=2		)
  													a_					=		tokenizer.batch_decode(_UpperCAmelCase			,    skip_special_tokens=_UpperCAmelCase		)
  													self.assertListEqual(_UpperCAmelCase			,    _UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_			,   a_			,   a_			,   a_					=		self._get_beam_search_test_objects()
  
  													a_					=		tokenizer(_UpperCAmelCase			,    return_tensors="""tf"""			,    padding=_UpperCAmelCase		)
  													a_					=		tokenizer(sentences[0]			,    return_tensors="""tf"""		)
  													a_					=		model.generate(**_UpperCAmelCase			,    do_sample=_UpperCAmelCase			,    num_beams=2		)
  													a_					=		(
  													    shape_list(inputs_non_padded["""input_ids"""]		)[-1]
  													    - tf.reduce_sum(tf.cast(inputs["""attention_mask"""][-1]			,    tf.intaa		)		).numpy()
  													)
  													a_					=		tokenizer(sentences[1]			,    return_tensors="""tf"""		)
  													a_					=		model.generate(
  													    **_UpperCAmelCase			,    do_sample=_UpperCAmelCase			,    num_beams=2			,    max_length=model.config.max_length - num_paddings		)
  													a_					=		tokenizer.decode(output_non_padded[0]			,    skip_special_tokens=_UpperCAmelCase		)
  													a_					=		tokenizer.decode(output_padded[0]			,    skip_special_tokens=_UpperCAmelCase		)
  													self.assertListEqual(_UpperCAmelCase			,    [non_padded_sentence, padded_sentence]		)
  
  
  
  
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_			,   a_			,   a_			,   a_					=		self._get_beam_search_test_objects()
  
  													a_					=		tokenizer(_UpperCAmelCase			,    return_tensors="""tf"""			,    padding=_UpperCAmelCase		)
  													a_					=		tf.function(model.generate			,    jit_compile=_UpperCAmelCase		)
  													a_					=		xla_generate(**_UpperCAmelCase			,    do_sample=_UpperCAmelCase			,    num_beams=2		)
  													a_					=		tokenizer.batch_decode(_UpperCAmelCase			,    skip_special_tokens=_UpperCAmelCase		)
  													self.assertListEqual(_UpperCAmelCase			,    _UpperCAmelCase		)
style_context: 
  
  
  import unittest
  
  import numpy as np
  import timeout_decorator  # noqa
  
  from transformers import MBartConfig, is_flax_available
  from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
  from transformers.utils import cached_property
  
  from ...generation.test_flax_utils import FlaxGenerationTesterMixin
  from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor
  
  
  if is_flax_available():
  	import os
  
  	# The slow tests are often failing with OOM error on GPU
  	# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
  	# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
  	A_    :					List[Any]             ="""platform"""
  
  	import jax
  	import jax.numpy as jnp
  
  	from transformers import AutoTokenizer
  	from transformers.models.mbart.modeling_flax_mbart import (
  	    FlaxMBartForConditionalGeneration,
  	    FlaxMBartForQuestionAnswering,
  	    FlaxMBartForSequenceClassification,
  	    FlaxMBartModel,
  	    shift_tokens_right,
  	)
  
  
  
  
  def   lowerCamelCase_       (					UpperCAmelCase__       ,    UpperCAmelCase__       ,    UpperCAmelCase__=None       ,    UpperCAmelCase__=None       ,    UpperCAmelCase__=None       ,    UpperCAmelCase__=None       ,    UpperCAmelCase__=None       ,    UpperCAmelCase__=None       ,    ):
  
  
  
  						"""simple docstring"""
  
  
  
  
  						if attention_mask is None:
  												a_					=		np.where(input_ids != config.pad_token_id       ,    1       ,    0      )
  						if decoder_attention_mask is None:
  												a_					=		np.where(decoder_input_ids != config.pad_token_id       ,    1       ,    0      )
  						if head_mask is None:
  												a_					=		np.ones((config.encoder_layers, config.encoder_attention_heads)      )
  						if decoder_head_mask is None:
  												a_					=		np.ones((config.decoder_layers, config.decoder_attention_heads)      )
  						if cross_attn_head_mask is None:
  												a_					=		np.ones((config.decoder_layers, config.decoder_attention_heads)      )
  						return {
  						    "input_ids": input_ids,
  						    "decoder_input_ids": decoder_input_ids,
  						    "attention_mask": attention_mask,
  						    "decoder_attention_mask": decoder_attention_mask,
  						}
  
  
  
  
  
  class        lowercase_				( unittest.TestCase):
  
  
  
  
  
  
  							"""simple docstring"""
  
  
  
  
  
  
  
  							def __init__(   self			,    _UpperCAmelCase			,    _UpperCAmelCase=13			,    _UpperCAmelCase=7			,    _UpperCAmelCase=True			,    _UpperCAmelCase=False			,    _UpperCAmelCase=99			,    _UpperCAmelCase=16			,    _UpperCAmelCase=2			,    _UpperCAmelCase=4			,    _UpperCAmelCase=4			,    _UpperCAmelCase="gelu"			,    _UpperCAmelCase=0.1			,    _UpperCAmelCase=0.1			,    _UpperCAmelCase=32			,    _UpperCAmelCase=2			,    _UpperCAmelCase=1			,    _UpperCAmelCase=0			,    _UpperCAmelCase=2			,    _UpperCAmelCase=0.0_2			,    ):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		parent
  													a_					=		batch_size
  													a_					=		seq_length
  													a_					=		is_training
  													a_					=		use_labels
  													a_					=		vocab_size
  													a_					=		hidden_size
  													a_					=		num_hidden_layers
  													a_					=		num_attention_heads
  													a_					=		intermediate_size
  													a_					=		hidden_act
  													a_					=		hidden_dropout_prob
  													a_					=		attention_probs_dropout_prob
  													a_					=		max_position_embeddings
  													a_					=		eos_token_id
  													a_					=		pad_token_id
  													a_					=		bos_token_id
  													a_					=		decoder_start_token_id
  													a_					=		initializer_range
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		np.clip(ids_tensor([self.batch_size, self.seq_length - 1]			,    self.vocab_size		)			,    3			,    self.vocab_size		)
  													a_					=		np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)			,    dtype=np.intaa		))			,    -1		)
  
  													a_					=		shift_tokens_right(_UpperCAmelCase			,    1		)
  
  													a_					=		MBartConfig(
  													    vocab_size=self.vocab_size			,    d_model=self.hidden_size			,    encoder_layers=self.num_hidden_layers			,    decoder_layers=self.num_hidden_layers			,    encoder_attention_heads=self.num_attention_heads			,    decoder_attention_heads=self.num_attention_heads			,    encoder_ffn_dim=self.intermediate_size			,    decoder_ffn_dim=self.intermediate_size			,    dropout=self.hidden_dropout_prob			,    attention_dropout=self.attention_probs_dropout_prob			,    max_position_embeddings=self.max_position_embeddings			,    eos_token_id=self.eos_token_id			,    bos_token_id=self.bos_token_id			,    pad_token_id=self.pad_token_id			,    decoder_start_token_id=self.decoder_start_token_id			,    initializer_range=self.initializer_range			,    use_cache=_UpperCAmelCase			,    )
  													a_					=		prepare_mbart_inputs_dict(_UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase		)
  													return config, inputs_dict
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_			,   a_					=		self.prepare_config_and_inputs()
  													return config, inputs_dict
  
  
  							def 						lowercase__     (   self			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		20
  													a_					=		model_class_name(_UpperCAmelCase		)
  
  													a_					=		model.encode(inputs_dict["""input_ids"""]		)
  
  													a_			,   a_					=		(
  													    inputs_dict["""decoder_input_ids"""],
  													    inputs_dict["""decoder_attention_mask"""],
  													)
  
  													a_					=		model.init_cache(decoder_input_ids.shape[0]			,    _UpperCAmelCase			,    _UpperCAmelCase		)
  													a_					=		jnp.ones((decoder_input_ids.shape[0], max_decoder_length)			,    dtype="""i4"""		)
  
  													a_					=		jnp.broadcast_to(
  													    jnp.arange(decoder_input_ids.shape[-1] - 1		)[None, :]			,    (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)			,    )
  													a_					=		model.decode(
  													    decoder_input_ids[:, :-1]			,    _UpperCAmelCase			,    decoder_attention_mask=_UpperCAmelCase			,    past_key_values=_UpperCAmelCase			,    decoder_position_ids=_UpperCAmelCase			,    )
  
  													a_					=		jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]			,    dtype="""i4"""		)
  													a_					=		model.decode(
  													    decoder_input_ids[:, -1:]			,    _UpperCAmelCase			,    decoder_attention_mask=_UpperCAmelCase			,    past_key_values=outputs_cache.past_key_values			,    decoder_position_ids=_UpperCAmelCase			,    )
  
  													a_					=		model.decode(_UpperCAmelCase			,    _UpperCAmelCase		)
  
  													a_					=		np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])		)		)
  													self.parent.assertTrue(diff < 1e-3			,    msg=f"Max diff is {diff}"		)
  
  
  
  
  
  
  							def 						lowercase__     (   self			,    _UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		20
  													a_					=		model_class_name(_UpperCAmelCase		)
  
  													a_					=		model.encode(inputs_dict["""input_ids"""]		)
  
  													a_			,   a_					=		(
  													    inputs_dict["""decoder_input_ids"""],
  													    inputs_dict["""decoder_attention_mask"""],
  													)
  
  													a_					=		jnp.concatenate(
  													    [
  													        decoder_attention_mask,
  													        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])		),
  													    ]			,    axis=-1			,    )
  
  													a_					=		model.init_cache(decoder_input_ids.shape[0]			,    _UpperCAmelCase			,    _UpperCAmelCase		)
  													a_					=		jnp.broadcast_to(
  													    jnp.arange(decoder_input_ids.shape[-1] - 1		)[None, :]			,    (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)			,    )
  
  													a_					=		model.decode(
  													    decoder_input_ids[:, :-1]			,    _UpperCAmelCase			,    decoder_attention_mask=_UpperCAmelCase			,    past_key_values=_UpperCAmelCase			,    decoder_position_ids=_UpperCAmelCase			,    )
  													a_					=		jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]			,    dtype="""i4"""		)
  													a_					=		model.decode(
  													    decoder_input_ids[:, -1:]			,    _UpperCAmelCase			,    past_key_values=outputs_cache.past_key_values			,    decoder_attention_mask=_UpperCAmelCase			,    decoder_position_ids=_UpperCAmelCase			,    )
  
  													a_					=		model.decode(_UpperCAmelCase			,    _UpperCAmelCase			,    decoder_attention_mask=_UpperCAmelCase		)
  
  													a_					=		np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])		)		)
  													self.parent.assertTrue(diff < 1e-3			,    msg=f"Max diff is {diff}"		)
  
  
  
  
  
  @require_flax
  class        lowercase_				( unittest.TestCase):
  
  
  
  
  
  
  							"""simple docstring"""
  
  
  
  
  
  
  
  							snake_case_        =		99
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		np.array(
  													    [
  													        [71, 82, 18, 33, 46, 91, 2],
  													        [68, 34, 26, 58, 30, 82, 2],
  													        [5, 97, 17, 39, 94, 40, 2],
  													        [76, 83, 94, 25, 70, 78, 2],
  													        [87, 59, 41, 35, 48, 66, 2],
  													        [55, 13, 16, 58, 5, 2, 1],  # note padding
  													        [64, 27, 31, 51, 12, 75, 2],
  													        [52, 64, 86, 17, 83, 39, 2],
  													        [48, 61, 9, 24, 71, 82, 2],
  													        [26, 1, 60, 48, 22, 13, 2],
  													        [21, 5, 62, 28, 14, 76, 2],
  													        [45, 98, 37, 86, 59, 48, 2],
  													        [70, 70, 50, 9, 28, 0, 2],
  													    ]			,    dtype=np.intaa			,    )
  
  													a_					=		input_ids.shape[0]
  													a_					=		MBartConfig(
  													    vocab_size=self.vocab_size			,    d_model=24			,    encoder_layers=2			,    decoder_layers=2			,    encoder_attention_heads=2			,    decoder_attention_heads=2			,    encoder_ffn_dim=32			,    decoder_ffn_dim=32			,    max_position_embeddings=48			,    eos_token_id=2			,    pad_token_id=1			,    bos_token_id=0			,    )
  													return config, input_ids, batch_size
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_			,   a_			,   a_					=		self._get_config_and_data()
  													a_					=		FlaxMBartForSequenceClassification(_UpperCAmelCase		)
  													a_					=		model(input_ids=_UpperCAmelCase			,    decoder_input_ids=_UpperCAmelCase		)
  													a_					=		(batch_size, config.num_labels)
  													self.assertEqual(outputs["""logits"""].shape			,    _UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_			,   a_			,   a_					=		self._get_config_and_data()
  													a_					=		FlaxMBartForQuestionAnswering(_UpperCAmelCase		)
  													a_					=		model(input_ids=_UpperCAmelCase		)
  
  													self.assertEqual(outputs["""start_logits"""].shape			,    input_ids.shape		)
  													self.assertEqual(outputs["""end_logits"""].shape			,    input_ids.shape		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_			,   a_			,   a_					=		self._get_config_and_data()
  													a_					=		FlaxMBartForConditionalGeneration(_UpperCAmelCase		)
  													a_					=		lm_model(input_ids=_UpperCAmelCase		)
  													a_					=		(batch_size, input_ids.shape[1], config.vocab_size)
  													self.assertEqual(outputs["""logits"""].shape			,    _UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		MBartConfig(
  													    vocab_size=self.vocab_size			,    d_model=14			,    encoder_layers=2			,    decoder_layers=2			,    encoder_attention_heads=2			,    decoder_attention_heads=2			,    encoder_ffn_dim=8			,    decoder_ffn_dim=8			,    max_position_embeddings=48			,    )
  													a_					=		FlaxMBartForConditionalGeneration(_UpperCAmelCase		)
  													a_					=		np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]			,    dtype=np.intaa		)
  													a_					=		np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]			,    dtype=np.intaa		)
  													a_					=		lm_model(input_ids=_UpperCAmelCase			,    decoder_input_ids=_UpperCAmelCase		)
  													a_					=		(*summary.shape, config.vocab_size)
  													self.assertEqual(outputs["""logits"""].shape			,    _UpperCAmelCase		)
  
  
  
  
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]			,    dtype=np.intaa		)
  													a_					=		shift_tokens_right(_UpperCAmelCase			,    1		)
  													a_					=		np.equal(_UpperCAmelCase			,    1		).astype(np.floataa		).sum()
  													a_					=		np.equal(_UpperCAmelCase			,    1		).astype(np.floataa		).sum()
  													self.assertEqual(shifted.shape			,    input_ids.shape		)
  													self.assertEqual(_UpperCAmelCase			,    n_pad_before - 1		)
  													self.assertTrue(np.equal(shifted[:, 0]			,    2		).all()		)
  
  
  
  
  
  @require_flax
  class        lowercase_				( UpperCamelCase__			,unittest.TestCase			,UpperCamelCase__):
  
  
  
  
  
  
  							"""simple docstring"""
  
  
  
  
  
  
  
  							snake_case_        =		True
  							snake_case_        =		(
  							    (
  							        FlaxMBartModel,
  							        FlaxMBartForConditionalGeneration,
  							        FlaxMBartForSequenceClassification,
  							        FlaxMBartForQuestionAnswering,
  							    )
  							    if is_flax_available()
  							    else ()
  							)
  							snake_case_        =		(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		FlaxMBartModelTester(self		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_			,   a_					=		self.model_tester.prepare_config_and_inputs()
  													for model_class in self.all_model_classes:
  																			self.model_tester.check_use_cache_forward(_UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_			,   a_					=		self.model_tester.prepare_config_and_inputs()
  													for model_class in self.all_model_classes:
  																			self.model_tester.check_use_cache_forward_with_attn_mask(_UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_			,   a_					=		self.model_tester.prepare_config_and_inputs_for_common()
  
  													for model_class in self.all_model_classes:
  																			with self.subTest(model_class.__name__		):
  																									a_					=		self._prepare_for_class(_UpperCAmelCase			,    _UpperCAmelCase		)
  																									a_					=		model_class(_UpperCAmelCase		)
  
  																									@jax.jit
  																									def encode_jitted(_UpperCAmelCase			,    _UpperCAmelCase=None			,    **_UpperCAmelCase		):
  																															return model.encode(input_ids=_UpperCAmelCase			,    attention_mask=_UpperCAmelCase		)
  
  																									with self.subTest("""JIT Enabled"""		):
  																															a_					=		encode_jitted(**_UpperCAmelCase		).to_tuple()
  
  																									with self.subTest("""JIT Disabled"""		):
  																															with jax.disable_jit():
  																																					a_					=		encode_jitted(**_UpperCAmelCase		).to_tuple()
  
  																									self.assertEqual(len(_UpperCAmelCase		)			,    len(_UpperCAmelCase		)		)
  																									for jitted_output, output in zip(_UpperCAmelCase			,    _UpperCAmelCase		):
  																															self.assertEqual(jitted_output.shape			,    output.shape		)
  
  
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_			,   a_					=		self.model_tester.prepare_config_and_inputs_for_common()
  
  													for model_class in self.all_model_classes:
  																			with self.subTest(model_class.__name__		):
  																									a_					=		model_class(_UpperCAmelCase		)
  																									a_					=		model.encode(inputs_dict["""input_ids"""]			,    inputs_dict["""attention_mask"""]		)
  
  																									a_					=		{
  																									    """decoder_input_ids""": inputs_dict["""decoder_input_ids"""],
  																									    """decoder_attention_mask""": inputs_dict["""decoder_attention_mask"""],
  																									    """encoder_outputs""": encoder_outputs,
  																									}
  
  																									@jax.jit
  																									def decode_jitted(_UpperCAmelCase			,    _UpperCAmelCase			,    _UpperCAmelCase		):
  																															return model.decode(
  																															    decoder_input_ids=_UpperCAmelCase			,    decoder_attention_mask=_UpperCAmelCase			,    encoder_outputs=_UpperCAmelCase			,    )
  
  																									with self.subTest("""JIT Enabled"""		):
  																															a_					=		decode_jitted(**_UpperCAmelCase		).to_tuple()
  
  																									with self.subTest("""JIT Disabled"""		):
  																															with jax.disable_jit():
  																																					a_					=		decode_jitted(**_UpperCAmelCase		).to_tuple()
  
  																									self.assertEqual(len(_UpperCAmelCase		)			,    len(_UpperCAmelCase		)		)
  																									for jitted_output, output in zip(_UpperCAmelCase			,    _UpperCAmelCase		):
  																															self.assertEqual(jitted_output.shape			,    output.shape		)
  
  
  
  
  
  
  							@slow
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													for model_class_name in self.all_model_classes:
  																			a_					=		model_class_name.from_pretrained("""facebook/mbart-large-cc25"""			,    from_pt=_UpperCAmelCase		)
  																			# FlaxMBartForSequenceClassification expects eos token in input_ids
  																			a_					=		np.ones((1, 1)		) * model.config.eos_token_id
  																			a_					=		model(_UpperCAmelCase		)
  																			self.assertIsNotNone(_UpperCAmelCase		)
  
  
  
  
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class        lowercase_				( unittest.TestCase):
  
  
  
  
  
  
  							"""simple docstring"""
  
  
  
  
  
  
  
  							snake_case_        =		[
  							    ''' UN Chief Says There Is No Military Solution in Syria''',
  							]
  							snake_case_        =		[
  							    '''Şeful ONU declară că nu există o soluţie militară în Siria''',
  							]
  							snake_case_        =		'''facebook/mbart-large-en-ro'''
  
  
  							@cached_property
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													return AutoTokenizer.from_pretrained(self.model_name		)
  
  
  							@cached_property
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		FlaxMBartForConditionalGeneration.from_pretrained(self.model_name			,    from_pt=_UpperCAmelCase		)
  													return model
  
  
  							def 						lowercase__     (   self			,    **_UpperCAmelCase		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		self.translate_src_text(**_UpperCAmelCase		)
  													self.assertListEqual(self.expected_text			,    _UpperCAmelCase		)
  
  
  							def 						lowercase__     (   self			,    **_UpperCAmelCase		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													a_					=		self.tokenizer(self.src_text			,    **_UpperCAmelCase			,    return_tensors="""np"""		)
  													a_					=		self.model.generate(
  													    model_inputs.input_ids			,    attention_mask=model_inputs.attention_mask			,    decoder_start_token_id=self.tokenizer.lang_code_to_id["""ro_RO"""]			,    early_stopping=_UpperCAmelCase			,    num_beams=2			,    ).sequences
  													a_					=		self.tokenizer.batch_decode(_UpperCAmelCase			,    skip_special_tokens=_UpperCAmelCase		)
  													return generated_words
  
  
  
  
  
  
  							@slow
  							def 						lowercase__     (   self		):
  
  
  
  
  
  
  
  													"""simple docstring"""
  
  
  													self._assert_generated_batch_equal_expected()