

"""simple docstring"""




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				:						Optional[int]    ,				__UpperCAmelCase				:						Optional[int]    ,				__UpperCAmelCase				:						str=1_3    ,				__UpperCAmelCase				:						List[Any]=7    ,				__UpperCAmelCase				:						Optional[Any]=True    ,				__UpperCAmelCase				:						Tuple=True    ,				__UpperCAmelCase				:						Union[str, Any]=True    ,				__UpperCAmelCase				:						Any=True    ,				__UpperCAmelCase				:						List[Any]=9_9    ,				__UpperCAmelCase				:						Optional[Any]=3_2    ,				__UpperCAmelCase				:						Dict=5    ,				__UpperCAmelCase				:						Union[str, Any]=4    ,				__UpperCAmelCase				:						int=3_7    ,				__UpperCAmelCase				:						List[str]="gelu"    ,				__UpperCAmelCase				:						List[Any]=0.1    ,				__UpperCAmelCase				:						List[str]=0.1    ,				__UpperCAmelCase				:						Optional[int]=1_2_8    ,				__UpperCAmelCase				:						Any=3_2    ,				__UpperCAmelCase				:						List[Any]=1_6    ,				__UpperCAmelCase				:						str=2    ,				__UpperCAmelCase				:						Union[str, Any]=0.02    ,				__UpperCAmelCase				:						Optional[int]=3    ,				__UpperCAmelCase				:						Optional[int]=4    ,				__UpperCAmelCase				:						Optional[int]=None    ,				)   ->    Any:
						SCREAMING_SNAKE_CASE__								=      parent
						SCREAMING_SNAKE_CASE__								=      batch_size
						SCREAMING_SNAKE_CASE__								=      seq_length
						SCREAMING_SNAKE_CASE__								=      is_training
						SCREAMING_SNAKE_CASE__								=      use_input_mask
						SCREAMING_SNAKE_CASE__								=      use_token_type_ids
						SCREAMING_SNAKE_CASE__								=      use_labels
						SCREAMING_SNAKE_CASE__								=      vocab_size
						SCREAMING_SNAKE_CASE__								=      hidden_size
						SCREAMING_SNAKE_CASE__								=      num_hidden_layers
						SCREAMING_SNAKE_CASE__								=      num_attention_heads
						SCREAMING_SNAKE_CASE__								=      intermediate_size
						SCREAMING_SNAKE_CASE__								=      hidden_act
						SCREAMING_SNAKE_CASE__								=      hidden_dropout_prob
						SCREAMING_SNAKE_CASE__								=      attention_probs_dropout_prob
						SCREAMING_SNAKE_CASE__								=      max_position_embeddings
						SCREAMING_SNAKE_CASE__								=      type_vocab_size
						SCREAMING_SNAKE_CASE__								=      type_sequence_label_size
						SCREAMING_SNAKE_CASE__								=      initializer_range
						SCREAMING_SNAKE_CASE__								=      num_labels
						SCREAMING_SNAKE_CASE__								=      num_choices
						SCREAMING_SNAKE_CASE__								=      scope

			def 					SCREAMING_SNAKE_CASE				(					self				:						int       )   ->    Tuple:
						SCREAMING_SNAKE_CASE__								=      ids_tensor([self.batch_size, self.seq_length]    ,				self.vocab_size       )

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

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

						SCREAMING_SNAKE_CASE__								=      None
						SCREAMING_SNAKE_CASE__								=      None
						SCREAMING_SNAKE_CASE__								=      None
						if self.use_labels:
									SCREAMING_SNAKE_CASE__								=      ids_tensor([self.batch_size]    ,				self.type_sequence_label_size       )
									SCREAMING_SNAKE_CASE__								=      ids_tensor([self.batch_size, self.seq_length]    ,				self.num_labels       )
									SCREAMING_SNAKE_CASE__								=      ids_tensor([self.batch_size]    ,				self.num_choices       )

						SCREAMING_SNAKE_CASE__								=      self.get_config()

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

			def 					SCREAMING_SNAKE_CASE				(					self				:						Tuple       )   ->    Optional[int]:
						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=__UpperCamelCase    ,				initializer_range=self.initializer_range    ,				)

			def 					SCREAMING_SNAKE_CASE				(					self				:						List[Any]       )   ->    Union[str, Any]:
						(
						    (
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		
						)								=      self.prepare_config_and_inputs()

						SCREAMING_SNAKE_CASE__								=      True
						SCREAMING_SNAKE_CASE__								=      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
						SCREAMING_SNAKE_CASE__								=      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 					SCREAMING_SNAKE_CASE				(					self				:						int    ,				__UpperCAmelCase				:						Dict    ,				__UpperCAmelCase				:						Optional[Any]    ,				__UpperCAmelCase				:						str    ,				__UpperCAmelCase				:						Union[str, Any]    ,				__UpperCAmelCase				:						Dict    ,				__UpperCAmelCase				:						Tuple    ,				__UpperCAmelCase				:						int       )   ->    Any:
						SCREAMING_SNAKE_CASE__								=      NezhaModel(config=__UpperCamelCase       )
						model.to(__UpperCamelCase       )
						model.eval()
						SCREAMING_SNAKE_CASE__								=      model(__UpperCamelCase    ,				attention_mask=__UpperCamelCase    ,				token_type_ids=__UpperCamelCase       )
						SCREAMING_SNAKE_CASE__								=      model(__UpperCamelCase    ,				token_type_ids=__UpperCamelCase       )
						SCREAMING_SNAKE_CASE__								=      model(__UpperCamelCase       )
						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 					SCREAMING_SNAKE_CASE				(					self				:						Optional[int]    ,				__UpperCAmelCase				:						Tuple    ,				__UpperCAmelCase				:						Dict    ,				__UpperCAmelCase				:						Optional[Any]    ,				__UpperCAmelCase				:						Dict    ,				__UpperCAmelCase				:						Union[str, Any]    ,				__UpperCAmelCase				:						int    ,				__UpperCAmelCase				:						Optional[Any]    ,				__UpperCAmelCase				:						List[Any]    ,				__UpperCAmelCase				:						Union[str, Any]    ,				)   ->    int:
						SCREAMING_SNAKE_CASE__								=      True
						SCREAMING_SNAKE_CASE__								=      NezhaModel(__UpperCamelCase       )
						model.to(__UpperCamelCase       )
						model.eval()
						SCREAMING_SNAKE_CASE__								=      model(
						    __UpperCamelCase    ,				attention_mask=__UpperCamelCase    ,				token_type_ids=__UpperCamelCase    ,				encoder_hidden_states=__UpperCamelCase    ,				encoder_attention_mask=__UpperCamelCase    ,				)
						SCREAMING_SNAKE_CASE__								=      model(
						    __UpperCamelCase    ,				attention_mask=__UpperCamelCase    ,				token_type_ids=__UpperCamelCase    ,				encoder_hidden_states=__UpperCamelCase    ,				)
						SCREAMING_SNAKE_CASE__								=      model(__UpperCamelCase    ,				attention_mask=__UpperCamelCase    ,				token_type_ids=__UpperCamelCase       )
						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 					SCREAMING_SNAKE_CASE				(					self				:						Dict    ,				__UpperCAmelCase				:						Dict    ,				__UpperCAmelCase				:						Optional[int]    ,				__UpperCAmelCase				:						Tuple    ,				__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						Union[str, Any]    ,				__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						Any       )   ->    Optional[int]:
						SCREAMING_SNAKE_CASE__								=      NezhaForMaskedLM(config=__UpperCamelCase       )
						model.to(__UpperCamelCase       )
						model.eval()
						SCREAMING_SNAKE_CASE__								=      model(__UpperCamelCase    ,				attention_mask=__UpperCamelCase    ,				token_type_ids=__UpperCamelCase    ,				labels=__UpperCamelCase       )
						self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.vocab_size)       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Tuple    ,				__UpperCAmelCase				:						Tuple    ,				__UpperCAmelCase				:						Any    ,				__UpperCAmelCase				:						Any    ,				__UpperCAmelCase				:						int    ,				__UpperCAmelCase				:						str    ,				__UpperCAmelCase				:						Dict    ,				__UpperCAmelCase				:						int       )   ->    Any:
						SCREAMING_SNAKE_CASE__								=      NezhaForNextSentencePrediction(config=__UpperCamelCase       )
						model.to(__UpperCamelCase       )
						model.eval()
						SCREAMING_SNAKE_CASE__								=      model(
						    __UpperCamelCase    ,				attention_mask=__UpperCamelCase    ,				token_type_ids=__UpperCamelCase    ,				labels=__UpperCamelCase    ,				)
						self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, 2)       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Any    ,				__UpperCAmelCase				:						Optional[Any]    ,				__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						Any    ,				__UpperCAmelCase				:						List[Any]    ,				__UpperCAmelCase				:						Optional[int]    ,				__UpperCAmelCase				:						Union[str, Any]       )   ->    List[Any]:
						SCREAMING_SNAKE_CASE__								=      NezhaForPreTraining(config=__UpperCamelCase       )
						model.to(__UpperCamelCase       )
						model.eval()
						SCREAMING_SNAKE_CASE__								=      model(
						    __UpperCamelCase    ,				attention_mask=__UpperCamelCase    ,				token_type_ids=__UpperCamelCase    ,				labels=__UpperCamelCase    ,				next_sentence_label=__UpperCamelCase    ,				)
						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 					SCREAMING_SNAKE_CASE				(					self				:						Union[str, Any]    ,				__UpperCAmelCase				:						Optional[int]    ,				__UpperCAmelCase				:						Tuple    ,				__UpperCAmelCase				:						Any    ,				__UpperCAmelCase				:						str    ,				__UpperCAmelCase				:						List[Any]    ,				__UpperCAmelCase				:						Dict    ,				__UpperCAmelCase				:						Tuple       )   ->    Any:
						SCREAMING_SNAKE_CASE__								=      NezhaForQuestionAnswering(config=__UpperCamelCase       )
						model.to(__UpperCamelCase       )
						model.eval()
						SCREAMING_SNAKE_CASE__								=      model(
						    __UpperCamelCase    ,				attention_mask=__UpperCamelCase    ,				token_type_ids=__UpperCamelCase    ,				start_positions=__UpperCamelCase    ,				end_positions=__UpperCamelCase    ,				)
						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 					SCREAMING_SNAKE_CASE				(					self				:						Optional[Any]    ,				__UpperCAmelCase				:						int    ,				__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						int    ,				__UpperCAmelCase				:						Optional[int]    ,				__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						List[str]       )   ->    Tuple:
						SCREAMING_SNAKE_CASE__								=      self.num_labels
						SCREAMING_SNAKE_CASE__								=      NezhaForSequenceClassification(__UpperCamelCase       )
						model.to(__UpperCamelCase       )
						model.eval()
						SCREAMING_SNAKE_CASE__								=      model(__UpperCamelCase    ,				attention_mask=__UpperCamelCase    ,				token_type_ids=__UpperCamelCase    ,				labels=__UpperCamelCase       )
						self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.num_labels)       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						List[Any]    ,				__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						Any    ,				__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						int    ,				__UpperCAmelCase				:						List[Any]    ,				__UpperCAmelCase				:						Union[str, Any]    ,				__UpperCAmelCase				:						Union[str, Any]       )   ->    Tuple:
						SCREAMING_SNAKE_CASE__								=      self.num_labels
						SCREAMING_SNAKE_CASE__								=      NezhaForTokenClassification(config=__UpperCamelCase       )
						model.to(__UpperCamelCase       )
						model.eval()
						SCREAMING_SNAKE_CASE__								=      model(__UpperCamelCase    ,				attention_mask=__UpperCamelCase    ,				token_type_ids=__UpperCamelCase    ,				labels=__UpperCamelCase       )
						self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.num_labels)       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Tuple    ,				__UpperCAmelCase				:						Optional[Any]    ,				__UpperCAmelCase				:						Any    ,				__UpperCAmelCase				:						Optional[Any]    ,				__UpperCAmelCase				:						Union[str, Any]    ,				__UpperCAmelCase				:						Union[str, Any]    ,				__UpperCAmelCase				:						Dict    ,				__UpperCAmelCase				:						str       )   ->    int:
						SCREAMING_SNAKE_CASE__								=      self.num_choices
						SCREAMING_SNAKE_CASE__								=      NezhaForMultipleChoice(config=__UpperCamelCase       )
						model.to(__UpperCamelCase       )
						model.eval()
						SCREAMING_SNAKE_CASE__								=      input_ids.unsqueeze(1       ).expand(-1    ,				self.num_choices    ,				-1       ).contiguous()
						SCREAMING_SNAKE_CASE__								=      token_type_ids.unsqueeze(1       ).expand(-1    ,				self.num_choices    ,				-1       ).contiguous()
						SCREAMING_SNAKE_CASE__								=      input_mask.unsqueeze(1       ).expand(-1    ,				self.num_choices    ,				-1       ).contiguous()
						SCREAMING_SNAKE_CASE__								=      model(
						    __UpperCamelCase    ,				attention_mask=__UpperCamelCase    ,				token_type_ids=__UpperCamelCase    ,				labels=__UpperCamelCase    ,				)
						self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.num_choices)       )





			def 					SCREAMING_SNAKE_CASE				(					self				:						Union[str, Any]       )   ->    Optional[int]:
						SCREAMING_SNAKE_CASE__								=      self.prepare_config_and_inputs()
						(
						    (
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		
						)								=      config_and_inputs
						SCREAMING_SNAKE_CASE__								=      {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
						return config, inputs_dict





@require_torch
class    __UpperCamelCase						(_UpperCAmelCase			,_UpperCAmelCase			,_UpperCAmelCase			,unittest.TestCase			):
			lowerCamelCase__   :		List[str]							=							(
			    (
			        NezhaModel,
			        NezhaForMaskedLM,
			        NezhaForMultipleChoice,
			        NezhaForNextSentencePrediction,
			        NezhaForPreTraining,
			        NezhaForQuestionAnswering,
			        NezhaForSequenceClassification,
			        NezhaForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			lowerCamelCase__   :		Union[str, Any]							=							(
			    {
			        'feature-extraction': NezhaModel,
			        'fill-mask': NezhaForMaskedLM,
			        'question-answering': NezhaForQuestionAnswering,
			        'text-classification': NezhaForSequenceClassification,
			        'token-classification': NezhaForTokenClassification,
			        'zero-shot': NezhaForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			lowerCamelCase__   :		int							=							True

			def 					SCREAMING_SNAKE_CASE				(					self				:						Union[str, Any]    ,				__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						List[Any]    ,				__UpperCAmelCase				:						List[Any]=False       )   ->    Dict:
						SCREAMING_SNAKE_CASE__								=      super()._prepare_for_class(__UpperCamelCase    ,				__UpperCamelCase    ,				return_labels=__UpperCamelCase       )

						if return_labels:
									if model_class in get_values(__UpperCamelCase       ):
												SCREAMING_SNAKE_CASE__								=      torch.zeros(
												    (self.model_tester.batch_size, self.model_tester.seq_length)    ,				dtype=torch.long    ,				device=__UpperCamelCase       )
												SCREAMING_SNAKE_CASE__								=      torch.zeros(
												    self.model_tester.batch_size    ,				dtype=torch.long    ,				device=__UpperCamelCase       )
						return inputs_dict

			def 					SCREAMING_SNAKE_CASE				(					self				:						Optional[int]       )   ->    Tuple:
						SCREAMING_SNAKE_CASE__								=      NezhaModelTester(self       )
						SCREAMING_SNAKE_CASE__								=      ConfigTester(self    ,				config_class=__UpperCamelCase    ,				hidden_size=3_7       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Any       )   ->    int:
						self.config_tester.run_common_tests()

			def 					SCREAMING_SNAKE_CASE				(					self				:						str       )   ->    Union[str, Any]:
						SCREAMING_SNAKE_CASE__								=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__UpperCamelCase       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Union[str, Any]       )   ->    Optional[int]:
						SCREAMING_SNAKE_CASE__								=      self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_model_as_decoder(*__UpperCamelCase       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						List[Any]       )   ->    List[str]:
						(
						    (
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		(
						    SCREAMING_SNAKE_CASE__
						)							,		
						)								=      self.model_tester.prepare_config_and_inputs_for_decoder()

						SCREAMING_SNAKE_CASE__								=      None

						self.model_tester.create_and_check_model_as_decoder(
						    __UpperCamelCase    ,				__UpperCamelCase    ,				__UpperCamelCase    ,				__UpperCamelCase    ,				__UpperCamelCase    ,				__UpperCamelCase    ,				__UpperCamelCase    ,				__UpperCamelCase    ,				__UpperCamelCase    ,				)

			def 					SCREAMING_SNAKE_CASE				(					self				:						List[str]       )   ->    List[str]:
						SCREAMING_SNAKE_CASE__								=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Dict       )   ->    str:
						SCREAMING_SNAKE_CASE__								=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*__UpperCamelCase       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Dict       )   ->    Dict:
						SCREAMING_SNAKE_CASE__								=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_next_sequence_prediction(*__UpperCamelCase       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Tuple       )   ->    Optional[Any]:
						SCREAMING_SNAKE_CASE__								=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*__UpperCamelCase       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Any       )   ->    List[Any]:
						SCREAMING_SNAKE_CASE__								=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*__UpperCamelCase       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						int       )   ->    Optional[Any]:
						SCREAMING_SNAKE_CASE__								=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*__UpperCamelCase       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Optional[int]       )   ->    Tuple:
						SCREAMING_SNAKE_CASE__								=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase       )

			@slow
			def 					SCREAMING_SNAKE_CASE				(					self				:						List[Any]       )   ->    Union[str, Any]:






						for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									SCREAMING_SNAKE_CASE__								=      NezhaModel.from_pretrained(__UpperCamelCase       )
									self.assertIsNotNone(__UpperCamelCase       )





			@slow
			@require_torch_gpu
			def 					SCREAMING_SNAKE_CASE				(					self				:						Dict       )   ->    Any:
						SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__								=      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

									SCREAMING_SNAKE_CASE__								=      True
									SCREAMING_SNAKE_CASE__								=      model_class(config=__UpperCamelCase       )

									SCREAMING_SNAKE_CASE__								=      self._prepare_for_class(__UpperCamelCase    ,				__UpperCamelCase       )
									SCREAMING_SNAKE_CASE__								=      torch.jit.trace(
									    __UpperCamelCase    ,				(inputs_dict["""input_ids"""].to("""cpu"""       ), inputs_dict["""attention_mask"""].to("""cpu"""       ))       )

									with tempfile.TemporaryDirectory() as tmp:
												torch.jit.save(__UpperCamelCase    ,				os.path.join(__UpperCamelCase    ,				"""bert.pt"""       )       )
												SCREAMING_SNAKE_CASE__								=      torch.jit.load(os.path.join(__UpperCamelCase    ,				"""bert.pt"""       )    ,				map_location=__UpperCamelCase       )
												loaded(inputs_dict["""input_ids"""].to(__UpperCamelCase       )    ,				inputs_dict["""attention_mask"""].to(__UpperCamelCase       )       )





