




'''simple docstring'''





import unittest

from transformers import YosoConfig, is_torch_available
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 (
       YosoForMaskedLM,
       YosoForMultipleChoice,
       YosoForQuestionAnswering,
       YosoForSequenceClassification,
       YosoForTokenClassification,
       YosoModel,
   )
   from transformers.models.yoso.modeling_yoso import YOSO_PRETRAINED_MODEL_ARCHIVE_LIST

class    _a				:






     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					,)							->			int:
          _snake_case							=      parent
          _snake_case							=      batch_size
          _snake_case							=      seq_length
          _snake_case							=      is_training
          _snake_case							=      use_input_mask
          _snake_case							=      use_token_type_ids
          _snake_case							=      use_labels
          _snake_case							=      vocab_size
          _snake_case							=      hidden_size
          _snake_case							=      num_hidden_layers
          _snake_case							=      num_attention_heads
          _snake_case							=      intermediate_size
          _snake_case							=      hidden_act
          _snake_case							=      hidden_dropout_prob
          _snake_case							=      attention_probs_dropout_prob
          _snake_case							=      max_position_embeddings
          _snake_case							=      type_vocab_size
          _snake_case							=      type_sequence_label_size
          _snake_case							=      initializer_range
          _snake_case							=      num_labels
          _snake_case							=      num_choices
          _snake_case							=      scope






     def 				_lowercase	( self    )							->			Any:
          _snake_case							=      ids_tensor([self.batch_size, self.seq_length]					,self.vocab_size    )

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

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

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

          _snake_case							=      self.get_config()

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






     def 				_lowercase	( self    )							->			Tuple:
          return YosoConfig(
              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 				_lowercase	( self    )							->			int:
          _snake_case							=      self.get_config()
          _snake_case							=      300
          return config






     def 				_lowercase	( self    )							->			Any:
          (
              (
              _snake_case						
          )							,					(
              _snake_case						
          )							,					(
              _snake_case						
          )							,					(
              _snake_case						
          )							,					(
              _snake_case						
          )							,					(
              _snake_case						
          )							,					(
              _snake_case						
          )							,											
          )							=      self.prepare_config_and_inputs()

          _snake_case							=      True
          _snake_case							=      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]    )
          _snake_case							=      ids_tensor([self.batch_size, self.seq_length]					,vocab_size=2    )

          return (
              config,
              input_ids,
              token_type_ids,
              input_mask,
              sequence_labels,
              token_labels,
              choice_labels,
              encoder_hidden_states,
              encoder_attention_mask,
          )






     def 				_lowercase	( self					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE    )							->			Dict:
          _snake_case							=      YosoModel(config=a__    )
          model.to(a__    )
          model.eval()
          _snake_case							=      model(a__					,attention_mask=a__					,token_type_ids=a__    )
          _snake_case							=      model(a__					,token_type_ids=a__    )
          _snake_case							=      model(a__    )
          self.parent.assertEqual(result.last_hidden_state.shape					,(self.batch_size, self.seq_length, self.hidden_size)    )






     def 				_lowercase	( 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]:
          _snake_case							=      True
          _snake_case							=      YosoModel(a__    )
          model.to(a__    )
          model.eval()
          _snake_case							=      model(
              a__					,attention_mask=a__					,token_type_ids=a__					,encoder_hidden_states=a__					,encoder_attention_mask=a__					,)
          _snake_case							=      model(
              a__					,attention_mask=a__					,token_type_ids=a__					,encoder_hidden_states=a__					,)
          _snake_case							=      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)    )






     def 				_lowercase	( self					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE    )							->			Optional[int]:
          _snake_case							=      YosoForMaskedLM(config=a__    )
          model.to(a__    )
          model.eval()
          _snake_case							=      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 				_lowercase	( self					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE    )							->			Optional[Any]:
          _snake_case							=      YosoForQuestionAnswering(config=a__    )
          model.to(a__    )
          model.eval()
          _snake_case							=      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 				_lowercase	( self					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE    )							->			Optional[Any]:
          _snake_case							=      self.num_labels
          _snake_case							=      YosoForSequenceClassification(a__    )
          model.to(a__    )
          model.eval()
          _snake_case							=      model(a__					,attention_mask=a__					,token_type_ids=a__					,labels=a__    )
          self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.num_labels)    )






     def 				_lowercase	( self					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE    )							->			Any:
          _snake_case							=      self.num_labels
          _snake_case							=      YosoForTokenClassification(config=a__    )
          model.to(a__    )
          model.eval()
          _snake_case							=      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 				_lowercase	( self					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE					,_SCREAMING_SNAKE_CASE    )							->			Dict:
          _snake_case							=      self.num_choices
          _snake_case							=      YosoForMultipleChoice(config=a__    )
          model.to(a__    )
          model.eval()
          _snake_case							=      input_ids.unsqueeze(1    ).expand(-1					,self.num_choices					,-1    ).contiguous()
          _snake_case							=      token_type_ids.unsqueeze(1    ).expand(-1					,self.num_choices					,-1    ).contiguous()
          _snake_case							=      input_mask.unsqueeze(1    ).expand(-1					,self.num_choices					,-1    ).contiguous()
          _snake_case							=      model(
              a__					,attention_mask=a__					,token_type_ids=a__					,labels=a__					,)
          self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.num_choices)    )






     def 				_lowercase	( self    )							->			List[str]:
          _snake_case							=      self.prepare_config_and_inputs()
          (
              (
              _snake_case						
          )							,					(
              _snake_case						
          )							,					(
              _snake_case						
          )							,					(
              _snake_case						
          )							,					(
              _snake_case						
          )							,					(
              _snake_case						
          )							,					(
              _snake_case						
          )							,											
          )							=      config_and_inputs
          _snake_case							=      {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
          return config, inputs_dict







@require_torch
class    _a				(						_UpperCamelCase					,					_UpperCamelCase					,					unittest.TestCase							):
     SCREAMING_SNAKE_CASE_							:    List[str]    				=						(
         (
             YosoModel,
             YosoForMaskedLM,
             YosoForMultipleChoice,
             YosoForQuestionAnswering,
             YosoForSequenceClassification,
             YosoForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     SCREAMING_SNAKE_CASE_							:    Optional[int]    				=						False
     SCREAMING_SNAKE_CASE_							:    str    				=						False
     SCREAMING_SNAKE_CASE_							:    Optional[int]    				=						False

     SCREAMING_SNAKE_CASE_							:    str    				=						()
     SCREAMING_SNAKE_CASE_							:    Optional[Any]    				=						(
         {
             """feature-extraction""": YosoModel,
             """fill-mask""": YosoForMaskedLM,
             """question-answering""": YosoForQuestionAnswering,
             """text-classification""": YosoForSequenceClassification,
             """token-classification""": YosoForTokenClassification,
             """zero-shot""": YosoForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )






     def 				_lowercase	( self    )							->			List[str]:
          _snake_case							=      YosoModelTester(self    )
          _snake_case							=      ConfigTester(self					,config_class=a__					,hidden_size=37    )






     def 				_lowercase	( self    )							->			Any:
          self.config_tester.run_common_tests()






     def 				_lowercase	( self    )							->			Any:
          _snake_case							=      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*a__    )






     def 				_lowercase	( self    )							->			Optional[int]:
          _snake_case							=      self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
               _snake_case							=      type
               self.model_tester.create_and_check_model(*a__    )






     def 				_lowercase	( self    )							->			List[str]:
          _snake_case							=      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*a__    )






     def 				_lowercase	( self    )							->			List[str]:
          _snake_case							=      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_multiple_choice(*a__    )






     def 				_lowercase	( self    )							->			Tuple:
          _snake_case							=      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*a__    )






     def 				_lowercase	( self    )							->			Tuple:
          _snake_case							=      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*a__    )






     def 				_lowercase	( self    )							->			Optional[int]:
          _snake_case							=      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*a__    )






     @slow
     def 				_lowercase	( self    )							->			Optional[int]:







          for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _snake_case							=      YosoModel.from_pretrained(a__    )
               self.assertIsNotNone(a__    )






     def 				_lowercase	( self    )							->			Tuple:
          return







