



import PIL.Image
import PIL.ImageOps
from packaging import version
from PIL import Image


if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("""9.1.0"""):
  a__			=		{
      "linear": PIL.Image.Resampling.BILINEAR,
      "bilinear": PIL.Image.Resampling.BILINEAR,
      "bicubic": PIL.Image.Resampling.BICUBIC,
      "lanczos": PIL.Image.Resampling.LANCZOS,
      "nearest": PIL.Image.Resampling.NEAREST,
  }
else:
  a__			=		{
      "linear": PIL.Image.LINEAR,
      "bilinear": PIL.Image.BILINEAR,
      "bicubic": PIL.Image.BICUBIC,
      "lanczos": PIL.Image.LANCZOS,
      "nearest": PIL.Image.NEAREST,
  }



def 					lowercase				(      SCREAMING_SNAKE_CASE__				:		Dict )     ->   List[str]:
 _snake_case			:     Any 		=    (images / 2 + 0.5).clamp(0 ,				1 )
 _snake_case			:     Union[str, Any] 		=    images.cpu().permute(0 ,				2 ,				3 ,				1 ).float().numpy()
 _snake_case			:     List[Any] 		=    numpy_to_pil(__lowerCamelCase )
 return images



def 					lowercase				(      SCREAMING_SNAKE_CASE__				:		Optional[int] )     ->   str:
 if images.ndim == 3:
  _snake_case			:     Tuple 		=    images[None, ...]
 _snake_case			:     Optional[Any] 		=    (images * 255).round().astype("""uint8""" )
 if images.shape[-1] == 1:
  # special case for grayscale (single channel) images
  _snake_case			:     Tuple 		=    [Image.fromarray(image.squeeze() ,				mode="""L""" ) for image in images]
 else:
  _snake_case			:     Any 		=    [Image.fromarray(__lowerCamelCase ) for image in images]

 return pil_images








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



      """simple docstring"""


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

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

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

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

        __UpperCAmelCase							:	str           =  self.get_config()

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
      def 	UpperCAmelCase      (      self					:			str					)							->      str:
        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=__lowercase		,      initializer_range=self.initializer_range		,      )
      def 	UpperCAmelCase      (      self					:			int					)							->      Optional[Any]:
        (
            (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , 
        )							:	List[str]           =  self.prepare_config_and_inputs()

        __UpperCAmelCase							:	str           =  True
        __UpperCAmelCase							:	Union[str, Any]           =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
        __UpperCAmelCase							:	Tuple           =  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 	UpperCAmelCase      (      self					:			str		,      __lowercase					:			Union[str, Any]		,      __lowercase					:			str		,      __lowercase					:			Optional[int]		,      __lowercase					:			Dict		,      __lowercase					:			Union[str, Any]		,      __lowercase					:			str		,      __lowercase					:			int					)							->      Any:
        __UpperCAmelCase							:	Union[str, Any]           =  NezhaModel(config=__lowercase					)
        model.to(__lowercase					)
        model.eval()
        __UpperCAmelCase							:	int           =  model(__lowercase		,      attention_mask=__lowercase		,      token_type_ids=__lowercase					)
        __UpperCAmelCase							:	Optional[Any]           =  model(__lowercase		,      token_type_ids=__lowercase					)
        __UpperCAmelCase							:	List[Any]           =  model(__lowercase					)
        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 	UpperCAmelCase      (      self					:			Optional[int]		,      __lowercase					:			Union[str, Any]		,      __lowercase					:			Optional[int]		,      __lowercase					:			Dict		,      __lowercase					:			Optional[int]		,      __lowercase					:			Any		,      __lowercase					:			str		,      __lowercase					:			Tuple		,      __lowercase					:			Optional[Any]		,      __lowercase					:			Tuple		,      )							->      Optional[int]:
        __UpperCAmelCase							:	List[Any]           =  True
        __UpperCAmelCase							:	Optional[Any]           =  NezhaModel(__lowercase					)
        model.to(__lowercase					)
        model.eval()
        __UpperCAmelCase							:	List[Any]           =  model(
            __lowercase		,      attention_mask=__lowercase		,      token_type_ids=__lowercase		,      encoder_hidden_states=__lowercase		,      encoder_attention_mask=__lowercase		,      )
        __UpperCAmelCase							:	Optional[Any]           =  model(
            __lowercase		,      attention_mask=__lowercase		,      token_type_ids=__lowercase		,      encoder_hidden_states=__lowercase		,      )
        __UpperCAmelCase							:	List[str]           =  model(__lowercase		,      attention_mask=__lowercase		,      token_type_ids=__lowercase					)
        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 	UpperCAmelCase      (      self					:			Any		,      __lowercase					:			int		,      __lowercase					:			str		,      __lowercase					:			Dict		,      __lowercase					:			Union[str, Any]		,      __lowercase					:			Tuple		,      __lowercase					:			Optional[int]		,      __lowercase					:			Optional[int]					)							->      Optional[int]:
        __UpperCAmelCase							:	Optional[int]           =  NezhaForMaskedLM(config=__lowercase					)
        model.to(__lowercase					)
        model.eval()
        __UpperCAmelCase							:	Optional[int]           =  model(__lowercase		,      attention_mask=__lowercase		,      token_type_ids=__lowercase		,      labels=__lowercase					)
        self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.vocab_size)					)
      def 	UpperCAmelCase      (      self					:			Dict		,      __lowercase					:			str		,      __lowercase					:			Optional[Any]		,      __lowercase					:			int		,      __lowercase					:			List[str]		,      __lowercase					:			Dict		,      __lowercase					:			str		,      __lowercase					:			str					)							->      Dict:
        __UpperCAmelCase							:	Optional[int]           =  NezhaForNextSentencePrediction(config=__lowercase					)
        model.to(__lowercase					)
        model.eval()
        __UpperCAmelCase							:	Union[str, Any]           =  model(
            __lowercase		,      attention_mask=__lowercase		,      token_type_ids=__lowercase		,      labels=__lowercase		,      )
        self.parent.assertEqual(result.logits.shape		,      (self.batch_size, 2)					)
      def 	UpperCAmelCase      (      self					:			Tuple		,      __lowercase					:			Any		,      __lowercase					:			Dict		,      __lowercase					:			Optional[int]		,      __lowercase					:			Tuple		,      __lowercase					:			Optional[int]		,      __lowercase					:			Optional[int]		,      __lowercase					:			List[str]					)							->      int:
        __UpperCAmelCase							:	Optional[Any]           =  NezhaForPreTraining(config=__lowercase					)
        model.to(__lowercase					)
        model.eval()
        __UpperCAmelCase							:	int           =  model(
            __lowercase		,      attention_mask=__lowercase		,      token_type_ids=__lowercase		,      labels=__lowercase		,      next_sentence_label=__lowercase		,      )
        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 	UpperCAmelCase      (      self					:			Tuple		,      __lowercase					:			List[str]		,      __lowercase					:			Any		,      __lowercase					:			Tuple		,      __lowercase					:			Optional[Any]		,      __lowercase					:			Dict		,      __lowercase					:			List[Any]		,      __lowercase					:			Dict					)							->      List[Any]:
        __UpperCAmelCase							:	Optional[Any]           =  NezhaForQuestionAnswering(config=__lowercase					)
        model.to(__lowercase					)
        model.eval()
        __UpperCAmelCase							:	Any           =  model(
            __lowercase		,      attention_mask=__lowercase		,      token_type_ids=__lowercase		,      start_positions=__lowercase		,      end_positions=__lowercase		,      )
        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 	UpperCAmelCase      (      self					:			Tuple		,      __lowercase					:			Dict		,      __lowercase					:			Any		,      __lowercase					:			Any		,      __lowercase					:			Union[str, Any]		,      __lowercase					:			int		,      __lowercase					:			Dict		,      __lowercase					:			List[Any]					)							->      Optional[Any]:
        __UpperCAmelCase							:	Optional[int]           =  self.num_labels
        __UpperCAmelCase							:	Union[str, Any]           =  NezhaForSequenceClassification(__lowercase					)
        model.to(__lowercase					)
        model.eval()
        __UpperCAmelCase							:	Optional[Any]           =  model(__lowercase		,      attention_mask=__lowercase		,      token_type_ids=__lowercase		,      labels=__lowercase					)
        self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_labels)					)
      def 	UpperCAmelCase      (      self					:			Tuple		,      __lowercase					:			List[Any]		,      __lowercase					:			Union[str, Any]		,      __lowercase					:			str		,      __lowercase					:			List[Any]		,      __lowercase					:			List[str]		,      __lowercase					:			Dict		,      __lowercase					:			Dict					)							->      str:
        __UpperCAmelCase							:	Union[str, Any]           =  self.num_labels
        __UpperCAmelCase							:	Dict           =  NezhaForTokenClassification(config=__lowercase					)
        model.to(__lowercase					)
        model.eval()
        __UpperCAmelCase							:	Dict           =  model(__lowercase		,      attention_mask=__lowercase		,      token_type_ids=__lowercase		,      labels=__lowercase					)
        self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.num_labels)					)
      def 	UpperCAmelCase      (      self					:			Tuple		,      __lowercase					:			Any		,      __lowercase					:			List[Any]		,      __lowercase					:			Dict		,      __lowercase					:			Tuple		,      __lowercase					:			List[Any]		,      __lowercase					:			Any		,      __lowercase					:			int					)							->      Optional[int]:
        __UpperCAmelCase							:	List[str]           =  self.num_choices
        __UpperCAmelCase							:	int           =  NezhaForMultipleChoice(config=__lowercase					)
        model.to(__lowercase					)
        model.eval()
        __UpperCAmelCase							:	List[str]           =  input_ids.unsqueeze(1					).expand(-1		,      self.num_choices		,      -1					).contiguous()
        __UpperCAmelCase							:	Tuple           =  token_type_ids.unsqueeze(1					).expand(-1		,      self.num_choices		,      -1					).contiguous()
        __UpperCAmelCase							:	int           =  input_mask.unsqueeze(1					).expand(-1		,      self.num_choices		,      -1					).contiguous()
        __UpperCAmelCase							:	Tuple           =  model(
            __lowercase		,      attention_mask=__lowercase		,      token_type_ids=__lowercase		,      labels=__lowercase		,      )
        self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_choices)					)




      def 	UpperCAmelCase      (      self					:			Dict					)							->      List[Any]:
        __UpperCAmelCase							:	Optional[int]           =  self.prepare_config_and_inputs()
        (
            (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , 
        )							:	Tuple           =  config_and_inputs
        __UpperCAmelCase							:	Optional[Any]           =  {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
        return config, inputs_dict


@require_torch
class     a    ( lowercase__							,       lowercase__							,       lowercase__							,       unittest.TestCase   ):



      """simple docstring"""


      a  : Union[str, Any]							       =	(
          (
              NezhaModel,
              NezhaForMaskedLM,
              NezhaForMultipleChoice,
              NezhaForNextSentencePrediction,
              NezhaForPreTraining,
              NezhaForQuestionAnswering,
              NezhaForSequenceClassification,
              NezhaForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      a  : Tuple							       =	(
          {
              'feature-extraction': NezhaModel,
              'fill-mask': NezhaForMaskedLM,
              'question-answering': NezhaForQuestionAnswering,
              'text-classification': NezhaForSequenceClassification,
              'token-classification': NezhaForTokenClassification,
              'zero-shot': NezhaForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      a  : Dict							       =	True
      def 	UpperCAmelCase      (      self					:			Optional[int]		,      __lowercase					:			str		,      __lowercase					:			Dict		,      __lowercase					:			int=False					)							->      Dict:
        __UpperCAmelCase							:	Optional[Any]           =  super()._prepare_for_class(__lowercase		,      __lowercase		,      return_labels=__lowercase					)

        if return_labels:
          if model_class in get_values(__lowercase					):
            __UpperCAmelCase							:	int           =  torch.zeros(
                (self.model_tester.batch_size, self.model_tester.seq_length)		,      dtype=torch.long		,      device=__lowercase					)
            __UpperCAmelCase							:	int           =  torch.zeros(
                self.model_tester.batch_size		,      dtype=torch.long		,      device=__lowercase					)
        return inputs_dict
      def 	UpperCAmelCase      (      self					:			Any					)							->      int:
        __UpperCAmelCase							:	Tuple           =  NezhaModelTester(self					)
        __UpperCAmelCase							:	List[Any]           =  ConfigTester(self		,      config_class=__lowercase		,      hidden_size=37					)
      def 	UpperCAmelCase      (      self					:			Dict					)							->      List[str]:
        self.config_tester.run_common_tests()
      def 	UpperCAmelCase      (      self					:			str					)							->      List[Any]:
        __UpperCAmelCase							:	Optional[int]           =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowercase					)
      def 	UpperCAmelCase      (      self					:			Optional[Any]					)							->      Tuple:
        __UpperCAmelCase							:	Optional[Any]           =  self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(*__lowercase					)
      def 	UpperCAmelCase      (      self					:			int					)							->      str:
        # This regression test was failing with PyTorch < 1.3
        (
            (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , (
            __UpperCAmelCase
        )      , 
        )							:	Tuple           =  self.model_tester.prepare_config_and_inputs_for_decoder()

        __UpperCAmelCase							:	int           =  None

        self.model_tester.create_and_check_model_as_decoder(
            __lowercase		,      __lowercase		,      __lowercase		,      __lowercase		,      __lowercase		,      __lowercase		,      __lowercase		,      __lowercase		,      __lowercase		,      )
      def 	UpperCAmelCase      (      self					:			Dict					)							->      Union[str, Any]:
        __UpperCAmelCase							:	List[Any]           =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowercase					)
      def 	UpperCAmelCase      (      self					:			int					)							->      Dict:
        __UpperCAmelCase							:	Optional[Any]           =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*__lowercase					)
      def 	UpperCAmelCase      (      self					:			List[str]					)							->      List[Any]:
        __UpperCAmelCase							:	int           =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_next_sequence_prediction(*__lowercase					)
      def 	UpperCAmelCase      (      self					:			Optional[int]					)							->      Optional[int]:
        __UpperCAmelCase							:	Optional[int]           =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*__lowercase					)
      def 	UpperCAmelCase      (      self					:			int					)							->      Dict:
        __UpperCAmelCase							:	Dict           =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*__lowercase					)
      def 	UpperCAmelCase      (      self					:			List[str]					)							->      str:
        __UpperCAmelCase							:	Union[str, Any]           =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*__lowercase					)
      def 	UpperCAmelCase      (      self					:			Optional[Any]					)							->      str:
        __UpperCAmelCase							:	int           =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*__lowercase					)
      @slow
      def 	UpperCAmelCase      (      self					:			str					)							->      Any:
        for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          __UpperCAmelCase							:	int           =  NezhaModel.from_pretrained(__lowercase					)
          self.assertIsNotNone(__lowercase					)




      @slow
      @require_torch_gpu
      def 	UpperCAmelCase      (      self					:			Union[str, Any]					)							->      List[Any]:
        __UpperCAmelCase      , __UpperCAmelCase							:	Optional[Any]           =  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							:	Optional[Any]           =  True
          __UpperCAmelCase							:	Dict           =  model_class(config=__lowercase					)

          __UpperCAmelCase							:	Union[str, Any]           =  self._prepare_for_class(__lowercase		,      __lowercase					)
          __UpperCAmelCase							:	Union[str, Any]           =  torch.jit.trace(
              __lowercase		,      (inputs_dict["""input_ids"""].to("""cpu"""					), inputs_dict["""attention_mask"""].to("""cpu"""					))					)

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


