

import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         MODEL_FOR_PRETRAINING_MAPPING,
         NezhaForMaskedLM,
         NezhaForMultipleChoice,
         NezhaForNextSentencePrediction,
         NezhaForPreTraining,
         NezhaForQuestionAnswering,
         NezhaForSequenceClassification,
         NezhaForTokenClassification,
         NezhaModel,
     )
     from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST





class  __UpperCAmelCase					:



   def __init__(							self							:					Dict,		__A							:					Optional[Any],		__A							:					Any=1_3,		__A							:					List[Any]=7,		__A							:					int=True,		__A							:					int=True,		__A							:					List[str]=True,		__A							:					Optional[int]=True,		__A							:					Tuple=9_9,		__A							:					str=3_2,		__A							:					List[str]=5,		__A							:					Optional[int]=4,		__A							:					int=3_7,		__A							:					Any="gelu",		__A							:					List[Any]=0.1,		__A							:					Union[str, Any]=0.1,		__A							:					Optional[Any]=1_2_8,		__A							:					Optional[Any]=3_2,		__A							:					Optional[int]=1_6,		__A							:					List[Any]=2,		__A							:					Optional[int]=0.0_2,		__A							:					Optional[int]=3,		__A							:					Optional[int]=4,		__A							:					List[Any]=None,		):
     UpperCAmelCase	:							List[Any]						=			parent
     UpperCAmelCase	:							Optional[Any]						=			batch_size
     UpperCAmelCase	:							Union[str, Any]						=			seq_length
     UpperCAmelCase	:							Any						=			is_training
     UpperCAmelCase	:							Any						=			use_input_mask
     UpperCAmelCase	:							List[str]						=			use_token_type_ids
     UpperCAmelCase	:							Dict						=			use_labels
     UpperCAmelCase	:							Union[str, Any]						=			vocab_size
     UpperCAmelCase	:							int						=			hidden_size
     UpperCAmelCase	:							Any						=			num_hidden_layers
     UpperCAmelCase	:							Any						=			num_attention_heads
     UpperCAmelCase	:							str						=			intermediate_size
     UpperCAmelCase	:							Optional[Any]						=			hidden_act
     UpperCAmelCase	:							Tuple						=			hidden_dropout_prob
     UpperCAmelCase	:							Any						=			attention_probs_dropout_prob
     UpperCAmelCase	:							Optional[int]						=			max_position_embeddings
     UpperCAmelCase	:							Optional[int]						=			type_vocab_size
     UpperCAmelCase	:							Tuple						=			type_sequence_label_size
     UpperCAmelCase	:							Optional[Any]						=			initializer_range
     UpperCAmelCase	:							Dict						=			num_labels
     UpperCAmelCase	:							Optional[int]						=			num_choices
     UpperCAmelCase	:							List[Any]						=			scope



   def 	__magic_name__   (							self							:					Optional[Any]						):
     UpperCAmelCase	:							Dict						=			ids_tensor([self.batch_size, self.seq_length],		self.vocab_size						)

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

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

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

     UpperCAmelCase	:							Dict						=			self.get_config()

     return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



   def 	__magic_name__   (							self							:					List[Any]						):
     return NezhaConfig(
         vocab_size=self.vocab_size,		hidden_size=self.hidden_size,		num_hidden_layers=self.num_hidden_layers,		num_attention_heads=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,		max_position_embeddings=self.max_position_embeddings,		type_vocab_size=self.type_vocab_size,		is_decoder=__A,		initializer_range=self.initializer_range,		)



   def 	__magic_name__   (							self							:					Any						):
     (
         (
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						
     )	:							Any						=			self.prepare_config_and_inputs()

     UpperCAmelCase	:							Tuple						=			True
     UpperCAmelCase	:							Dict						=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
     UpperCAmelCase	:							Optional[int]						=			ids_tensor([self.batch_size, self.seq_length],		vocab_size=2						)

     return (
         config,
         input_ids,
         token_type_ids,
         input_mask,
         sequence_labels,
         token_labels,
         choice_labels,
         encoder_hidden_states,
         encoder_attention_mask,
     )



   def 	__magic_name__   (							self							:					List[Any],		__A							:					Any,		__A							:					Optional[int],		__A							:					Any,		__A							:					Union[str, Any],		__A							:					List[Any],		__A							:					Union[str, Any],		__A							:					Union[str, Any]						):
     UpperCAmelCase	:							Optional[Any]						=			NezhaModel(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Dict						=			model(__A,		attention_mask=__A,		token_type_ids=__A						)
     UpperCAmelCase	:							Optional[Any]						=			model(__A,		token_type_ids=__A						)
     UpperCAmelCase	:							int						=			model(__A						)
     self.parent.assertEqual(result.last_hidden_state.shape,		(self.batch_size, self.seq_length, self.hidden_size)						)
     self.parent.assertEqual(result.pooler_output.shape,		(self.batch_size, self.hidden_size)						)



   def 	__magic_name__   (							self							:					int,		__A							:					List[str],		__A							:					Tuple,		__A							:					str,		__A							:					Optional[int],		__A							:					Optional[int],		__A							:					Optional[int],		__A							:					Optional[int],		__A							:					List[Any],		__A							:					List[str],		):
     UpperCAmelCase	:							Optional[Any]						=			True
     UpperCAmelCase	:							List[str]						=			NezhaModel(__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							str						=			model(
         __A,		attention_mask=__A,		token_type_ids=__A,		encoder_hidden_states=__A,		encoder_attention_mask=__A,		)
     UpperCAmelCase	:							int						=			model(
         __A,		attention_mask=__A,		token_type_ids=__A,		encoder_hidden_states=__A,		)
     UpperCAmelCase	:							Optional[int]						=			model(__A,		attention_mask=__A,		token_type_ids=__A						)
     self.parent.assertEqual(result.last_hidden_state.shape,		(self.batch_size, self.seq_length, self.hidden_size)						)
     self.parent.assertEqual(result.pooler_output.shape,		(self.batch_size, self.hidden_size)						)



   def 	__magic_name__   (							self							:					Dict,		__A							:					Any,		__A							:					Optional[Any],		__A							:					Any,		__A							:					Optional[int],		__A							:					Union[str, Any],		__A							:					Union[str, Any],		__A							:					Union[str, Any]						):
     UpperCAmelCase	:							List[str]						=			NezhaForMaskedLM(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Optional[int]						=			model(__A,		attention_mask=__A,		token_type_ids=__A,		labels=__A						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.vocab_size)						)



   def 	__magic_name__   (							self							:					Union[str, Any],		__A							:					Tuple,		__A							:					Any,		__A							:					Optional[int],		__A							:					Dict,		__A							:					Any,		__A							:					List[Any],		__A							:					Any						):
     UpperCAmelCase	:							List[str]						=			NezhaForNextSentencePrediction(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Dict						=			model(
         __A,		attention_mask=__A,		token_type_ids=__A,		labels=__A,		)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, 2)						)



   def 	__magic_name__   (							self							:					Tuple,		__A							:					Optional[Any],		__A							:					Optional[Any],		__A							:					Any,		__A							:					Optional[int],		__A							:					List[str],		__A							:					List[Any],		__A							:					Union[str, Any]						):
     UpperCAmelCase	:							str						=			NezhaForPreTraining(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Any						=			model(
         __A,		attention_mask=__A,		token_type_ids=__A,		labels=__A,		next_sentence_label=__A,		)
     self.parent.assertEqual(result.prediction_logits.shape,		(self.batch_size, self.seq_length, self.vocab_size)						)
     self.parent.assertEqual(result.seq_relationship_logits.shape,		(self.batch_size, 2)						)



   def 	__magic_name__   (							self							:					Any,		__A							:					Optional[int],		__A							:					List[Any],		__A							:					str,		__A							:					Tuple,		__A							:					str,		__A							:					Any,		__A							:					int						):
     UpperCAmelCase	:							Optional[Any]						=			NezhaForQuestionAnswering(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							List[str]						=			model(
         __A,		attention_mask=__A,		token_type_ids=__A,		start_positions=__A,		end_positions=__A,		)
     self.parent.assertEqual(result.start_logits.shape,		(self.batch_size, self.seq_length)						)
     self.parent.assertEqual(result.end_logits.shape,		(self.batch_size, self.seq_length)						)



   def 	__magic_name__   (							self							:					Optional[int],		__A							:					int,		__A							:					Optional[Any],		__A							:					Dict,		__A							:					Tuple,		__A							:					List[str],		__A							:					int,		__A							:					Optional[Any]						):
     UpperCAmelCase	:							Optional[Any]						=			self.num_labels
     UpperCAmelCase	:							str						=			NezhaForSequenceClassification(__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Union[str, Any]						=			model(__A,		attention_mask=__A,		token_type_ids=__A,		labels=__A						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.num_labels)						)



   def 	__magic_name__   (							self							:					Optional[Any],		__A							:					Dict,		__A							:					str,		__A							:					List[str],		__A							:					Dict,		__A							:					Optional[Any],		__A							:					Optional[Any],		__A							:					int						):
     UpperCAmelCase	:							Union[str, Any]						=			self.num_labels
     UpperCAmelCase	:							List[str]						=			NezhaForTokenClassification(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Optional[Any]						=			model(__A,		attention_mask=__A,		token_type_ids=__A,		labels=__A						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.num_labels)						)



   def 	__magic_name__   (							self							:					List[Any],		__A							:					Optional[int],		__A							:					Optional[int],		__A							:					List[Any],		__A							:					str,		__A							:					List[str],		__A							:					Optional[int],		__A							:					List[str]						):
     UpperCAmelCase	:							Any						=			self.num_choices
     UpperCAmelCase	:							Optional[int]						=			NezhaForMultipleChoice(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							List[str]						=			input_ids.unsqueeze(1						).expand(-1,		self.num_choices,		-1						).contiguous()
     UpperCAmelCase	:							str						=			token_type_ids.unsqueeze(1						).expand(-1,		self.num_choices,		-1						).contiguous()
     UpperCAmelCase	:							Union[str, Any]						=			input_mask.unsqueeze(1						).expand(-1,		self.num_choices,		-1						).contiguous()
     UpperCAmelCase	:							Optional[int]						=			model(
         __A,		attention_mask=__A,		token_type_ids=__A,		labels=__A,		)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.num_choices)						)





   def 	__magic_name__   (							self							:					Dict						):
     UpperCAmelCase	:							List[Any]						=			self.prepare_config_and_inputs()
     (
         (
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						
     )	:							Tuple						=			config_and_inputs
     UpperCAmelCase	:							Any						=			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
     return config, inputs_dict






@require_torch
class  __UpperCAmelCase					(	lowerCamelCase__      ,  lowerCamelCase__      ,  lowerCamelCase__      ,  unittest.TestCase						):
   UpperCamelCase       					=					(
       (
           NezhaModel,
           NezhaForMaskedLM,
           NezhaForMultipleChoice,
           NezhaForNextSentencePrediction,
           NezhaForPreTraining,
           NezhaForQuestionAnswering,
           NezhaForSequenceClassification,
           NezhaForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   UpperCamelCase       					=					(
       {
           """feature-extraction""": NezhaModel,
           """fill-mask""": NezhaForMaskedLM,
           """question-answering""": NezhaForQuestionAnswering,
           """text-classification""": NezhaForSequenceClassification,
           """token-classification""": NezhaForTokenClassification,
           """zero-shot""": NezhaForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   UpperCamelCase       					=					True



   def 	__magic_name__   (							self							:					str,		__A							:					Union[str, Any],		__A							:					int,		__A							:					List[str]=False						):
     UpperCAmelCase	:							Tuple						=			super()._prepare_for_class(__A,		__A,		return_labels=__A						)

     if return_labels:
       if model_class in get_values(__A						):
         UpperCAmelCase	:							Optional[Any]						=			torch.zeros(
             (self.model_tester.batch_size, self.model_tester.seq_length),		dtype=torch.long,		device=__A						)
         UpperCAmelCase	:							Optional[int]						=			torch.zeros(
             self.model_tester.batch_size,		dtype=torch.long,		device=__A						)
     return inputs_dict



   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase	:							Dict						=			NezhaModelTester(self						)
     UpperCAmelCase	:							Any						=			ConfigTester(self,		config_class=__A,		hidden_size=3_7						)



   def 	__magic_name__   (							self							:					List[Any]						):
     self.config_tester.run_common_tests()



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							Tuple						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__A						)



   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase	:							List[Any]						=			self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_model_as_decoder(*__A						)



   def 	__magic_name__   (							self							:					Tuple						):
     # This regression test was failing with PyTorch < 1.3
     (
         (
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						
     )	:							int						=			self.model_tester.prepare_config_and_inputs_for_decoder()

     UpperCAmelCase	:							str						=			None

     self.model_tester.create_and_check_model_as_decoder(
         __A,		__A,		__A,		__A,		__A,		__A,		__A,		__A,		__A,		)



   def 	__magic_name__   (							self							:					Dict						):
     UpperCAmelCase	:							int						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*__A						)



   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase	:							str						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*__A						)



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							Tuple						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_next_sequence_prediction(*__A						)



   def 	__magic_name__   (							self							:					Tuple						):
     UpperCAmelCase	:							Any						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_pretraining(*__A						)



   def 	__magic_name__   (							self							:					Dict						):
     UpperCAmelCase	:							str						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*__A						)



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							Optional[Any]						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*__A						)



   def 	__magic_name__   (							self							:					Optional[Any]						):
     UpperCAmelCase	:							Dict						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*__A						)



   @slow
   def 	__magic_name__   (							self							:					str						):
     for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       UpperCAmelCase	:							Optional[Any]						=			NezhaModel.from_pretrained(__A						)
       self.assertIsNotNone(__A						)





   @slow
   @require_torch_gpu
   def 	__magic_name__   (							self							:					Any						):
     UpperCAmelCase			,						UpperCAmelCase	:							Dict						=			self.model_tester.prepare_config_and_inputs_for_common()
     for model_class in self.all_model_classes:
       # NezhaForMultipleChoice behaves incorrectly in JIT environments.
       if model_class == NezhaForMultipleChoice:
         return

       UpperCAmelCase	:							Any						=			True
       UpperCAmelCase	:							Optional[Any]						=			model_class(config=__A						)

       UpperCAmelCase	:							Union[str, Any]						=			self._prepare_for_class(__A,		__A						)
       UpperCAmelCase	:							Dict						=			torch.jit.trace(
           __A,		(inputs_dict['''input_ids'''].to('''cpu'''						), inputs_dict['''attention_mask'''].to('''cpu'''						))						)

       with tempfile.TemporaryDirectory() as tmp:
         torch.jit.save(__A,		os.path.join(__A,		'''bert.pt'''						)						)
         UpperCAmelCase	:							Union[str, Any]						=			torch.jit.load(os.path.join(__A,		'''bert.pt'''						),		map_location=__A						)
         loaded(inputs_dict['''input_ids'''].to(__A						),		inputs_dict['''attention_mask'''].to(__A						)						)






@require_torch
class  __UpperCAmelCase					(	unittest.TestCase						):



   @slow
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase	:							Dict						=			NezhaModel.from_pretrained('''sijunhe/nezha-cn-base'''						)
     UpperCAmelCase	:							List[str]						=			torch.tensor([[0, 1, 2, 3, 4, 5]]						)
     UpperCAmelCase	:							Optional[int]						=			torch.tensor([[0, 1, 1, 1, 1, 1]]						)
     with torch.no_grad():
       UpperCAmelCase	:							List[str]						=			model(__A,		attention_mask=__A						)[0]
     UpperCAmelCase	:							Tuple						=			torch.Size((1, 6, 7_6_8)						)
     self.assertEqual(output.shape,		__A						)
     UpperCAmelCase	:							Optional[Any]						=			torch.tensor([[[0.0_6_8_5, 0.2_4_4_1, 0.1_1_0_2], [0.0_6_0_0, 0.1_9_0_6, 0.1_3_4_9], [0.0_2_2_1, 0.0_8_1_9, 0.0_5_8_6]]]						)

     self.assertTrue(torch.allclose(output[:, 1:4, 1:4],		__A,		atol=1E-4						)						)





   @slow
   def 	__magic_name__   (							self							:					Dict						):
     UpperCAmelCase	:							Optional[int]						=			NezhaForMaskedLM.from_pretrained('''sijunhe/nezha-cn-base'''						)
     UpperCAmelCase	:							int						=			torch.tensor([[0, 1, 2, 3, 4, 5]]						)
     UpperCAmelCase	:							Tuple						=			torch.tensor([[1, 1, 1, 1, 1, 1]]						)
     with torch.no_grad():
       UpperCAmelCase	:							Union[str, Any]						=			model(__A,		attention_mask=__A						)[0]
     UpperCAmelCase	:							Dict						=			torch.Size((1, 6, 2_1_1_2_8)						)
     self.assertEqual(output.shape,		__A						)
     UpperCAmelCase	:							List[Any]						=			torch.tensor(
         [[-2.7_9_3_9, -1.7_9_0_2, -2.2_1_8_9], [-2.8_5_8_5, -1.8_9_0_8, -2.3_7_2_3], [-2.6_4_9_9, -1.7_7_5_0, -2.2_5_5_8]]						)

     self.assertTrue(torch.allclose(output[:, 1:4, 1:4],		__A,		atol=1E-4						)						)








import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
     import os

     import jax.numpy as jnp
     from jax import jit

     from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
     from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

     _lowerCamelCase    :       Any        =							"0.12"  # assumed parallelism: 8


if is_torch_available():
     import torch

def       a__					(		UpperCAmelCase							:      int    ,			UpperCAmelCase							:      List[str]    ,			UpperCAmelCase							:      List[str]=None      )  ->						List[Any]:
  if rng is None:
    UpperCAmelCase	:							Dict						=			random.Random()

  UpperCAmelCase	:							Optional[Any]						=			1
  for dim in shape:
    total_dims *= dim

  UpperCAmelCase	:							List[str]						=			[]
  for _ in range(UpperCAmelCase      ):
    values.append(rng.randint(0    ,			vocab_size - 1      )      )

  UpperCAmelCase	:							List[str]						=			np.array(UpperCAmelCase    ,			dtype=jnp.intaa      ).reshape(UpperCAmelCase      )

  return output

def       a__					(		UpperCAmelCase							:      List[str]    ,			UpperCAmelCase							:      Optional[int]=None      )  ->						List[str]:
  UpperCAmelCase	:							Optional[int]						=			ids_tensor(UpperCAmelCase    ,			vocab_size=2    ,			rng=UpperCAmelCase      )
  # make sure that at least one token is attended to for each batch
  UpperCAmelCase	:							str						=			1
  return attn_mask





@require_flax
class  __UpperCAmelCase					:
   UpperCamelCase       					=					None
   UpperCamelCase       					=					()



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase			,						UpperCAmelCase	:							Tuple						=			self.model_tester.prepare_config_and_inputs_for_common()

     # cut to half length & take max batch_size 3
     UpperCAmelCase	:							Optional[int]						=			2
     UpperCAmelCase	:							Dict						=			inputs['''input_ids'''].shape[-1] // 2
     UpperCAmelCase	:							Dict						=			inputs['''input_ids'''][:max_batch_size, :sequence_length]

     UpperCAmelCase	:							Optional[int]						=			jnp.ones_like(__A						)
     UpperCAmelCase	:							Optional[int]						=			attention_mask[:max_batch_size, :sequence_length]

     # generate max 5 tokens
     UpperCAmelCase	:							Optional[Any]						=			input_ids.shape[-1] + 5
     if config.eos_token_id is not None and config.pad_token_id is None:
       # hack to allow generate for models such as GPT2 as is done in `generate()`
       UpperCAmelCase	:							Optional[Any]						=			config.eos_token_id
     return config, input_ids, attention_mask, max_length



   @is_pt_flax_cross_test
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Any						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							Optional[Any]						=			False
     UpperCAmelCase	:							Any						=			max_length
     UpperCAmelCase	:							List[Any]						=			0

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Union[str, Any]						=			model_class(__A						)

       UpperCAmelCase	:							List[str]						=			model_class.__name__[4:]  # Skip the "Flax" at the beginning
       UpperCAmelCase	:							List[Any]						=			getattr(__A,		__A						)
       UpperCAmelCase	:							Union[str, Any]						=			pt_model_class(__A						).eval()
       UpperCAmelCase	:							Tuple						=			load_flax_weights_in_pytorch_model(__A,		flax_model.params						)

       UpperCAmelCase	:							Dict						=			flax_model.generate(__A						).sequences
       UpperCAmelCase	:							str						=			pt_model.generate(torch.tensor(__A,		dtype=torch.long						)						)

       if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
         UpperCAmelCase	:							Any						=			flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

       self.assertListEqual(pt_generation_outputs.numpy().tolist(),		flax_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					Tuple						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Union[str, Any]						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							str						=			False
     UpperCAmelCase	:							Dict						=			max_length

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Union[str, Any]						=			model_class(__A						)

       UpperCAmelCase	:							Optional[int]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							List[Any]						=			jit(model.generate						)
       UpperCAmelCase	:							Optional[Any]						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							List[Any]						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							str						=			True
     UpperCAmelCase	:							Dict						=			max_length

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Union[str, Any]						=			model_class(__A						)

       UpperCAmelCase	:							Optional[Any]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							str						=			jit(model.generate						)
       UpperCAmelCase	:							List[Any]						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							int						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							Dict						=			False
     UpperCAmelCase	:							Union[str, Any]						=			max_length
     UpperCAmelCase	:							List[Any]						=			2

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							int						=			model_class(__A						)

       UpperCAmelCase	:							str						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							int						=			jit(model.generate						)
       UpperCAmelCase	:							Union[str, Any]						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					List[str]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Optional[Any]						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							Any						=			False
     UpperCAmelCase	:							Optional[int]						=			max_length
     UpperCAmelCase	:							Optional[int]						=			2
     UpperCAmelCase	:							str						=			2

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							int						=			model_class(__A						)

       UpperCAmelCase	:							Optional[Any]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[0],		input_ids.shape[0] * config.num_return_sequences						)



   def 	__magic_name__   (							self							:					Any						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							str						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							str						=			True
     UpperCAmelCase	:							Union[str, Any]						=			max_length
     UpperCAmelCase	:							Union[str, Any]						=			0.8
     UpperCAmelCase	:							str						=			1_0
     UpperCAmelCase	:							Any						=			0.3
     UpperCAmelCase	:							str						=			1
     UpperCAmelCase	:							Union[str, Any]						=			8
     UpperCAmelCase	:							Optional[Any]						=			9

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							int						=			model_class(__A						)

       UpperCAmelCase	:							List[Any]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							Optional[int]						=			jit(model.generate						)
       UpperCAmelCase	:							Any						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					List[Any]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Dict						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							Optional[Any]						=			max_length
     UpperCAmelCase	:							Tuple						=			1
     UpperCAmelCase	:							Optional[Any]						=			8
     UpperCAmelCase	:							Optional[int]						=			9

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							int						=			model_class(__A						)

       UpperCAmelCase	:							List[str]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							Dict						=			jit(model.generate						)
       UpperCAmelCase	:							List[str]						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					Optional[Any]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Union[str, Any]						=			self._get_input_ids_and_config()
     UpperCAmelCase	:							List[str]						=			max_length
     UpperCAmelCase	:							Union[str, Any]						=			2
     UpperCAmelCase	:							List[Any]						=			1
     UpperCAmelCase	:							List[str]						=			8
     UpperCAmelCase	:							int						=			9

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Optional[Any]						=			model_class(__A						)

       UpperCAmelCase	:							Union[str, Any]						=			model.generate(__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							List[str]						=			jit(model.generate						)
       UpperCAmelCase	:							Tuple						=			jit_generate(__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Optional[int]						=			self._get_input_ids_and_config()

     # pad attention mask on the left
     UpperCAmelCase	:							Union[str, Any]						=			attention_mask.at[(0, 0)].set(0						)

     UpperCAmelCase	:							Tuple						=			False
     UpperCAmelCase	:							str						=			max_length

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Any						=			model_class(__A						)

       UpperCAmelCase	:							Union[str, Any]						=			model.generate(__A,		attention_mask=__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							List[Any]						=			jit(model.generate						)
       UpperCAmelCase	:							Optional[Any]						=			jit_generate(__A,		attention_mask=__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)



   def 	__magic_name__   (							self							:					Optional[int]						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Optional[int]						=			self._get_input_ids_and_config()

     # pad attention mask on the left
     UpperCAmelCase	:							Union[str, Any]						=			attention_mask.at[(0, 0)].set(0						)

     UpperCAmelCase	:							Union[str, Any]						=			True
     UpperCAmelCase	:							Dict						=			max_length

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							Any						=			model_class(__A						)

       UpperCAmelCase	:							Optional[Any]						=			model.generate(__A,		attention_mask=__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							Optional[Any]						=			jit(model.generate						)
       UpperCAmelCase	:							Optional[Any]						=			jit_generate(__A,		attention_mask=__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)





   def 	__magic_name__   (							self							:					Tuple						):
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Union[str, Any]						=			self._get_input_ids_and_config()

     # pad attention mask on the left
     UpperCAmelCase	:							Dict						=			attention_mask.at[(0, 0)].set(0						)

     UpperCAmelCase	:							Union[str, Any]						=			2
     UpperCAmelCase	:							str						=			max_length

     for model_class in self.all_generative_model_classes:
       UpperCAmelCase	:							str						=			model_class(__A						)

       UpperCAmelCase	:							int						=			model.generate(__A,		attention_mask=__A						).sequences
       self.assertEqual(generation_outputs.shape[-1],		__A						)

       UpperCAmelCase	:							Optional[Any]						=			jit(model.generate						)
       UpperCAmelCase	:							Dict						=			jit_generate(__A,		attention_mask=__A						).sequences

       self.assertListEqual(generation_outputs.tolist(),		jit_generation_outputs.tolist()						)






@require_flax
class  __UpperCAmelCase					(	unittest.TestCase						):



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							Union[str, Any]						=			AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''						)
     UpperCAmelCase	:							List[str]						=			FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''						)

     UpperCAmelCase	:							int						=			'''Hello world'''
     UpperCAmelCase	:							Optional[int]						=			tokenizer(__A,		return_tensors='''np'''						).input_ids

     # typos are quickly detected (the correct argument is `do_sample`)
     with self.assertRaisesRegex(__A,		'''do_samples'''						):
       model.generate(__A,		do_samples=__A						)

     # arbitrary arguments that will not be used anywhere are also not accepted
     with self.assertRaisesRegex(__A,		'''foo'''						):
       UpperCAmelCase	:							Any						=			{'''foo''': '''bar'''}
       model.generate(__A,		**__A						)