@require_torch
class    __UpperCamelCase						(unittest.TestCase			):

			@slow
			def 					SCREAMING_SNAKE_CASE				(					self				:						Union[str, Any]       )   ->    Any:
						SCREAMING_SNAKE_CASE__								=      NezhaModel.from_pretrained("""sijunhe/nezha-cn-base"""       )
						SCREAMING_SNAKE_CASE__								=      torch.tensor([[0, 1, 2, 3, 4, 5]]       )
						SCREAMING_SNAKE_CASE__								=      torch.tensor([[0, 1, 1, 1, 1, 1]]       )
						with torch.no_grad():
									SCREAMING_SNAKE_CASE__								=      model(__UpperCamelCase    ,				attention_mask=__UpperCamelCase       )[0]
						SCREAMING_SNAKE_CASE__								=      torch.Size((1, 6, 7_6_8)       )
						self.assertEqual(output.shape    ,				__UpperCamelCase       )
						SCREAMING_SNAKE_CASE__								=      torch.tensor([[[0.0_685, 0.2_441, 0.1_102], [0.0_600, 0.1_906, 0.1_349], [0.0_221, 0.0_819, 0.0_586]]]       )

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





			@slow
			def 					SCREAMING_SNAKE_CASE				(					self				:						str       )   ->    Any:
						SCREAMING_SNAKE_CASE__								=      NezhaForMaskedLM.from_pretrained("""sijunhe/nezha-cn-base"""       )
						SCREAMING_SNAKE_CASE__								=      torch.tensor([[0, 1, 2, 3, 4, 5]]       )
						SCREAMING_SNAKE_CASE__								=      torch.tensor([[1, 1, 1, 1, 1, 1]]       )
						with torch.no_grad():
									SCREAMING_SNAKE_CASE__								=      model(__UpperCamelCase    ,				attention_mask=__UpperCamelCase       )[0]
						SCREAMING_SNAKE_CASE__								=      torch.Size((1, 6, 2_1_1_2_8)       )
						self.assertEqual(output.shape    ,				__UpperCamelCase       )
						SCREAMING_SNAKE_CASE__								=      torch.tensor(
						    [[-2.7_939, -1.7_902, -2.2_189], [-2.8_585, -1.8_908, -2.3_723], [-2.6_499, -1.7_750, -2.2_558]]       )

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





