
from manim import *







class        lowerCAmelCase__      (   a__	):

      def  A_				(			self				)					->      Union[str, Any]:


        '''simple docstring'''

        _UpperCamelCase						    =	Rectangle(height=0.5       ,							width=0.5				)
        _UpperCamelCase						    =	Rectangle(height=0.46       ,							width=0.46				).set_stroke(width=0				)

        _UpperCamelCase						    =	[mem.copy() for i in range(6				)]
        _UpperCamelCase						    =	[mem.copy() for i in range(6				)]
        _UpperCamelCase						    =	VGroup(*lowerCamelCase_				).arrange(lowerCamelCase_       ,							buff=0				)
        _UpperCamelCase						    =	VGroup(*lowerCamelCase_				).arrange(lowerCamelCase_       ,							buff=0				)
        _UpperCamelCase						    =	VGroup(lowerCamelCase_       ,							lowerCamelCase_				).arrange(lowerCamelCase_       ,							buff=0				)
        _UpperCamelCase						    =	Text("""CPU"""       ,							font_size=24				)
        _UpperCamelCase						    =	Group(lowerCamelCase_       ,							lowerCamelCase_				).arrange(lowerCamelCase_       ,							buff=0.5       ,							aligned_edge=lowerCamelCase_				)
        cpu.move_to([-2.5, -0.5, 0]				)
        self.add(lowerCamelCase_				)

        _UpperCamelCase						    =	[mem.copy() for i in range(4				)]
        _UpperCamelCase						    =	VGroup(*lowerCamelCase_				).arrange(lowerCamelCase_       ,							buff=0				)
        _UpperCamelCase						    =	Text("""GPU"""       ,							font_size=24				)
        _UpperCamelCase						    =	Group(lowerCamelCase_       ,							lowerCamelCase_				).arrange(lowerCamelCase_       ,							buff=0.5       ,							aligned_edge=lowerCamelCase_				)
        gpu.move_to([-1, -1, 0]				)
        self.add(lowerCamelCase_				)

        _UpperCamelCase						    =	[mem.copy() for i in range(6				)]
        _UpperCamelCase						    =	VGroup(*lowerCamelCase_				).arrange(lowerCamelCase_       ,							buff=0				)

        _UpperCamelCase						    =	Text("""Model"""       ,							font_size=24				)
        _UpperCamelCase						    =	Group(lowerCamelCase_       ,							lowerCamelCase_				).arrange(lowerCamelCase_       ,							buff=0.5       ,							aligned_edge=lowerCamelCase_				)
        model.move_to([3, -1.0, 0]				)
        self.add(lowerCamelCase_				)

        _UpperCamelCase						    =	[]
        for i, rect in enumerate(lowerCamelCase_				):
          rect.set_stroke(lowerCamelCase_				)
          # target = fill.copy().set_fill(YELLOW, opacity=0.7)
          # target.move_to(rect)
          # self.add(target)

          _UpperCamelCase						    =	Rectangle(height=0.46 / 4       ,							width=0.46 / 3				).set_stroke(width=0.0				).set_fill(lowerCamelCase_       ,							opacity=0.7				)

          if i == 0:
            cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT				)       ,							buff=0.02       ,							direction=lowerCamelCase_				)
            cpu_target.set_x(cpu_target.get_x() + 0.1				)
          elif i == 3:
            cpu_target.next_to(cpu_targs[0]       ,							direction=lowerCamelCase_       ,							buff=0.0				)
          else:
            cpu_target.next_to(cpu_targs[i - 1]       ,							direction=lowerCamelCase_       ,							buff=0.0				)
          self.add(lowerCamelCase_				)
          cpu_targs.append(lowerCamelCase_				)

        _UpperCamelCase						    =	[mem.copy() for i in range(6				)]
        _UpperCamelCase						    =	VGroup(*lowerCamelCase_				).arrange(lowerCamelCase_       ,							buff=0				)

        _UpperCamelCase						    =	Text("""Loaded Checkpoint"""       ,							font_size=24				)
        _UpperCamelCase						    =	Group(lowerCamelCase_       ,							lowerCamelCase_				).arrange(lowerCamelCase_       ,							aligned_edge=lowerCamelCase_       ,							buff=0.4				)
        checkpoint.move_to([3, 0.5, 0]				)

        _UpperCamelCase						    =	Square(side_length=2.2				)
        key.move_to([-5, 2, 0]				)

        _UpperCamelCase						    =	MarkupText(
            F'<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'       ,							font_size=18       ,							)

        key_text.move_to([-5, 2.4, 0]				)

        self.add(lowerCamelCase_       ,							lowerCamelCase_				)

        _UpperCamelCase						    =	MarkupText(
            F'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'       ,							font_size=18       ,							)

        blue_text.next_to(lowerCamelCase_       ,							DOWN * 2.4       ,							aligned_edge=key_text.get_left()				)

        _UpperCamelCase						    =	MarkupText(
            F'Next, a <i><span fgcolor=\"{BLUE}\">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor=\"{BLUE}\">single shard</span>.'       ,							font_size=24       ,							)
        step_a.move_to([2, 2, 0]				)
        self.play(Write(lowerCamelCase_				)       ,							Write(lowerCamelCase_				)				)

        self.play(Write(lowerCamelCase_       ,							run_time=1				)       ,							Create(lowerCamelCase_       ,							run_time=1				)				)

        _UpperCamelCase						    =	[]
        _UpperCamelCase						    =	[]
        for i, rect in enumerate(lowerCamelCase_				):
          _UpperCamelCase						    =	fill.copy().set_fill(lowerCamelCase_       ,							opacity=0.7				)
          target.move_to(lowerCamelCase_				)
          first_animations.append(GrowFromCenter(lowerCamelCase_       ,							run_time=1				)				)

          _UpperCamelCase						    =	target.copy()
          cpu_target.generate_target()
          if i < 5:
            cpu_target.target.move_to(cpu_left_col_base[i + 1]				)
          else:
            cpu_target.target.move_to(cpu_right_col_base[i - 5]				)
          second_animations.append(MoveToTarget(lowerCamelCase_       ,							run_time=1.5				)				)

        self.play(*lowerCamelCase_				)
        self.play(*lowerCamelCase_				)
        self.wait()





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

						def __init__(			self       ,							a       ,							a=13       ,							a=7       ,							a=True       ,							a=True       ,							a=True       ,							a=True       ,							a=99       ,							a=32       ,							a=5       ,							a=4       ,							a=37       ,							a="gelu"       ,							a=0.1       ,							a=0.1       ,							a=1_28       ,							a=32       ,							a=16       ,							a=2       ,							a=0.02       ,							a=3       ,							a=4       ,							a=None       ,							)					->      Optional[int]:


								'''simple docstring'''

								_UpperCamelCase						    =	parent
								_UpperCamelCase						    =	batch_size
								_UpperCamelCase						    =	seq_length
								_UpperCamelCase						    =	is_training
								_UpperCamelCase						    =	use_input_mask
								_UpperCamelCase						    =	use_token_type_ids
								_UpperCamelCase						    =	use_labels
								_UpperCamelCase						    =	vocab_size
								_UpperCamelCase						    =	hidden_size
								_UpperCamelCase						    =	num_hidden_layers
								_UpperCamelCase						    =	num_attention_heads
								_UpperCamelCase						    =	intermediate_size
								_UpperCamelCase						    =	hidden_act
								_UpperCamelCase						    =	hidden_dropout_prob
								_UpperCamelCase						    =	attention_probs_dropout_prob
								_UpperCamelCase						    =	max_position_embeddings
								_UpperCamelCase						    =	type_vocab_size
								_UpperCamelCase						    =	type_sequence_label_size
								_UpperCamelCase						    =	initializer_range
								_UpperCamelCase						    =	num_labels
								_UpperCamelCase						    =	num_choices
								_UpperCamelCase						    =	scope

						def  A_				(			self				)					->      List[Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size				)

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

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

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

								_UpperCamelCase						    =	self.get_config()

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

						def  A_				(			self				)					->      Optional[Any]:


								'''simple docstring'''

								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				)					->      int:


								'''simple docstring'''

								(
								    (
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	
								)						    =	self.prepare_config_and_inputs()

								_UpperCamelCase						    =	True
								_UpperCamelCase						    =	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
								_UpperCamelCase						    =	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       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Union[str, Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaModel(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	model(a       ,							attention_mask=a       ,							token_type_ids=a				)
								_UpperCamelCase						    =	model(a       ,							token_type_ids=a				)
								_UpperCamelCase						    =	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       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							)					->      List[Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	True
								_UpperCamelCase						    =	NezhaModel(a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	model(
								    a       ,							attention_mask=a       ,							token_type_ids=a       ,							encoder_hidden_states=a       ,							encoder_attention_mask=a       ,							)
								_UpperCamelCase						    =	model(
								    a       ,							attention_mask=a       ,							token_type_ids=a       ,							encoder_hidden_states=a       ,							)
								_UpperCamelCase						    =	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       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Dict:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaForMaskedLM(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	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       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Optional[int]:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaForNextSentencePrediction(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	model(
								    a       ,							attention_mask=a       ,							token_type_ids=a       ,							labels=a       ,							)
								self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, 2)				)

						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Dict:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaForPreTraining(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	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       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaForQuestionAnswering(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	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       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Optional[int]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.num_labels
								_UpperCamelCase						    =	NezhaForSequenceClassification(a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	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       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Dict:


								'''simple docstring'''

								_UpperCamelCase						    =	self.num_labels
								_UpperCamelCase						    =	NezhaForTokenClassification(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	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       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Any:


								'''simple docstring'''

								_UpperCamelCase						    =	self.num_choices
								_UpperCamelCase						    =	NezhaForMultipleChoice(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	input_ids.unsqueeze(1				).expand(-1       ,							self.num_choices       ,							-1				).contiguous()
								_UpperCamelCase						    =	token_type_ids.unsqueeze(1				).expand(-1       ,							self.num_choices       ,							-1				).contiguous()
								_UpperCamelCase						    =	input_mask.unsqueeze(1				).expand(-1       ,							self.num_choices       ,							-1				).contiguous()
								_UpperCamelCase						    =	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				)					->      Tuple:


								'''simple docstring'''

								_UpperCamelCase						    =	self.prepare_config_and_inputs()
								(
								    (
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	
								)						    =	config_and_inputs
								_UpperCamelCase						    =	{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
								return config, inputs_dict







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

						def  A_				(			self       ,							a       ,							a       ,							a=False				)					->      List[Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	super()._prepare_for_class(a       ,							a       ,							return_labels=a				)

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

						def  A_				(			self				)					->      Optional[int]:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaModelTester(self				)
								_UpperCamelCase						    =	ConfigTester(self       ,							config_class=a       ,							hidden_size=37				)

						def  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								self.config_tester.run_common_tests()

						def  A_				(			self				)					->      Any:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*a				)

						def  A_				(			self				)					->      List[Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_model_as_decoder(*a				)

						def  A_				(			self				)					->      Dict:


								'''simple docstring'''

								(
								    (
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	(
								    _UpperCamelCase
								)    ,	
								)						    =	self.model_tester.prepare_config_and_inputs_for_decoder()

								_UpperCamelCase						    =	None

								self.model_tester.create_and_check_model_as_decoder(
								    a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							)

						def  A_				(			self				)					->      str:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*a				)

						def  A_				(			self				)					->      Union[str, Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*a				)

						def  A_				(			self				)					->      str:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_next_sequence_prediction(*a				)

						def  A_				(			self				)					->      int:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*a				)

						def  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*a				)

						def  A_				(			self				)					->      Any:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*a				)

						def  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*a				)

						@slow
						def  A_				(			self				)					->      List[Any]:


								'''simple docstring'''

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







						@slow
						@require_torch_gpu
						def  A_				(			self				)					->      Optional[Any]:


								'''simple docstring'''

								_UpperCamelCase    ,	_UpperCamelCase						    =	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						    =	True
										_UpperCamelCase						    =	model_class(config=a				)

										_UpperCamelCase						    =	self._prepare_for_class(a       ,							a				)
										_UpperCamelCase						    =	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						    =	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        lowerCAmelCase__      (   unittest.TestCase	):

						@slow
						def  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaModel.from_pretrained("""sijunhe/nezha-cn-base"""				)
								_UpperCamelCase						    =	torch.tensor([[0, 1, 2, 3, 4, 5]]				)
								_UpperCamelCase						    =	torch.tensor([[0, 1, 1, 1, 1, 1]]				)
								with torch.no_grad():
										_UpperCamelCase						    =	model(a       ,							attention_mask=a				)[0]
								_UpperCamelCase						    =	torch.Size((1, 6, 7_68)				)
								self.assertEqual(output.shape       ,							a				)
								_UpperCamelCase						    =	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				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaForMaskedLM.from_pretrained("""sijunhe/nezha-cn-base"""				)
								_UpperCamelCase						    =	torch.tensor([[0, 1, 2, 3, 4, 5]]				)
								_UpperCamelCase						    =	torch.tensor([[1, 1, 1, 1, 1, 1]]				)
								with torch.no_grad():
										_UpperCamelCase						    =	model(a       ,							attention_mask=a				)[0]
								_UpperCamelCase						    =	torch.Size((1, 6, 2_11_28)				)
								self.assertEqual(output.shape       ,							a				)
								_UpperCamelCase						    =	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				)				)



