


'''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 			UpperCAmelCase			:

 def __init__(			self			:						Union[str, Any],       a_			:						str,       a_			:						Optional[Any]=13,       a_			:						Optional[int]=7,       a_			:						Optional[int]=True,       a_			:						List[Any]=True,       a_			:						Union[str, Any]=True,       a_			:						int=True,       a_			:						Optional[Any]=99,       a_			:						List[str]=32,       a_			:						int=5,       a_			:						Optional[int]=4,       a_			:						Optional[Any]=37,       a_			:						List[Any]="gelu",       a_			:						Dict=0.1,       a_			:						Any=0.1,       a_			:						Any=512,       a_			:						Optional[Any]=16,       a_			:						Tuple=2,       a_			:						Union[str, Any]=0.02,       a_			:						List[str]=3,       a_			:						Tuple=4,       a_			:						str=None,       ):





    """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 		lowercase_    (			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 		lowercase_    (			self			:						int       ):





    """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 		lowercase_    (			self			:						List[Any]       ):





    """simple docstring"""



    UpperCamelCase__			    =    self.get_config()
    UpperCamelCase__			    =    300
    return config

 def 		lowercase_    (			self			:						str       ):





    """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 		lowercase_    (			self			:						Optional[Any],       a_			:						int,       a_			:						Dict,       a_			:						List[Any],       a_			:						List[str],       a_			:						List[str],       a_			:						Any,       a_			:						List[Any]       ):





    """simple docstring"""



    UpperCamelCase__			    =    YosoModel(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)       )

 def 		lowercase_    (			self			:						str,       a_			:						int,       a_			:						Optional[int],       a_			:						Tuple,       a_			:						int,       a_			:						Optional[Any],       a_			:						Optional[int],       a_			:						str,       a_			:						Union[str, Any],       a_			:						Union[str, Any],       ):





    """simple docstring"""



    UpperCamelCase__			    =    True
    UpperCamelCase__			    =    YosoModel(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)       )

 def 		lowercase_    (			self			:						Union[str, Any],       a_			:						List[str],       a_			:						Optional[Any],       a_			:						List[str],       a_			:						Tuple,       a_			:						Dict,       a_			:						int,       a_			:						int       ):





    """simple docstring"""



    UpperCamelCase__			    =    YosoForMaskedLM(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 		lowercase_    (			self			:						int,       a_			:						str,       a_			:						Optional[Any],       a_			:						Dict,       a_			:						Any,       a_			:						Optional[Any],       a_			:						Union[str, Any],       a_			:						str       ):





    """simple docstring"""



    UpperCamelCase__			    =    YosoForQuestionAnswering(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 		lowercase_    (			self			:						Dict,       a_			:						Union[str, Any],       a_			:						Union[str, Any],       a_			:						str,       a_			:						Tuple,       a_			:						Union[str, Any],       a_			:						Any,       a_			:						Optional[Any]       ):





    """simple docstring"""



    UpperCamelCase__			    =    self.num_labels
    UpperCamelCase__			    =    YosoForSequenceClassification(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 		lowercase_    (			self			:						str,       a_			:						Dict,       a_			:						List[str],       a_			:						List[str],       a_			:						Any,       a_			:						str,       a_			:						Any,       a_			:						Tuple       ):





    """simple docstring"""



    UpperCamelCase__			    =    self.num_labels
    UpperCamelCase__			    =    YosoForTokenClassification(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 		lowercase_    (			self			:						Any,       a_			:						Dict,       a_			:						int,       a_			:						Optional[Any],       a_			:						List[str],       a_			:						str,       a_			:						Dict,       a_			:						Optional[int]       ):





    """simple docstring"""



    UpperCamelCase__			    =    self.num_choices
    UpperCamelCase__			    =    YosoForMultipleChoice(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 		lowercase_    (			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 			UpperCAmelCase			(							SCREAMING_SNAKE_CASE__		,      SCREAMING_SNAKE_CASE__		,      unittest.TestCase):
 _lowerCamelCase			:       Any            =       (
     (
         YosoModel,
         YosoForMaskedLM,
         YosoForMultipleChoice,
         YosoForQuestionAnswering,
         YosoForSequenceClassification,
         YosoForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 _lowerCamelCase			:       List[Any]            =       False
 _lowerCamelCase			:       List[Any]            =       False
 _lowerCamelCase			:       List[str]            =       False

 _lowerCamelCase			:       List[str]            =       ()
 _lowerCamelCase			:       Dict            =       (
     {
         '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			:						Optional[int]       ):





    """simple docstring"""



    UpperCamelCase__			    =    YosoModelTester(self       )
    UpperCamelCase__			    =    ConfigTester(self,       config_class=a_,       hidden_size=37       )

 def 		lowercase_    (			self			:						Tuple       ):





    """simple docstring"""



    self.config_tester.run_common_tests()

 def 		lowercase_    (			self			:						List[str]       ):





    """simple docstring"""



    UpperCamelCase__			    =    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*a_       )

 def 		lowercase_    (			self			:						Any       ):





    """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_model(*a_       )

 def 		lowercase_    (			self			:						List[Any]       ):





    """simple docstring"""



    UpperCamelCase__			    =    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*a_       )

 def 		lowercase_    (			self			:						Optional[int]       ):





    """simple docstring"""



    UpperCamelCase__			    =    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_multiple_choice(*a_       )

 def 		lowercase_    (			self			:						Tuple       ):





    """simple docstring"""



    UpperCamelCase__			    =    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_question_answering(*a_       )

 def 		lowercase_    (			self			:						Any       ):





    """simple docstring"""



    UpperCamelCase__			    =    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_sequence_classification(*a_       )

 def 		lowercase_    (			self			:						List[Any]       ):





    """simple docstring"""



    UpperCamelCase__			    =    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*a_       )

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





    """simple docstring"""



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

 def 		lowercase_    (			self			:						List[str]       ):





    """simple docstring"""



    return







@require_torch
class 			UpperCAmelCase			(							unittest.TestCase):

 @slow
 def 		lowercase_    (			self			:						Tuple       ):





    """simple docstring"""



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

    with torch.no_grad():
       UpperCamelCase__			    =    model(a_       )[0]

    UpperCamelCase__			    =    torch.Size((1, 6, 768)       )
    self.assertEqual(output.shape,       a_       )

    UpperCamelCase__			    =    torch.tensor(
        [[[-0.0_611, 0.1_242, 0.0_840], [0.0_280, -0.0_048, 0.1_125], [0.0_106, 0.0_226, 0.0_751]]]       )

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

 @slow
 def 		lowercase_    (			self			:						Union[str, Any]       ):





    """simple docstring"""



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

    with torch.no_grad():
       UpperCamelCase__			    =    model(a_       )[0]

    UpperCamelCase__			    =    5_0265

    UpperCamelCase__			    =    torch.Size((1, 6, vocab_size)       )
    self.assertEqual(output.shape,       a_       )

    UpperCamelCase__			    =    torch.tensor(
        [[[-2.1_313, -3.7_285, -2.2_407], [-2.7_047, -3.3_314, -2.6_408], [0.0_629, -2.5_166, -0.3_356]]]       )

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

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





    """simple docstring"""



    UpperCamelCase__			    =    YosoForMaskedLM.from_pretrained("uw-madison/yoso-4096"       )
    UpperCamelCase__			    =    torch.arange(4096       ).unsqueeze(0       )

    with torch.no_grad():
       UpperCamelCase__			    =    model(a_       )[0]

    UpperCamelCase__			    =    5_0265

    UpperCamelCase__			    =    torch.Size((1, 4096, vocab_size)       )
    self.assertEqual(output.shape,       a_       )

    UpperCamelCase__			    =    torch.tensor(
        [[[-2.3_914, -4.3_742, -5.0_956], [-4.0_988, -4.2_384, -7.0_406], [-3.1_427, -3.7_192, -6.6_800]]]       )

    self.assertTrue(torch.allclose(output[:, :3, :3],       a_,       atol=1e-4       )       )
"""simple docstring"""


from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCamelCase										=     logging.get_logger(__name__)

__UpperCamelCase										=     {
    '''uw-madison/nystromformer-512''': '''https://huggingface.co/uw-madison/nystromformer-512/resolve/main/config.json''',
    # See all Nystromformer models at https://huggingface.co/models?filter=nystromformer
}