"""simple docstring"""




import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

from ...tokenization_utils import BatchEncoding, PreTrainedTokenizer
from ...utils import logging


A_							:			Union[str, Any]											= logging.get_logger(__name__)

A_							:			Tuple											= "▁"

A_							:			int											= {
    "vocab_file": "vocab.json",
    "spm_file": "sentencepiece.bpe.model",
    "tokenizer_config_file": "tokenizer_config.json",
}

A_							:			Dict											= {
    "vocab_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/vocab.json",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/vocab.json",
    },
    "spm_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/sentencepiece.bpe.model",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/sentencepiece.bpe.model",
    },
    "tokenizer_config_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/tokenizer_config.json",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/tokenizer_config.json",
    },
}

A_							:			Dict											= {
    "facebook/m2m100_418M": 1_024,
}

# fmt: off
A_							:			Any											= {
    "m2m100": ["af", "am", "ar", "ast", "az", "ba", "be", "bg", "bn", "br", "bs", "ca", "ceb", "cs", "cy", "da", "de", "el", "en", "es", "et", "fa", "ff", "fi", "fr", "fy", "ga", "gd", "gl", "gu", "ha", "he", "hi", "hr", "ht", "hu", "hy", "id", "ig", "ilo", "is", "it", "ja", "jv", "ka", "kk", "km", "kn", "ko", "lb", "lg", "ln", "lo", "lt", "lv", "mg", "mk", "ml", "mn", "mr", "ms", "my", "ne", "nl", "no", "ns", "oc", "or", "pa", "pl", "ps", "pt", "ro", "ru", "sd", "si", "sk", "sl", "so", "sq", "sr", "ss", "su", "sv", "sw", "ta", "th", "tl", "tn", "tr", "uk", "ur", "uz", "vi", "wo", "xh", "yi", "yo", "zh", "zu"],
    "wmt21": ["en", "ha", "is", "ja", "cs", "ru", "zh", "de"]
}