@require_torch
class    _a				(						unittest.TestCase							):






     @slow
     def 				_lowercase	( self    )							->			Optional[Any]:
          _snake_case							=      YosoModel.from_pretrained("uw-madison/yoso-4096"    )
          _snake_case							=      torch.tensor([[0, 1, 2, 3, 4, 5]]    )

          with torch.no_grad():
               _snake_case							=      model(a__    )[0]

          _snake_case							=      torch.Size((1, 6, 768)    )
          self.assertEqual(output.shape					,a__    )

          _snake_case							=      torch.tensor(
              [[[-0.0_6_1_1, 0.1_2_4_2, 0.0_8_4_0], [0.0_2_8_0, -0.0_0_4_8, 0.1_1_2_5], [0.0_1_0_6, 0.0_2_2_6, 0.0_7_5_1]]]    )

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






     @slow
     def 				_lowercase	( self    )							->			Union[str, Any]:
          _snake_case							=      YosoForMaskedLM.from_pretrained("uw-madison/yoso-4096"    )
          _snake_case							=      torch.tensor([[0, 1, 2, 3, 4, 5]]    )

          with torch.no_grad():
               _snake_case							=      model(a__    )[0]

          _snake_case							=      50_265

          _snake_case							=      torch.Size((1, 6, vocab_size)    )
          self.assertEqual(output.shape					,a__    )

          _snake_case							=      torch.tensor(
              [[[-2.1_3_1_3, -3.7_2_8_5, -2.2_4_0_7], [-2.7_0_4_7, -3.3_3_1_4, -2.6_4_0_8], [0.0_6_2_9, -2.5_1_6_6, -0.3_3_5_6]]]    )

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






     @slow
     def 				_lowercase	( self    )							->			Optional[int]:
          _snake_case							=      YosoForMaskedLM.from_pretrained("uw-madison/yoso-4096"    )
          _snake_case							=      torch.arange(4_096    ).unsqueeze(0    )

          with torch.no_grad():
               _snake_case							=      model(a__    )[0]

          _snake_case							=      50_265

          _snake_case							=      torch.Size((1, 4_096, vocab_size)    )
          self.assertEqual(output.shape					,a__    )

          _snake_case							=      torch.tensor(
              [[[-2.3_9_1_4, -4.3_7_4_2, -5.0_9_5_6], [-4.0_9_8_8, -4.2_3_8_4, -7.0_4_0_6], [-3.1_4_2_7, -3.7_1_9_2, -6.6_8_0_0]]]    )

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






