
import unittest

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

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,
				    ElectraForCausalLM,
				    ElectraForMaskedLM,
				    ElectraForMultipleChoice,
				    ElectraForPreTraining,
				    ElectraForQuestionAnswering,
				    ElectraForSequenceClassification,
				    ElectraForTokenClassification,
				    ElectraModel,
				)
				from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST






class 							SCREAMING_SNAKE_CASE__     :

							"""simple docstring"""






							def __init__(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_=13  ,						SCREAMING_SNAKE_CASE_=7  ,						SCREAMING_SNAKE_CASE_=True  ,						SCREAMING_SNAKE_CASE_=True  ,						SCREAMING_SNAKE_CASE_=True  ,						SCREAMING_SNAKE_CASE_=True  ,						SCREAMING_SNAKE_CASE_=99  ,						SCREAMING_SNAKE_CASE_=32  ,						SCREAMING_SNAKE_CASE_=5  ,						SCREAMING_SNAKE_CASE_=4  ,						SCREAMING_SNAKE_CASE_=37  ,						SCREAMING_SNAKE_CASE_="gelu"  ,						SCREAMING_SNAKE_CASE_=0.1  ,						SCREAMING_SNAKE_CASE_=0.1  ,						SCREAMING_SNAKE_CASE_=512  ,						SCREAMING_SNAKE_CASE_=16  ,						SCREAMING_SNAKE_CASE_=2  ,						SCREAMING_SNAKE_CASE_=0.0_2  ,						SCREAMING_SNAKE_CASE_=3  ,						SCREAMING_SNAKE_CASE_=4  ,						SCREAMING_SNAKE_CASE_=None  ,						)-> Union[str, Any]:





													'''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      )-> Optional[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  	=      ids_tensor([self.batch_size, self.seq_length]  ,						1      )

													__UpperCamelCase  	=      self.get_config()

													return (
													    config,
													    input_ids,
													    token_type_ids,
													    input_mask,
													    sequence_labels,
													    token_labels,
													    choice_labels,
													    fake_token_labels,
													)



							def      A__		(	self      )-> int:





													'''simple docstring'''






													return ElectraConfig(
													    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=SCREAMING_SNAKE_CASE_  ,						initializer_range=self.initializer_range  ,						)



							def      A__		(	self      )-> Tuple:





													'''simple docstring'''






													(
													    (
													    __UpperCamelCase
													)	,				(
													    __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  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> str:





													'''simple docstring'''






													__UpperCamelCase  	=      ElectraModel(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.hidden_size)      )



							def      A__		(	self  ,						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_  ,						)-> Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      True
													__UpperCamelCase  	=      ElectraModel(SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(
													    SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						encoder_hidden_states=SCREAMING_SNAKE_CASE_  ,						encoder_attention_mask=SCREAMING_SNAKE_CASE_  ,						)
													__UpperCamelCase  	=      model(
													    SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						encoder_hidden_states=SCREAMING_SNAKE_CASE_  ,						)
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.hidden_size)      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> Optional[int]:





													'''simple docstring'''






													__UpperCamelCase  	=      ElectraForMaskedLM(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						labels=SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length, self.vocab_size)      )



							def      A__		(	self  ,						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_  ,						)-> Union[str, Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      ElectraForCausalLM(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						labels=SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length, self.vocab_size)      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> int:





													'''simple docstring'''






													__UpperCamelCase  	=      self.num_labels
													__UpperCamelCase  	=      ElectraForTokenClassification(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						labels=SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length, self.num_labels)      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> Optional[int]:





													'''simple docstring'''






													__UpperCamelCase  	=      self.num_labels
													__UpperCamelCase  	=      ElectraForPreTraining(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						labels=SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length)      )



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





													'''simple docstring'''






													__UpperCamelCase  	=      self.num_labels
													__UpperCamelCase  	=      ElectraForSequenceClassification(SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						labels=SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_labels)      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> Tuple:





													'''simple docstring'''






													__UpperCamelCase  	=      ElectraForQuestionAnswering(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(
													    SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						start_positions=SCREAMING_SNAKE_CASE_  ,						end_positions=SCREAMING_SNAKE_CASE_  ,						)
													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  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> List[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      self.num_choices
													__UpperCamelCase  	=      ElectraForMultipleChoice(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													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(
													    SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						labels=SCREAMING_SNAKE_CASE_  ,						)
													self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_choices)      )



							def      A__		(	self      )-> Dict:





													'''simple docstring'''






													__UpperCamelCase  	=      self.prepare_config_and_inputs()
													(
													    (
													    __UpperCamelCase
													)	,				(
													    __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 							SCREAMING_SNAKE_CASE__     (				SCREAMING_SNAKE_CASE_ ,		SCREAMING_SNAKE_CASE_ ,		unittest.TestCase     ):

							"""simple docstring"""






							_snake_case  		=	(
							    (
							        ElectraModel,
							        ElectraForPreTraining,
							        ElectraForMaskedLM,
							        ElectraForCausalLM,
							        ElectraForMultipleChoice,
							        ElectraForTokenClassification,
							        ElectraForSequenceClassification,
							        ElectraForQuestionAnswering,
							    )
							    if is_torch_available()
							    else ()
							)
							_snake_case  		=	(
							    {
							        'feature-extraction': ElectraModel,
							        'fill-mask': ElectraForMaskedLM,
							        'question-answering': ElectraForQuestionAnswering,
							        'text-classification': ElectraForSequenceClassification,
							        'text-generation': ElectraForCausalLM,
							        'token-classification': ElectraForTokenClassification,
							        'zero-shot': ElectraForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							_snake_case  		=	True



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_=False      )-> Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      super()._prepare_for_class(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						return_labels=SCREAMING_SNAKE_CASE_      )

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



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





													'''simple docstring'''






													__UpperCamelCase  	=      ElectraModelTester(self      )
													__UpperCamelCase  	=      ConfigTester(self  ,						config_class=SCREAMING_SNAKE_CASE_  ,						hidden_size=37      )



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





													'''simple docstring'''






													self.config_tester.run_common_tests()



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





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_model(*SCREAMING_SNAKE_CASE_      )



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





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs_for_decoder()
													self.model_tester.create_and_check_electra_model_as_decoder(*SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> Tuple:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													for type in ["absolute", "relative_key", "relative_key_query"]:
																			__UpperCamelCase  	=      type
																			self.model_tester.create_and_check_electra_model(*SCREAMING_SNAKE_CASE_      )



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





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_masked_lm(*SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> Tuple:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_token_classification(*SCREAMING_SNAKE_CASE_      )



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





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_pretraining(*SCREAMING_SNAKE_CASE_      )



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





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_sequence_classification(*SCREAMING_SNAKE_CASE_      )



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





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_question_answering(*SCREAMING_SNAKE_CASE_      )



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





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_electra_for_multiple_choice(*SCREAMING_SNAKE_CASE_      )



							@slow
							def      A__		(	self      )-> Any:





													'''simple docstring'''






													for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			__UpperCamelCase  	=      ElectraModel.from_pretrained(SCREAMING_SNAKE_CASE_      )
																			self.assertIsNotNone(SCREAMING_SNAKE_CASE_      )



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





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs_for_decoder()
													self.model_tester.create_and_check_electra_for_causal_lm(*SCREAMING_SNAKE_CASE_      )






@require_torch
class 							SCREAMING_SNAKE_CASE__     (				unittest.TestCase     ):

							"""simple docstring"""






							@slow
							def      A__		(	self      )-> str:





													'''simple docstring'''






													__UpperCamelCase  	=      ElectraModel.from_pretrained('''google/electra-small-discriminator'''      )
													__UpperCamelCase  	=      torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]      )
													__UpperCamelCase  	=      torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]      )
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_      )[0]
													__UpperCamelCase  	=      torch.Size((1, 11, 256)      )
													self.assertEqual(output.shape  ,						SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      torch.tensor(
													    [[[0.4_4_7_1, 0.6_8_2_1, -0.3_2_6_5], [0.4_6_2_7, 0.5_2_5_5, -0.3_6_6_8], [0.4_5_3_2, 0.3_3_1_3, -0.4_3_4_4]]]      )

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









import numpy as np






class 	_a			:






			def __init__(    self    ,  lowercase_=5    ,  lowercase_=5 )	->					Union[str, Any]:
									lowerCAmelCase	:   Union[str, Any]					     =					depth
									lowerCAmelCase	:   Tuple					     =					0
									lowerCAmelCase	:   Tuple					     =					None
									lowerCAmelCase	:   Union[str, Any]					     =					None
									lowerCAmelCase	:   Any					     =					min_leaf_size
									lowerCAmelCase	:   Optional[int]					     =					None






			def 							_snake_case   (    self    ,  lowercase_    ,  lowercase_ )	->					Any:
									if labels.ndim != 1:
															print("""Error: Input labels must be one dimensional""" )

									return np.mean((labels - prediction) ** 2 )






			def 							_snake_case   (    self    ,  lowercase_    ,  lowercase_ )	->					Dict:
									if x.ndim != 1:
															print("""Error: Input data set must be one dimensional""" )
															return
									if len(lowercase_ ) != len(lowercase_ ):
															print("""Error: X and y have different lengths""" )
															return
									if y.ndim != 1:
															print("""Error: Data set labels must be one dimensional""" )
															return

									if len(lowercase_ ) < 2 * self.min_leaf_size:
															lowerCAmelCase	:   Tuple					     =					np.mean(lowercase_ )
															return

									if self.depth == 1:
															lowerCAmelCase	:   Tuple					     =					np.mean(lowercase_ )
															return

									lowerCAmelCase	:   List[str]					     =					0
									lowerCAmelCase	:   List[Any]					     =					self.mean_squared_error(lowercase_    ,  np.mean(lowercase_ ) ) * 2
									for i in range(len(lowercase_ ) ):
															if len(x[:i] ) < self.min_leaf_size:
																					continue
															elif len(x[i:] ) < self.min_leaf_size:
																					continue
															else:
																					lowerCAmelCase	:   int					     =					self.mean_squared_error(x[:i]    ,  np.mean(y[:i] ) )
																					lowerCAmelCase	:   Union[str, Any]					     =					self.mean_squared_error(x[i:]    ,  np.mean(y[i:] ) )
																					lowerCAmelCase	:   Tuple					     =					error_left + error_right
																					if error < min_error:
																											lowerCAmelCase	:   Tuple					     =					i
																											lowerCAmelCase	:   Tuple					     =					error

									if best_split != 0:
															lowerCAmelCase	:   int					     =					x[:best_split]
															lowerCAmelCase	:   List[str]					     =					y[:best_split]
															lowerCAmelCase	:   Dict					     =					x[best_split:]
															lowerCAmelCase	:   Tuple					     =					y[best_split:]

															lowerCAmelCase	:   Any					     =					x[best_split]
															lowerCAmelCase	:   str					     =					DecisionTree(
															    depth=self.depth - 1    ,  min_leaf_size=self.min_leaf_size )
															lowerCAmelCase	:   List[str]					     =					DecisionTree(
															    depth=self.depth - 1    ,  min_leaf_size=self.min_leaf_size )
															self.left.train(lowercase_    ,  lowercase_ )
															self.right.train(lowercase_    ,  lowercase_ )
									else:
															lowerCAmelCase	:   List[str]					     =					np.mean(lowercase_ )

									return






			def 							_snake_case   (    self    ,  lowercase_ )	->					int:
									if self.prediction is not None:
															return self.prediction
									elif self.left or self.right is not None:
															if x >= self.decision_boundary:
																					return self.right.predict(lowercase_ )
															else:
																					return self.left.predict(lowercase_ )
									else:
															print("""Error: Decision tree not yet trained""" )
															return None






class 	_a			:






			@staticmethod
			def 							_snake_case   (    lowercase_    ,  lowercase_ )	->					Optional[Any]:
									lowerCAmelCase	:   Optional[int]					     =					float(0 )
									for label in labels:
															squared_error_sum += (label - prediction) ** 2

									return float(squared_error_sum / labels.size )







def  _UpperCAmelCase  (				):


						'''simple docstring'''




						lowerCAmelCase	:   List[str]					     =					np.arange(-1.0			,1.0			,0.005  )
						lowerCAmelCase	:   List[str]					     =					np.sin(SCREAMING_SNAKE_CASE__  )

						lowerCAmelCase	:   Tuple					     =					DecisionTree(depth=1_0			,min_leaf_size=1_0  )
						tree.train(SCREAMING_SNAKE_CASE__			,SCREAMING_SNAKE_CASE__  )

						lowerCAmelCase	:   Dict					     =					(np.random.rand(1_0  ) * 2) - 1
						lowerCAmelCase	:   str					     =					np.array([tree.predict(SCREAMING_SNAKE_CASE__  ) for x in test_cases]  )
						lowerCAmelCase	:   Tuple					     =					np.mean((predictions - test_cases) ** 2  )

						print("""Test values: """ + str(SCREAMING_SNAKE_CASE__  )  )
						print("""Predictions: """ + str(SCREAMING_SNAKE_CASE__  )  )
						print("""Average error: """ + str(SCREAMING_SNAKE_CASE__  )  )


if __name__ == "__main__":
		main()
		import doctest

		doctest.testmod(name='mean_squarred_error', verbose=True)