@require_torch
class     a    ( unittest.TestCase   ):



      """simple docstring"""


      @slow
      def 	UpperCAmelCase      (      self					:			Any					)							->      Optional[Any]:
        __UpperCAmelCase							:	Tuple           =  NezhaModel.from_pretrained("""sijunhe/nezha-cn-base"""					)
        __UpperCAmelCase							:	List[Any]           =  torch.tensor([[0, 1, 2, 3, 4, 5]]					)
        __UpperCAmelCase							:	Any           =  torch.tensor([[0, 1, 1, 1, 1, 1]]					)
        with torch.no_grad():
          __UpperCAmelCase							:	Optional[Any]           =  model(__lowercase		,      attention_mask=__lowercase					)[0]
        __UpperCAmelCase							:	Dict           =  torch.Size((1, 6, 768)					)
        self.assertEqual(output.shape		,      __lowercase					)
        __UpperCAmelCase							:	Any           =  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]		,      __lowercase		,      atol=1e-4					)					)




      @slow
      def 	UpperCAmelCase      (      self					:			str					)							->      List[str]:
        __UpperCAmelCase							:	int           =  NezhaForMaskedLM.from_pretrained("""sijunhe/nezha-cn-base"""					)
        __UpperCAmelCase							:	Tuple           =  torch.tensor([[0, 1, 2, 3, 4, 5]]					)
        __UpperCAmelCase							:	Any           =  torch.tensor([[1, 1, 1, 1, 1, 1]]					)
        with torch.no_grad():
          __UpperCAmelCase							:	Dict           =  model(__lowercase		,      attention_mask=__lowercase					)[0]
        __UpperCAmelCase							:	Union[str, Any]           =  torch.Size((1, 6, 21128)					)
        self.assertEqual(output.shape		,      __lowercase					)
        __UpperCAmelCase							:	int           =  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]		,      __lowercase		,      atol=1e-4					)					)


