'''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	:	List[Any]						,       a	:	str						,       a	:	Any=13						,       a	:	Any=7						,       a	:	int=True						,       a	:	Optional[int]=True						,       a	:	str=True						,       a	:	str=True						,       a	:	List[str]=99						,       a	:	List[Any]=32						,       a	:	int=5						,       a	:	Optional[int]=4						,       a	:	Tuple=37						,       a	:	List[Any]="gelu"						,       a	:	str=0.1						,       a	:	List[str]=0.1						,       a	:	Optional[int]=5_12						,       a	:	List[Any]=16						,       a	:	Tuple=2						,       a	:	Optional[Any]=0.02						,       a	:	int=3						,       a	:	List[Any]=4						,       a	:	Any=None						,       ):





         """simple docstring"""



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






     def      SCREAMING_SNAKE_CASE__     (				self	:	Optional[Any]					):





         """simple docstring"""



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

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

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

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

         __lowerCamelCase						=						self.get_config()

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






     def      SCREAMING_SNAKE_CASE__     (				self	:	Dict					):





         """simple docstring"""



         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      SCREAMING_SNAKE_CASE__     (				self	:	Optional[Any]					):





         """simple docstring"""



         __lowerCamelCase						=						self.get_config()
         __lowerCamelCase						=						3_00
         return config






     def      SCREAMING_SNAKE_CASE__     (				self	:	Tuple					):





         """simple docstring"""



         (
             (
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,												
         )						=						self.prepare_config_and_inputs()

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





         """simple docstring"""



         __lowerCamelCase						=						YosoModel(config=a					)
         model.to(a					)
         model.eval()
         __lowerCamelCase						=						model(a						,       attention_mask=a						,       token_type_ids=a					)
         __lowerCamelCase						=						model(a						,       token_type_ids=a					)
         __lowerCamelCase						=						model(a					)
         self.parent.assertEqual(result.last_hidden_state.shape						,       (self.batch_size, self.seq_length, self.hidden_size)					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	Optional[int]						,       a	:	List[str]						,       a	:	int						,       a	:	List[Any]						,       a	:	List[Any]						,       a	:	Dict						,       a	:	List[Any]						,       a	:	Any						,       a	:	Any						,       a	:	Optional[Any]						,       ):





         """simple docstring"""



         __lowerCamelCase						=						True
         __lowerCamelCase						=						YosoModel(a					)
         model.to(a					)
         model.eval()
         __lowerCamelCase						=						model(
             a						,       attention_mask=a						,       token_type_ids=a						,       encoder_hidden_states=a						,       encoder_attention_mask=a						,       )
         __lowerCamelCase						=						model(
             a						,       attention_mask=a						,       token_type_ids=a						,       encoder_hidden_states=a						,       )
         __lowerCamelCase						=						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      SCREAMING_SNAKE_CASE__     (				self	:	Optional[int]						,       a	:	Dict						,       a	:	Tuple						,       a	:	Union[str, Any]						,       a	:	Tuple						,       a	:	Union[str, Any]						,       a	:	Union[str, Any]						,       a	:	Optional[int]					):





         """simple docstring"""



         __lowerCamelCase						=						YosoForMaskedLM(config=a					)
         model.to(a					)
         model.eval()
         __lowerCamelCase						=						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      SCREAMING_SNAKE_CASE__     (				self	:	int						,       a	:	Union[str, Any]						,       a	:	Optional[int]						,       a	:	Any						,       a	:	str						,       a	:	Optional[Any]						,       a	:	int						,       a	:	List[str]					):





         """simple docstring"""



         __lowerCamelCase						=						YosoForQuestionAnswering(config=a					)
         model.to(a					)
         model.eval()
         __lowerCamelCase						=						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      SCREAMING_SNAKE_CASE__     (				self	:	List[Any]						,       a	:	Union[str, Any]						,       a	:	List[str]						,       a	:	Tuple						,       a	:	List[Any]						,       a	:	Optional[int]						,       a	:	Union[str, Any]						,       a	:	List[str]					):





         """simple docstring"""



         __lowerCamelCase						=						self.num_labels
         __lowerCamelCase						=						YosoForSequenceClassification(a					)
         model.to(a					)
         model.eval()
         __lowerCamelCase						=						model(a						,       attention_mask=a						,       token_type_ids=a						,       labels=a					)
         self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.num_labels)					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	str						,       a	:	Any						,       a	:	Dict						,       a	:	Union[str, Any]						,       a	:	Any						,       a	:	List[Any]						,       a	:	Tuple						,       a	:	Tuple					):





         """simple docstring"""



         __lowerCamelCase						=						self.num_labels
         __lowerCamelCase						=						YosoForTokenClassification(config=a					)
         model.to(a					)
         model.eval()
         __lowerCamelCase						=						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      SCREAMING_SNAKE_CASE__     (				self	:	Optional[Any]						,       a	:	List[str]						,       a	:	int						,       a	:	Any						,       a	:	List[Any]						,       a	:	str						,       a	:	List[str]						,       a	:	Any					):





         """simple docstring"""



         __lowerCamelCase						=						self.num_choices
         __lowerCamelCase						=						YosoForMultipleChoice(config=a					)
         model.to(a					)
         model.eval()
         __lowerCamelCase						=						input_ids.unsqueeze(1					).expand(-1						,       self.num_choices						,       -1					).contiguous()
         __lowerCamelCase						=						token_type_ids.unsqueeze(1					).expand(-1						,       self.num_choices						,       -1					).contiguous()
         __lowerCamelCase						=						input_mask.unsqueeze(1					).expand(-1						,       self.num_choices						,       -1					).contiguous()
         __lowerCamelCase						=						model(
             a						,       attention_mask=a						,       token_type_ids=a						,       labels=a						,       )
         self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.num_choices)					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	Dict					):





         """simple docstring"""



         __lowerCamelCase						=						self.prepare_config_and_inputs()
         (
             (
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,						(
             __lowerCamelCase						
         )			,												
         )						=						config_and_inputs
         __lowerCamelCase						=						{'''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				):
     lowerCamelCase						:  Dict							=(
         (
             YosoModel,
             YosoForMaskedLM,
             YosoForMultipleChoice,
             YosoForQuestionAnswering,
             YosoForSequenceClassification,
             YosoForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     lowerCamelCase						:  List[Any]							=False
     lowerCamelCase						:  str							=False
     lowerCamelCase						:  List[str]							=False

     lowerCamelCase						:  Tuple							=()
     lowerCamelCase						:  Tuple							=(
         {
             "feature-extraction": YosoModel,
             "fill-mask": YosoForMaskedLM,
             "question-answering": YosoForQuestionAnswering,
             "text-classification": YosoForSequenceClassification,
             "token-classification": YosoForTokenClassification,
             "zero-shot": YosoForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )






     def      SCREAMING_SNAKE_CASE__     (				self	:	int					):





         """simple docstring"""



         __lowerCamelCase						=						YosoModelTester(self					)
         __lowerCamelCase						=						ConfigTester(self						,       config_class=a						,       hidden_size=37					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	Union[str, Any]					):





         """simple docstring"""



         self.config_tester.run_common_tests()






     def      SCREAMING_SNAKE_CASE__     (				self	:	Union[str, Any]					):





         """simple docstring"""



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






     def      SCREAMING_SNAKE_CASE__     (				self	:	Union[str, Any]					):





         """simple docstring"""



         __lowerCamelCase						=						self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
             __lowerCamelCase						=						type
             self.model_tester.create_and_check_model(*a					)






     def      SCREAMING_SNAKE_CASE__     (				self	:	int					):





         """simple docstring"""



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






     def      SCREAMING_SNAKE_CASE__     (				self	:	Any					):





         """simple docstring"""



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






     def      SCREAMING_SNAKE_CASE__     (				self	:	Union[str, Any]					):





         """simple docstring"""



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






     def      SCREAMING_SNAKE_CASE__     (				self	:	Tuple					):





         """simple docstring"""



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






     def      SCREAMING_SNAKE_CASE__     (				self	:	Dict					):





         """simple docstring"""



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






     @slow
     def      SCREAMING_SNAKE_CASE__     (				self	:	List[Any]					):





         """simple docstring"""



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






     def      SCREAMING_SNAKE_CASE__     (				self	:	List[Any]					):





         """simple docstring"""



         return




@require_torch
class 					a__   (     unittest.TestCase				):






     @slow
     def      SCREAMING_SNAKE_CASE__     (				self	:	str					):





         """simple docstring"""



         __lowerCamelCase						=						YosoModel.from_pretrained('''uw-madison/yoso-4096'''					)
         __lowerCamelCase						=						torch.tensor([[0, 1, 2, 3, 4, 5]]					)

         with torch.no_grad():
             __lowerCamelCase						=						model(a					)[0]

         __lowerCamelCase						=						torch.Size((1, 6, 7_68)					)
         self.assertEqual(output.shape						,       a					)

         __lowerCamelCase						=						torch.tensor(
             [[[-0.06_11, 0.12_42, 0.08_40], [0.02_80, -0.00_48, 0.11_25], [0.01_06, 0.02_26, 0.07_51]]]					)

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






     @slow
     def      SCREAMING_SNAKE_CASE__     (				self	:	Any					):





         """simple docstring"""



         __lowerCamelCase						=						YosoForMaskedLM.from_pretrained('''uw-madison/yoso-4096'''					)
         __lowerCamelCase						=						torch.tensor([[0, 1, 2, 3, 4, 5]]					)

         with torch.no_grad():
             __lowerCamelCase						=						model(a					)[0]

         __lowerCamelCase						=						5_02_65

         __lowerCamelCase						=						torch.Size((1, 6, vocab_size)					)
         self.assertEqual(output.shape						,       a					)

         __lowerCamelCase						=						torch.tensor(
             [[[-2.13_13, -3.72_85, -2.24_07], [-2.70_47, -3.33_14, -2.64_08], [0.06_29, -2.51_66, -0.33_56]]]					)

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






     @slow
     def      SCREAMING_SNAKE_CASE__     (				self	:	List[str]					):





         """simple docstring"""



         __lowerCamelCase						=						YosoForMaskedLM.from_pretrained('''uw-madison/yoso-4096'''					)
         __lowerCamelCase						=						torch.arange(40_96					).unsqueeze(0					)

         with torch.no_grad():
             __lowerCamelCase						=						model(a					)[0]

         __lowerCamelCase						=						5_02_65

         __lowerCamelCase						=						torch.Size((1, 40_96, vocab_size)					)
         self.assertEqual(output.shape						,       a					)

         __lowerCamelCase						=						torch.tensor(
             [[[-2.39_14, -4.37_42, -5.09_56], [-4.09_88, -4.23_84, -7.04_06], [-3.14_27, -3.71_92, -6.68_00]]]					)

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

'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


__UpperCAmelCase  							={
    "configuration_rwkv": ["RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP", "RwkvConfig", "RwkvOnnxConfig"],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase  							=[
      "RWKV_PRETRAINED_MODEL_ARCHIVE_LIST",
      "RwkvForCausalLM",
      "RwkvModel",
      "RwkvPreTrainedModel",
  ]


if TYPE_CHECKING:
  from .configuration_rwkv import RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP, RwkvConfig, RwkvOnnxConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_rwkv import (
        RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
        RwkvForCausalLM,
        RwkvModel,
        RwkvPreTrainedModel,
    )
else:
  import sys

  __UpperCAmelCase  							=_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