class 					UpperCamelCase   (							lowerCAmelCase__  ):
       SCREAMING_SNAKE_CASE_						      = "nystromformer"







       def __init__(					self,				lowerCAmelCase__=3_0000,				lowerCAmelCase__=768,				lowerCAmelCase__=12,				lowerCAmelCase__=12,				lowerCAmelCase__=3072,				lowerCAmelCase__="gelu_new",				lowerCAmelCase__=0.1,				lowerCAmelCase__=0.1,				lowerCAmelCase__=510,				lowerCAmelCase__=2,				lowerCAmelCase__=64,				lowerCAmelCase__=64,				lowerCAmelCase__=65,				lowerCAmelCase__=False,				lowerCAmelCase__=0.02,				lowerCAmelCase__=1e-5,				lowerCAmelCase__=1,				lowerCAmelCase__=0,				lowerCAmelCase__=2,				**lowerCAmelCase__,				) -> Optional[Any]:
            snake_case_   			=   vocab_size
            snake_case_   			=   max_position_embeddings
            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_   			=   initializer_range
            snake_case_   			=   type_vocab_size
            snake_case_   			=   segment_means_seq_len
            snake_case_   			=   num_landmarks
            snake_case_   			=   conv_kernel_size
            snake_case_   			=   inv_coeff_init_option
            snake_case_   			=   layer_norm_eps
            super().__init__(pad_token_id=lowerCAmelCase__,				bos_token_id=lowerCAmelCase__,				eos_token_id=lowerCAmelCase__,				**lowerCAmelCase__)