"""simple docstring"""


from __future__ import annotations

a_											=			[2, 1, 4, 5, 6, 0, 8, 9, 1, 2, 0, 6, 4, 2, 0, 6, 5, 3, 2, 7]
class snake_case					:







		def __init__(    self					:				Optional[int]			,      a__					:				int						)				->  None:





						'''simple docstring'''







						_A         =	-1
						_A         =	-1
						_A         =	[-1] * length
						_A         =	None
						_A         =	None







		def __repr__(    self					:				Any						)				->  str:





						'''simple docstring'''







						return F"""Node(min_value={self.minn} max_value={self.maxx})"""






def    a__       ( __lowercase      )		->		Node | None:
				_A         =	Node(len(__lowercase      )      )
				_A     ,						_A         =	min(__lowercase      ), max(__lowercase      )
				# Leaf node case where the node contains only one unique value
				if root.minn == root.maxx:
								return root
				_A         =	(root.minn + root.maxx) // 2

				_A         =	[]
				_A         =	[]

				for index, num in enumerate(__lowercase      ):
								if num <= pivot:
												left_arr.append(__lowercase      )
								else:
												right_arr.append(__lowercase      )
								_A         =	len(__lowercase      )
				_A         =	build_tree(__lowercase      )
				_A         =	build_tree(__lowercase      )
				return root






def    a__       ( __lowercase   ,	__lowercase   ,	__lowercase      )		->		int:
				if index < 0 or node is None:
								return 0
				# Leaf node cases
				if node.minn == node.maxx:
								return index + 1 if node.minn == num else 0
				_A         =	(node.minn + node.maxx) // 2
				if num <= pivot:
								# go the left subtree and map index to the left subtree
								return rank_till_index(node.left   ,	__lowercase   ,	node.map_left[index] - 1      )
				else:
								# go to the right subtree and map index to the right subtree
								return rank_till_index(node.right   ,	__lowercase   ,	index - node.map_left[index]      )






def    a__       ( __lowercase   ,	__lowercase   ,	__lowercase   ,	__lowercase      )		->		int:
				if start > end:
								return 0
				_A         =	rank_till_index(__lowercase   ,	__lowercase   ,	__lowercase      )
				_A         =	rank_till_index(__lowercase   ,	__lowercase   ,	start - 1      )
				return rank_till_end - rank_before_start






def    a__       ( __lowercase   ,	__lowercase   ,	__lowercase   ,	__lowercase      )		->		int:
				if index > (end - start) or start > end or node is None:
								return -1
				# Leaf node case
				if node.minn == node.maxx:
								return node.minn
				# Number of elements in the left subtree in interval [start, end]
				_A         =	node.map_left[end] - (
				    node.map_left[start - 1] if start else 0
				)
				if num_elements_in_left_tree > index:
								return quantile(
								    node.left   ,	__lowercase   ,	(node.map_left[start - 1] if start else 0)   ,	node.map_left[end] - 1   ,	)
				else:
								return quantile(
								    node.right   ,	index - num_elements_in_left_tree   ,	start - (node.map_left[start - 1] if start else 0)   ,	end - node.map_left[end]   ,	)






def    a__       ( __lowercase   ,	__lowercase   ,	__lowercase   ,	__lowercase   ,	__lowercase      )		->		int:
				if (
				    start > end
				    or node is None
				    or start_num > end_num
				    or node.minn > end_num
				    or node.maxx < start_num
				):
								return 0
				if start_num <= node.minn and node.maxx <= end_num:
								return end - start + 1
				_A         =	range_counting(
				    node.left   ,	(node.map_left[start - 1] if start else 0)   ,	node.map_left[end] - 1   ,	__lowercase   ,	__lowercase   ,	)
				_A         =	range_counting(
				    node.right   ,	start - (node.map_left[start - 1] if start else 0)   ,	end - node.map_left[end]   ,	__lowercase   ,	__lowercase   ,	)
				return left + right


if __name__ == "__main__":
		import doctest

		doctest.testmod()