class    lowerCamelCase						(A__			):
			lowerCamelCase__   :		Optional[int]							=							VOCAB_FILES_NAMES
			lowerCamelCase__   :		Any							=							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			lowerCamelCase__   :		List[Any]							=							PRETRAINED_VOCAB_FILES_MAP
			lowerCamelCase__   :		Tuple							=							['input_ids', 'attention_mask']

			lowerCamelCase__   :		List[int]							=							[]
			lowerCamelCase__   :		List[int]							=							[]

			def __init__(					self				:						List[str]    ,				__UpperCAmelCase				:						int    ,				__UpperCAmelCase				:						Union[str, Any]    ,				__UpperCAmelCase				:						Dict=None    ,				__UpperCAmelCase				:						Optional[Any]=None    ,				__UpperCAmelCase				:						int="<s>"    ,				__UpperCAmelCase				:						Optional[int]="</s>"    ,				__UpperCAmelCase				:						Any="</s>"    ,				__UpperCAmelCase				:						Any="<pad>"    ,				__UpperCAmelCase				:						str="<unk>"    ,				__UpperCAmelCase				:						Dict="m2m100"    ,				__UpperCAmelCase				:						Optional[Dict[str, Any]] = None    ,				__UpperCAmelCase				:						str=8    ,				**__UpperCAmelCase				:						str    ,				)   ->    None:
						SCREAMING_SNAKE_CASE__								=      {} if sp_model_kwargs is None else sp_model_kwargs

						SCREAMING_SNAKE_CASE__								=      language_codes
						SCREAMING_SNAKE_CASE__								=      FAIRSEQ_LANGUAGE_CODES[language_codes]
						SCREAMING_SNAKE_CASE__								=      {lang_code: F"""__{lang_code}__""" for lang_code in fairseq_language_code}

						SCREAMING_SNAKE_CASE__								=      kwargs.get("""additional_special_tokens"""    ,				[]       )
						kwargs["additional_special_tokens"] += [
						    self.get_lang_token(__UpperCAmelCase       )
						    for lang_code in fairseq_language_code
						    if self.get_lang_token(__UpperCAmelCase       ) not in kwargs["additional_special_tokens"]
						]

						super().__init__(
						    src_lang=__UpperCAmelCase    ,				tgt_lang=__UpperCAmelCase    ,				bos_token=__UpperCAmelCase    ,				eos_token=__UpperCAmelCase    ,				sep_token=__UpperCAmelCase    ,				unk_token=__UpperCAmelCase    ,				pad_token=__UpperCAmelCase    ,				language_codes=__UpperCAmelCase    ,				sp_model_kwargs=self.sp_model_kwargs    ,				num_madeup_words=__UpperCAmelCase    ,				**__UpperCAmelCase    ,				)

						SCREAMING_SNAKE_CASE__								=      vocab_file
						SCREAMING_SNAKE_CASE__								=      load_json(__UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      {v: k for k, v in self.encoder.items()}
						SCREAMING_SNAKE_CASE__								=      spm_file
						SCREAMING_SNAKE_CASE__								=      load_spm(__UpperCAmelCase    ,				self.sp_model_kwargs       )

						SCREAMING_SNAKE_CASE__								=      len(self.encoder       )

						SCREAMING_SNAKE_CASE__								=      {
						    self.get_lang_token(__UpperCAmelCase       ): self.encoder_size + i for i, lang_code in enumerate(__UpperCAmelCase       )
						}
						SCREAMING_SNAKE_CASE__								=      {lang_code: self.encoder_size + i for i, lang_code in enumerate(__UpperCAmelCase       )}
						SCREAMING_SNAKE_CASE__								=      {v: k for k, v in self.lang_token_to_id.items()}

						SCREAMING_SNAKE_CASE__								=      src_lang if src_lang is not None else """en"""
						SCREAMING_SNAKE_CASE__								=      tgt_lang
						SCREAMING_SNAKE_CASE__								=      self.get_lang_id(self._src_lang       )
						self.set_src_lang_special_tokens(self._src_lang       )

						SCREAMING_SNAKE_CASE__								=      num_madeup_words

			@property
			def 					SCREAMING_SNAKE_CASE				(					self				:						int       )   ->    int:
						return len(self.encoder       ) + len(self.lang_token_to_id       )

			@property
			def 					SCREAMING_SNAKE_CASE				(					self				:						Any       )   ->    str:
						return self._src_lang

			@src_lang.setter
			def 					SCREAMING_SNAKE_CASE				(					self				:						str    ,				__UpperCAmelCase				:						str       )   ->    None:
						SCREAMING_SNAKE_CASE__								=      new_src_lang
						self.set_src_lang_special_tokens(self._src_lang       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						int    ,				__UpperCAmelCase				:						str       )   ->    List[str]:
						return self.sp_model.encode(__UpperCAmelCase    ,				out_type=__UpperCAmelCase       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						str    ,				__UpperCAmelCase				:						Tuple       )   ->    Tuple:
						if token in self.lang_token_to_id:
									return self.lang_token_to_id[token]
						return self.encoder.get(__UpperCAmelCase    ,				self.encoder[self.unk_token]       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Dict    ,				__UpperCAmelCase				:						int       )   ->    str:
						if index in self.id_to_lang_token:
									return self.id_to_lang_token[index]
						return self.decoder.get(__UpperCAmelCase    ,				self.unk_token       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Optional[Any]    ,				__UpperCAmelCase				:						Optional[Any]       )   ->    Dict:
						SCREAMING_SNAKE_CASE__								=      []
						SCREAMING_SNAKE_CASE__								=      """"""
						for token in tokens:
									# make sure that special tokens are not decoded using sentencepiece model
									if token in self.all_special_tokens:
												out_string += self.sp_model.decode(__UpperCAmelCase       ) + token
												SCREAMING_SNAKE_CASE__								=      []
									else:
												current_sub_tokens.append(__UpperCAmelCase       )
						out_string += self.sp_model.decode(__UpperCAmelCase       )
						return out_string.strip()

			def 					SCREAMING_SNAKE_CASE				(					self				:						int    ,				__UpperCAmelCase				:						List[int]    ,				__UpperCAmelCase				:						Optional[List[int]] = None    ,				__UpperCAmelCase				:						bool = False       )   ->    List[int]:

						if already_has_special_tokens:
									return super().get_special_tokens_mask(
									    token_ids_a=__UpperCAmelCase    ,				token_ids_a=__UpperCAmelCase    ,				already_has_special_tokens=__UpperCAmelCase       )

						SCREAMING_SNAKE_CASE__								=      [1] * len(self.prefix_tokens       )
						SCREAMING_SNAKE_CASE__								=      [1] * len(self.suffix_tokens       )
						if token_ids_a is None:
									return prefix_ones + ([0] * len(__UpperCAmelCase       )) + suffix_ones
						return prefix_ones + ([0] * len(__UpperCAmelCase       )) + ([0] * len(__UpperCAmelCase       )) + suffix_ones

			def 					SCREAMING_SNAKE_CASE				(					self				:						List[str]    ,				__UpperCAmelCase				:						List[int]    ,				__UpperCAmelCase				:						Optional[List[int]] = None       )   ->    List[int]:
						if token_ids_a is None:
									return self.prefix_tokens + token_ids_a + self.suffix_tokens
						# We don't expect to process pairs, but leave the pair logic for API consistency
						return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens

			def 					SCREAMING_SNAKE_CASE				(					self				:						Dict       )   ->    Dict:
						SCREAMING_SNAKE_CASE__								=      {self.convert_ids_to_tokens(__UpperCAmelCase       ): i for i in range(self.vocab_size       )}
						vocab.update(self.added_tokens_encoder       )
						return vocab

			def __getstate__(					self				:						Union[str, Any]       )   ->    Dict:
						SCREAMING_SNAKE_CASE__								=      self.__dict__.copy()
						SCREAMING_SNAKE_CASE__								=      None
						return state

			def __setstate__(					self				:						Union[str, Any]    ,				__UpperCAmelCase				:						Dict       )   ->    None:
						SCREAMING_SNAKE_CASE__								=      d

						# for backward compatibility
						if not hasattr(self    ,				"""sp_model_kwargs"""       ):
									SCREAMING_SNAKE_CASE__								=      {}

						SCREAMING_SNAKE_CASE__								=      load_spm(self.spm_file    ,				self.sp_model_kwargs       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Tuple    ,				__UpperCAmelCase				:						str    ,				__UpperCAmelCase				:						Optional[str] = None       )   ->    Tuple[str]:
						SCREAMING_SNAKE_CASE__								=      Path(__UpperCAmelCase       )
						if not save_dir.is_dir():
									raise OSError(F"""{save_directory} should be a directory"""       )
						SCREAMING_SNAKE_CASE__								=      save_dir / (
						    (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""vocab_file"""]
						)
						SCREAMING_SNAKE_CASE__								=      save_dir / (
						    (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""spm_file"""]
						)

						save_json(self.encoder    ,				__UpperCAmelCase       )

						if os.path.abspath(self.spm_file       ) != os.path.abspath(__UpperCAmelCase       ) and os.path.isfile(self.spm_file       ):
									copyfile(self.spm_file    ,				__UpperCAmelCase       )
						elif not os.path.isfile(self.spm_file       ):
									with open(__UpperCAmelCase    ,				"""wb"""       ) as fi:
												SCREAMING_SNAKE_CASE__								=      self.sp_model.serialized_model_proto()
												fi.write(__UpperCAmelCase       )

						return (str(__UpperCAmelCase       ), str(__UpperCAmelCase       ))

			def 					SCREAMING_SNAKE_CASE				(					self				:						Optional[Any]    ,				__UpperCAmelCase				:						List[str]    ,				__UpperCAmelCase				:						str = "en"    ,				__UpperCAmelCase				:						Optional[List[str]] = None    ,				__UpperCAmelCase				:						str = "ro"    ,				**__UpperCAmelCase				:						str    ,				)   ->    BatchEncoding:
						SCREAMING_SNAKE_CASE__								=      src_lang
						SCREAMING_SNAKE_CASE__								=      tgt_lang
						self.set_src_lang_special_tokens(self.src_lang       )
						return super().prepare_seqaseq_batch(__UpperCAmelCase    ,				__UpperCAmelCase    ,				**__UpperCAmelCase       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						Dict    ,				__UpperCAmelCase				:						int    ,				__UpperCAmelCase				:						Optional[str]    ,				__UpperCAmelCase				:						Optional[str]    ,				**__UpperCAmelCase				:						Tuple       )   ->    str:
						if src_lang is None or tgt_lang is None:
									raise ValueError("""Translation requires a `src_lang` and a `tgt_lang` for this model"""       )
						SCREAMING_SNAKE_CASE__								=      src_lang
						SCREAMING_SNAKE_CASE__								=      self(__UpperCAmelCase    ,				add_special_tokens=__UpperCAmelCase    ,				**__UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      self.get_lang_id(__UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      tgt_lang_id
						return inputs

			def 					SCREAMING_SNAKE_CASE				(					self				:						List[str]       )   ->    int:
						self.set_src_lang_special_tokens(self.src_lang       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						str       )   ->    str:
						self.set_tgt_lang_special_tokens(self.tgt_lang       )

			def 					SCREAMING_SNAKE_CASE				(					self				:						List[Any]    ,				__UpperCAmelCase				:						str       )   ->    None:
						SCREAMING_SNAKE_CASE__								=      self.get_lang_token(__UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      self.lang_token_to_id[lang_token]
						SCREAMING_SNAKE_CASE__								=      [self.cur_lang_id]
						SCREAMING_SNAKE_CASE__								=      [self.eos_token_id]

			def 					SCREAMING_SNAKE_CASE				(					self				:						Dict    ,				__UpperCAmelCase				:						str       )   ->    None:
						SCREAMING_SNAKE_CASE__								=      self.get_lang_token(__UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      self.lang_token_to_id[lang_token]
						SCREAMING_SNAKE_CASE__								=      [self.cur_lang_id]
						SCREAMING_SNAKE_CASE__								=      [self.eos_token_id]

			def 					SCREAMING_SNAKE_CASE				(					self				:						Dict    ,				__UpperCAmelCase				:						str       )   ->    str:
						return self.lang_code_to_token[lang]





			def 					SCREAMING_SNAKE_CASE				(					self				:						int    ,				__UpperCAmelCase				:						str       )   ->    int:
						SCREAMING_SNAKE_CASE__								=      self.get_lang_token(__UpperCAmelCase       )
						return self.lang_token_to_id[lang_token]

def 	A       (     snake_case__		,				snake_case__    ):




			'''simple docstring'''


			SCREAMING_SNAKE_CASE__								=      sentencepiece.SentencePieceProcessor(**snake_case__    )
			spm.Load(str(snake_case__    )    )
			return spm

def 	A       (     snake_case__    ):




			'''simple docstring'''


			with open(snake_case__		,				"""r"""    ) as f:
						return json.load(snake_case__    )




def 	A       (     snake_case__		,				snake_case__    ):




			'''simple docstring'''


			with open(snake_case__		,				"""w"""    ) as f:
						json.dump(snake_case__		,				snake_case__		,				indent=2    )


