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






    """simple docstring"""







    def __init__(			self    , _lowercase    , _lowercase=13    , _lowercase=7    , _lowercase=True    , _lowercase=True    , _lowercase=True    , _lowercase=True    , _lowercase=99    , _lowercase=32    , _lowercase=5    , _lowercase=4    , _lowercase=37    , _lowercase="gelu"    , _lowercase=0.1    , _lowercase=0.1    , _lowercase=128    , _lowercase=32    , _lowercase=16    , _lowercase=2    , _lowercase=0.02    , _lowercase=3    , _lowercase=4    , _lowercase=None    , )	->	Optional[Any]:
        _lowerCamelCase	: Any       =       parent
        _lowerCamelCase	: List[Any]       =       batch_size
        _lowerCamelCase	: List[str]       =       seq_length
        _lowerCamelCase	: int       =       is_training
        _lowerCamelCase	: Optional[Any]       =       use_input_mask
        _lowerCamelCase	: int       =       use_token_type_ids
        _lowerCamelCase	: Optional[Any]       =       use_labels
        _lowerCamelCase	: Optional[Any]       =       vocab_size
        _lowerCamelCase	: int       =       hidden_size
        _lowerCamelCase	: Optional[int]       =       num_hidden_layers
        _lowerCamelCase	: Any       =       num_attention_heads
        _lowerCamelCase	: Optional[Any]       =       intermediate_size
        _lowerCamelCase	: Any       =       hidden_act
        _lowerCamelCase	: Optional[int]       =       hidden_dropout_prob
        _lowerCamelCase	: str       =       attention_probs_dropout_prob
        _lowerCamelCase	: Any       =       max_position_embeddings
        _lowerCamelCase	: Optional[Any]       =       type_vocab_size
        _lowerCamelCase	: Optional[Any]       =       type_sequence_label_size
        _lowerCamelCase	: str       =       initializer_range
        _lowerCamelCase	: List[Any]       =       num_labels
        _lowerCamelCase	: Union[str, Any]       =       num_choices
        _lowerCamelCase	: List[str]       =       scope







    def 				a__					(			self					)	->	List[str]:
        _lowerCamelCase	: Dict       =       ids_tensor([self.batch_size, self.seq_length]    , self.vocab_size					)

        _lowerCamelCase	: Optional[int]       =       None
        if self.use_input_mask:
            _lowerCamelCase	: List[Any]       =       random_attention_mask([self.batch_size, self.seq_length]					)

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

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

        _lowerCamelCase	: Optional[int]       =       self.get_config()

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







    def 				a__					(			self					)	->	Union[str, 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 				a__					(			self					)	->	Any:
        (
            (
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),							
        )	: Dict       =       self.prepare_config_and_inputs()

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







    def 				a__					(			self    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase					)	->	List[str]:
        _lowerCamelCase	: List[str]       =       NezhaForPreTraining(config=A_					)
        model.to(A_					)
        model.eval()
        _lowerCamelCase	: Tuple       =       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 				a__					(			self    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase					)	->	Optional[int]:
        _lowerCamelCase	: Optional[int]       =       NezhaForQuestionAnswering(config=A_					)
        model.to(A_					)
        model.eval()
        _lowerCamelCase	: List[Any]       =       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 				a__					(			self    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase					)	->	Union[str, Any]:
        _lowerCamelCase	: Any       =       self.num_labels
        _lowerCamelCase	: List[str]       =       NezhaForSequenceClassification(A_					)
        model.to(A_					)
        model.eval()
        _lowerCamelCase	: 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 				a__					(			self    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase					)	->	List[str]:
        _lowerCamelCase	: Optional[Any]       =       self.num_labels
        _lowerCamelCase	: Union[str, Any]       =       NezhaForTokenClassification(config=A_					)
        model.to(A_					)
        model.eval()
        _lowerCamelCase	: 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 				a__					(			self    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase    , _lowercase					)	->	str:
        _lowerCamelCase	: List[str]       =       self.num_choices
        _lowerCamelCase	: List[str]       =       NezhaForMultipleChoice(config=A_					)
        model.to(A_					)
        model.eval()
        _lowerCamelCase	: str       =       input_ids.unsqueeze(1					).expand(-1    , self.num_choices    , -1					).contiguous()
        _lowerCamelCase	: List[Any]       =       token_type_ids.unsqueeze(1					).expand(-1    , self.num_choices    , -1					).contiguous()
        _lowerCamelCase	: Optional[Any]       =       input_mask.unsqueeze(1					).expand(-1    , self.num_choices    , -1					).contiguous()
        _lowerCamelCase	: Any       =       model(
            A_    , attention_mask=A_    , token_type_ids=A_    , labels=A_    , )
        self.parent.assertEqual(result.logits.shape    , (self.batch_size, self.num_choices)					)







    def 				a__					(			self					)	->	List[str]:
        _lowerCamelCase	: int       =       self.prepare_config_and_inputs()
        (
            (
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),							
        )	: int       =       config_and_inputs
        _lowerCamelCase	: Tuple       =       {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
        return config, inputs_dict





@require_torch
class      _UpperCAmelCase						( __SCREAMING_SNAKE_CASE ,			__SCREAMING_SNAKE_CASE ,			__SCREAMING_SNAKE_CASE ,			unittest.TestCase						):






    """simple docstring"""
    __snake_case    =       (
        (
            NezhaModel,
            NezhaForMaskedLM,
            NezhaForMultipleChoice,
            NezhaForNextSentencePrediction,
            NezhaForPreTraining,
            NezhaForQuestionAnswering,
            NezhaForSequenceClassification,
            NezhaForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    __snake_case    =       (
        {
            """feature-extraction""": NezhaModel,
            """fill-mask""": NezhaForMaskedLM,
            """question-answering""": NezhaForQuestionAnswering,
            """text-classification""": NezhaForSequenceClassification,
            """token-classification""": NezhaForTokenClassification,
            """zero-shot""": NezhaForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    __snake_case    =       True







    def 				a__					(			self    , _lowercase    , _lowercase    , _lowercase=False					)	->	List[Any]:
        _lowerCamelCase	: Tuple       =       super()._prepare_for_class(A_    , A_    , return_labels=A_					)

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







    def 				a__					(			self					)	->	Any:
        _lowerCamelCase	: Dict       =       NezhaModelTester(self					)
        _lowerCamelCase	: int       =       ConfigTester(self    , config_class=A_    , hidden_size=37					)







    def 				a__					(			self					)	->	Optional[Any]:
        self.config_tester.run_common_tests()







    def 				a__					(			self					)	->	Optional[Any]:
        _lowerCamelCase	: Optional[int]       =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*A_					)







    def 				a__					(			self					)	->	Tuple:
        _lowerCamelCase	: Tuple       =       self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(*A_					)







    def 				a__					(			self					)	->	List[str]:
        (
            (
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),					(
            _lowerCamelCase		
        ),							
        )	: Union[str, Any]       =       self.model_tester.prepare_config_and_inputs_for_decoder()

        _lowerCamelCase	: Union[str, Any]       =       None

        self.model_tester.create_and_check_model_as_decoder(
            A_    , A_    , A_    , A_    , A_    , A_    , A_    , A_    , A_    , )







    def 				a__					(			self					)	->	str:
        _lowerCamelCase	: Dict       =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*A_					)







    def 				a__					(			self					)	->	Dict:
        _lowerCamelCase	: Tuple       =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*A_					)







    def 				a__					(			self					)	->	int:
        _lowerCamelCase	: Optional[Any]       =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_next_sequence_prediction(*A_					)







    def 				a__					(			self					)	->	str:
        _lowerCamelCase	: int       =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*A_					)







    def 				a__					(			self					)	->	Dict:
        _lowerCamelCase	: Any       =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*A_					)







    def 				a__					(			self					)	->	Union[str, Any]:
        _lowerCamelCase	: Dict       =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*A_					)







    def 				a__					(			self					)	->	Dict:
        _lowerCamelCase	: Optional[Any]       =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*A_					)







    @slow
    def 				a__					(			self					)	->	str:

        for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            _lowerCamelCase	: str       =       NezhaModel.from_pretrained(A_					)
            self.assertIsNotNone(A_					)







    @slow
    @require_torch_gpu
    def 				a__					(			self					)	->	List[str]:
        _lowerCamelCase,					_lowerCamelCase	: Optional[int]       =       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

            _lowerCamelCase	: List[Any]       =       True
            _lowerCamelCase	: int       =       model_class(config=A_					)

            _lowerCamelCase	: int       =       self._prepare_for_class(A_    , A_					)
            _lowerCamelCase	: str       =       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'''					)					)
                _lowerCamelCase	: List[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						):






    """simple docstring"""







    @slow
    def 				a__					(			self					)	->	Optional[Any]:
        _lowerCamelCase	: Optional[int]       =       NezhaModel.from_pretrained('''sijunhe/nezha-cn-base'''					)
        _lowerCamelCase	: List[str]       =       torch.tensor([[0, 1, 2, 3, 4, 5]]					)
        _lowerCamelCase	: Optional[int]       =       torch.tensor([[0, 1, 1, 1, 1, 1]]					)
        with torch.no_grad():
            _lowerCamelCase	: Union[str, Any]       =       model(A_    , attention_mask=A_					)[0]
        _lowerCamelCase	: List[Any]       =       torch.Size((1, 6, 768)					)
        self.assertEqual(output.shape    , A_					)
        _lowerCamelCase	: Any       =       torch.tensor([[[0.0685, 0.2441, 0.1102], [0.0600, 0.1906, 0.1349], [0.0221, 0.0819, 0.0586]]]					)

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







    @slow
    def 				a__					(			self					)	->	int:
        _lowerCamelCase	: Optional[Any]       =       NezhaForMaskedLM.from_pretrained('''sijunhe/nezha-cn-base'''					)
        _lowerCamelCase	: Optional[Any]       =       torch.tensor([[0, 1, 2, 3, 4, 5]]					)
        _lowerCamelCase	: List[str]       =       torch.tensor([[1, 1, 1, 1, 1, 1]]					)
        with torch.no_grad():
            _lowerCamelCase	: Union[str, Any]       =       model(A_    , attention_mask=A_					)[0]
        _lowerCamelCase	: List[str]       =       torch.Size((1, 6, 21128)					)
        self.assertEqual(output.shape    , A_					)
        _lowerCamelCase	: Dict       =       torch.tensor(
            [[-2.7939, -1.7902, -2.2189], [-2.8585, -1.8908, -2.3723], [-2.6499, -1.7750, -2.2558]]					)

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









from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Value
from .base import TaskTemplate



@dataclass(frozen=__SCREAMING_SNAKE_CASE    )
class 					__snake_case	(  __SCREAMING_SNAKE_CASE    ):


    '''simple docstring'''







    lowerCamelCase__    :		str				=	field(default="""summarization"""			,     metadata={"""include_in_asdict_even_if_is_default""": True}    )
    lowerCamelCase__    :		ClassVar[Features]				=	Features({"""text""": Value("""string"""    )}    )
    lowerCamelCase__    :		ClassVar[Features]				=	Features({"""summary""": Value("""string"""    )}    )
    lowerCamelCase__    :		str				=	"text"
    lowerCamelCase__    :		str				=	"summary"



    @property
    def  lowercase_					(       self     ):





        '''simple docstring'''

        return {self.text_column: "text", self.summary_column: "summary"}
