





from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import logging
from .configuration_regnet import RegNetConfig


UpperCAmelCase     :     Union[str, Any]							=	logging.get_logger(__name__)

# General docstring
UpperCAmelCase     :     str							=	"""RegNetConfig"""

# Base docstring
UpperCAmelCase     :     int							=	"""facebook/regnet-y-040"""
UpperCAmelCase     :     Union[str, Any]							=	[1, 1088, 7, 7]

# Image classification docstring
UpperCAmelCase     :     List[str]							=	"""facebook/regnet-y-040"""
UpperCAmelCase     :     str							=	"""tabby, tabby cat"""

UpperCAmelCase     :     Optional[int]							=	[
    """facebook/regnet-y-040""",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]







class     __lowerCAmelCase	(   nn.Module):







      def __init__(   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__ = 3		,				lowerCAmelCase__ = 1		,				lowerCAmelCase__ = 1		,				lowerCAmelCase__ = "relu"		,				)       ->						Any:



        '''simple docstring'''




        super().__init__()
        a__     : Dict  =nn.Convad(
            lowerCAmelCase__		,				lowerCAmelCase__		,				kernel_size=lowerCAmelCase__		,				stride=lowerCAmelCase__		,				padding=kernel_size // 2		,				groups=lowerCAmelCase__		,				bias=lowerCAmelCase__		,				)
        a__     : List[Any]  =nn.BatchNormad(lowerCAmelCase__							)
        a__     : Optional[Any]  =ACTaFN[activation] if activation is not None else nn.Identity()







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						Dict:



        '''simple docstring'''




        a__     : str  =self.convolution(lowerCAmelCase__							)
        a__     : str  =self.normalization(lowerCAmelCase__							)
        a__     : str  =self.activation(lowerCAmelCase__							)
        return hidden_state







class     __lowerCAmelCase	(   nn.Module):







      def __init__(   self		,				lowerCAmelCase__							)       ->						Optional[Any]:



        '''simple docstring'''




        super().__init__()
        a__     : Dict  =RegNetConvLayer(
            config.num_channels		,				config.embedding_size		,				kernel_size=3		,				stride=2		,				activation=config.hidden_act							)
        a__     : Optional[int]  =config.num_channels







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						List[Any]:



        '''simple docstring'''




        a__     : int  =pixel_values.shape[1]
        if num_channels != self.num_channels:
          raise ValueError(
              "Make sure that the channel dimension of the pixel values match with the one set in the configuration."							)
        a__     : Tuple  =self.embedder(lowerCAmelCase__							)
        return hidden_state







class     __lowerCAmelCase	(   nn.Module):







      def __init__(   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__ = 2							)       ->						int:



        '''simple docstring'''




        super().__init__()
        a__     : Optional[Any]  =nn.Convad(lowerCAmelCase__		,				lowerCAmelCase__		,				kernel_size=1		,				stride=lowerCAmelCase__		,				bias=lowerCAmelCase__							)
        a__     : int  =nn.BatchNormad(lowerCAmelCase__							)







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						Tensor:



        '''simple docstring'''




        a__     : Union[str, Any]  =self.convolution(lowerCAmelCase__							)
        a__     : Optional[Any]  =self.normalization(lowerCAmelCase__							)
        return hidden_state







class     __lowerCAmelCase	(   nn.Module):







      def __init__(   self		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						List[str]:



        '''simple docstring'''




        super().__init__()

        a__     : str  =nn.AdaptiveAvgPoolad((1, 1)							)
        a__     : str  =nn.Sequential(
            nn.Convad(lowerCAmelCase__		,				lowerCAmelCase__		,				kernel_size=1							)		,				nn.ReLU()		,				nn.Convad(lowerCAmelCase__		,				lowerCAmelCase__		,				kernel_size=1							)		,				nn.Sigmoid()		,				)







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						Optional[int]:



        '''simple docstring'''




        a__     : int  =self.pooler(lowerCAmelCase__							)
        a__     : Union[str, Any]  =self.attention(lowerCAmelCase__							)
        a__     : Any  =hidden_state * attention
        return hidden_state







class     __lowerCAmelCase	(   nn.Module):







      def __init__(   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__ = 1							)       ->						Tuple:



        '''simple docstring'''




        super().__init__()
        a__     : List[Any]  =in_channels != out_channels or stride != 1
        a__     : Dict  =max(1		,				out_channels // config.groups_width							)
        a__     : Optional[int]  =(
            RegNetShortCut(lowerCAmelCase__		,				lowerCAmelCase__		,				stride=lowerCAmelCase__							) if should_apply_shortcut else nn.Identity()
        )
        a__     : Dict  =nn.Sequential(
            RegNetConvLayer(lowerCAmelCase__		,				lowerCAmelCase__		,				kernel_size=1		,				activation=config.hidden_act							)		,				RegNetConvLayer(lowerCAmelCase__		,				lowerCAmelCase__		,				stride=lowerCAmelCase__		,				groups=lowerCAmelCase__		,				activation=config.hidden_act							)		,				RegNetConvLayer(lowerCAmelCase__		,				lowerCAmelCase__		,				kernel_size=1		,				activation=lowerCAmelCase__							)		,				)
        a__     : int  =ACTaFN[config.hidden_act]







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						Any:



        '''simple docstring'''




        a__     : Optional[Any]  =hidden_state
        a__     : List[str]  =self.layer(lowerCAmelCase__							)
        a__     : Optional[Any]  =self.shortcut(lowerCAmelCase__							)
        hidden_state += residual
        a__     : str  =self.activation(lowerCAmelCase__							)
        return hidden_state







class     __lowerCAmelCase	(   nn.Module):







      def __init__(   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__ = 1							)       ->						List[str]:



        '''simple docstring'''




        super().__init__()
        a__     : Optional[int]  =in_channels != out_channels or stride != 1
        a__     : Dict  =max(1		,				out_channels // config.groups_width							)
        a__     : List[Any]  =(
            RegNetShortCut(lowerCAmelCase__		,				lowerCAmelCase__		,				stride=lowerCAmelCase__							) if should_apply_shortcut else nn.Identity()
        )
        a__     : Dict  =nn.Sequential(
            RegNetConvLayer(lowerCAmelCase__		,				lowerCAmelCase__		,				kernel_size=1		,				activation=config.hidden_act							)		,				RegNetConvLayer(lowerCAmelCase__		,				lowerCAmelCase__		,				stride=lowerCAmelCase__		,				groups=lowerCAmelCase__		,				activation=config.hidden_act							)		,				RegNetSELayer(lowerCAmelCase__		,				reduced_channels=int(round(in_channels / 4							)							)							)		,				RegNetConvLayer(lowerCAmelCase__		,				lowerCAmelCase__		,				kernel_size=1		,				activation=lowerCAmelCase__							)		,				)
        a__     : Optional[int]  =ACTaFN[config.hidden_act]







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						List[str]:



        '''simple docstring'''




        a__     : List[str]  =hidden_state
        a__     : Any  =self.layer(lowerCAmelCase__							)
        a__     : List[Any]  =self.shortcut(lowerCAmelCase__							)
        hidden_state += residual
        a__     : Optional[int]  =self.activation(lowerCAmelCase__							)
        return hidden_state







class     __lowerCAmelCase	(   nn.Module):







      def __init__(   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__ = 2		,				lowerCAmelCase__ = 2		,				)       ->						Union[str, Any]:



        '''simple docstring'''




        super().__init__()

        a__     : Optional[int]  =RegNetXLayer if config.layer_type == "x" else RegNetYLayer

        a__     : List[Any]  =nn.Sequential(
            # downsampling is done in the first layer with stride of 2
            layer(
                lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				stride=lowerCAmelCase__		,				)		,				*[layer(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							) for _ in range(depth - 1							)]		,				)







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						Tuple:



        '''simple docstring'''




        a__     : Any  =self.layers(lowerCAmelCase__							)
        return hidden_state







class     __lowerCAmelCase	(   nn.Module):







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



        '''simple docstring'''




        super().__init__()
        a__     : List[Any]  =nn.ModuleList([]							)
        # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
        self.stages.append(
            RegNetStage(
                lowerCAmelCase__		,				config.embedding_size		,				config.hidden_sizes[0]		,				stride=2 if config.downsample_in_first_stage else 1		,				depth=config.depths[0]		,				)							)
        a__     : str  =zip(config.hidden_sizes		,				config.hidden_sizes[1:]							)
        for (in_channels, out_channels), depth in zip(lowerCAmelCase__		,				config.depths[1:]							):
          self.stages.append(RegNetStage(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				depth=lowerCAmelCase__							)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__ = False		,				lowerCAmelCase__ = True							)       ->						BaseModelOutputWithNoAttention:



        '''simple docstring'''




        a__     : int  =() if output_hidden_states else None

        for stage_module in self.stages:
          if output_hidden_states:
            a__     : str  =hidden_states + (hidden_state,)

          a__     : str  =stage_module(lowerCAmelCase__							)

        if output_hidden_states:
          a__     : Optional[Any]  =hidden_states + (hidden_state,)

        if not return_dict:
          return tuple(v for v in [hidden_state, hidden_states] if v is not None							)

        return BaseModelOutputWithNoAttention(last_hidden_state=lowerCAmelCase__		,				hidden_states=lowerCAmelCase__							)







class     __lowerCAmelCase	(   UpperCamelCase__):
      _lowercase    : Optional[int]            =      RegNetConfig
      _lowercase    : Optional[int]            =      """regnet"""
      _lowercase    : str            =      """pixel_values"""
      _lowercase    : List[Any]            =      True







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						Optional[Any]:



        '''simple docstring'''




        if isinstance(lowerCAmelCase__		,				nn.Convad							):
          nn.init.kaiming_normal_(module.weight		,				mode="fan_out"		,				nonlinearity="relu"							)
        elif isinstance(lowerCAmelCase__		,				(nn.BatchNormad, nn.GroupNorm)							):
          nn.init.constant_(module.weight		,				1							)
          nn.init.constant_(module.bias		,				0							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__=False							)       ->						int:



        '''simple docstring'''




        if isinstance(lowerCAmelCase__		,				lowerCAmelCase__							):
          a__     : Dict  =value
UpperCAmelCase     :     List[Any]							=	r"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

UpperCAmelCase     :     List[str]							=	r"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
"""







@add_start_docstrings(
    """The bare RegNet model outputting raw features without any specific head on top."""						,    UpperCamelCase__						,    )
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class     __lowerCAmelCase	(   UpperCamelCase__):







      def __init__(   self		,				lowerCAmelCase__							)       ->						Dict:



        '''simple docstring'''




        super().__init__(lowerCAmelCase__							)
        a__     : List[str]  =config
        a__     : Dict  =RegNetEmbeddings(lowerCAmelCase__							)
        a__     : Any  =RegNetEncoder(lowerCAmelCase__							)
        a__     : Tuple  =nn.AdaptiveAvgPoolad((1, 1)							)
        # Initialize weights and apply final processing
        self.post_init()







      @add_start_docstrings_to_model_forward(lowerCAmelCase__							)
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC		,				output_type=lowerCAmelCase__		,				config_class=_CONFIG_FOR_DOC		,				modality="vision"		,				expected_output=_EXPECTED_OUTPUT_SHAPE		,				)
      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__ = None		,				lowerCAmelCase__ = None							)       ->						BaseModelOutputWithPoolingAndNoAttention:



        '''simple docstring'''




        a__     : Any  =(
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        a__     : Optional[int]  =return_dict if return_dict is not None else self.config.use_return_dict

        a__     : List[str]  =self.embedder(lowerCAmelCase__							)

        a__     : Dict  =self.encoder(
            lowerCAmelCase__		,				output_hidden_states=lowerCAmelCase__		,				return_dict=lowerCAmelCase__							)

        a__     : Tuple  =encoder_outputs[0]

        a__     : List[Any]  =self.pooler(lowerCAmelCase__							)

        if not return_dict:
          return (last_hidden_state, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPoolingAndNoAttention(
            last_hidden_state=lowerCAmelCase__		,				pooler_output=lowerCAmelCase__		,				hidden_states=encoder_outputs.hidden_states		,				)







@add_start_docstrings(
    """
    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    """						,    UpperCamelCase__						,    )
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class     __lowerCAmelCase	(   UpperCamelCase__):







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



        '''simple docstring'''




        super().__init__(lowerCAmelCase__							)
        a__     : List[str]  =config.num_labels
        a__     : Optional[Any]  =RegNetModel(lowerCAmelCase__							)
        # classification head
        a__     : List[str]  =nn.Sequential(
            nn.Flatten()		,				nn.Linear(config.hidden_sizes[-1]		,				config.num_labels							) if config.num_labels > 0 else nn.Identity()		,				)
        # initialize weights and apply final processing
        self.post_init()







      @add_start_docstrings_to_model_forward(lowerCAmelCase__							)
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT		,				output_type=lowerCAmelCase__		,				config_class=_CONFIG_FOR_DOC		,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT		,				)
      def        _lowercase  (   self		,				lowerCAmelCase__ = None		,				lowerCAmelCase__ = None		,				lowerCAmelCase__ = None		,				lowerCAmelCase__ = None		,				)       ->						ImageClassifierOutputWithNoAttention:



        '''simple docstring'''




        a__     : Optional[Any]  =return_dict if return_dict is not None else self.config.use_return_dict

        a__     : Union[str, Any]  =self.regnet(lowerCAmelCase__		,				output_hidden_states=lowerCAmelCase__		,				return_dict=lowerCAmelCase__							)

        a__     : int  =outputs.pooler_output if return_dict else outputs[1]

        a__     : int  =self.classifier(lowerCAmelCase__							)

        a__     : Tuple  =None

        if labels is not None:
          if self.config.problem_type is None:
            if self.num_labels == 1:
              a__     : List[str]  ="regression"
            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
              a__     : Union[str, Any]  ="single_label_classification"
            else:
              a__     : Optional[int]  ="multi_label_classification"
          if self.config.problem_type == "regression":
            a__     : str  =MSELoss()
            if self.num_labels == 1:
              a__     : Dict  =loss_fct(logits.squeeze()		,				labels.squeeze()							)
            else:
              a__     : Dict  =loss_fct(lowerCAmelCase__		,				lowerCAmelCase__							)
          elif self.config.problem_type == "single_label_classification":
            a__     : Dict  =CrossEntropyLoss()
            a__     : int  =loss_fct(logits.view(-1		,				self.num_labels							)		,				labels.view(-1							)							)
          elif self.config.problem_type == "multi_label_classification":
            a__     : List[str]  =BCEWithLogitsLoss()
            a__     : Any  =loss_fct(lowerCAmelCase__		,				lowerCAmelCase__							)

        if not return_dict:
          a__     : List[Any]  =(logits,) + outputs[2:]
          return (loss,) + output if loss is not None else output

        return ImageClassifierOutputWithNoAttention(loss=lowerCAmelCase__		,				logits=lowerCAmelCase__		,				hidden_states=outputs.hidden_states							)













import unittest

from transformers import AutoTokenizer, NystromformerConfig, 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, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         NystromformerForMaskedLM,
         NystromformerForMultipleChoice,
         NystromformerForQuestionAnswering,
         NystromformerForSequenceClassification,
         NystromformerForTokenClassification,
         NystromformerModel,
     )
     from transformers.models.nystromformer.modeling_nystromformer import NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST







class     __lowerCAmelCase	:







      def __init__(   self		,				lowerCAmelCase__		,				lowerCAmelCase__=1_3		,				lowerCAmelCase__=7		,				lowerCAmelCase__=True		,				lowerCAmelCase__=True		,				lowerCAmelCase__=True		,				lowerCAmelCase__=True		,				lowerCAmelCase__=9_9		,				lowerCAmelCase__=3_2		,				lowerCAmelCase__=5		,				lowerCAmelCase__=4		,				lowerCAmelCase__=3_7		,				lowerCAmelCase__="gelu"		,				lowerCAmelCase__=0.1		,				lowerCAmelCase__=0.1		,				lowerCAmelCase__=5_1_2		,				lowerCAmelCase__=1_6		,				lowerCAmelCase__=2		,				lowerCAmelCase__=0.02		,				lowerCAmelCase__=3		,				lowerCAmelCase__=4		,				lowerCAmelCase__=None		,				)       ->						List[Any]:



        '''simple docstring'''




        a__     : Optional[Any]  =parent
        a__     : Tuple  =batch_size
        a__     : List[Any]  =seq_length
        a__     : Dict  =is_training
        a__     : Any  =use_input_mask
        a__     : int  =use_token_type_ids
        a__     : Optional[Any]  =use_labels
        a__     : Optional[Any]  =vocab_size
        a__     : List[str]  =hidden_size
        a__     : int  =num_hidden_layers
        a__     : Tuple  =num_attention_heads
        a__     : Union[str, Any]  =intermediate_size
        a__     : Optional[int]  =hidden_act
        a__     : int  =hidden_dropout_prob
        a__     : Union[str, Any]  =attention_probs_dropout_prob
        a__     : List[Any]  =max_position_embeddings
        a__     : str  =type_vocab_size
        a__     : Optional[Any]  =type_sequence_label_size
        a__     : Union[str, Any]  =initializer_range
        a__     : List[Any]  =num_labels
        a__     : str  =num_choices
        a__     : int  =scope







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        a__     : List[Any]  =ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size							)

        a__     : str  =None
        if self.use_input_mask:
          a__     : List[Any]  =random_attention_mask([self.batch_size, self.seq_length]							)

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

        a__     : Dict  =None
        a__     : str  =None
        a__     : str  =None
        if self.use_labels:
          a__     : List[str]  =ids_tensor([self.batch_size]		,				self.type_sequence_label_size							)
          a__     : Tuple  =ids_tensor([self.batch_size, self.seq_length]		,				self.num_labels							)
          a__     : Dict  =ids_tensor([self.batch_size]		,				self.num_choices							)

        a__     : Tuple  =self.get_config()

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







      def        _lowercase  (   self							)       ->						Tuple:



        '''simple docstring'''




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







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						List[Any]:



        '''simple docstring'''




        a__     : Tuple  =NystromformerModel(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : Optional[Any]  =model(lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__							)
        a__     : str  =model(lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__							)
        a__     : Optional[int]  =model(lowerCAmelCase__							)
        self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, self.seq_length, self.hidden_size)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						int:



        '''simple docstring'''




        a__     : int  =NystromformerForMaskedLM(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : Dict  =model(lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__		,				labels=lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.vocab_size)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Any:



        '''simple docstring'''




        a__     : Optional[int]  =NystromformerForQuestionAnswering(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : str  =model(
            lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__		,				start_positions=lowerCAmelCase__		,				end_positions=lowerCAmelCase__		,				)
        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		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Optional[int]:



        '''simple docstring'''




        a__     : Optional[Any]  =self.num_labels
        a__     : Dict  =NystromformerForSequenceClassification(lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : List[str]  =model(lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__		,				labels=lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_labels)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						List[str]:



        '''simple docstring'''




        a__     : Tuple  =self.num_labels
        a__     : List[str]  =NystromformerForTokenClassification(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : List[Any]  =model(lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__		,				labels=lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.num_labels)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Tuple:



        '''simple docstring'''




        a__     : List[Any]  =self.num_choices
        a__     : Optional[Any]  =NystromformerForMultipleChoice(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : List[str]  =input_ids.unsqueeze(1							).expand(-1		,				self.num_choices		,				-1							).contiguous()
        a__     : List[Any]  =token_type_ids.unsqueeze(1							).expand(-1		,				self.num_choices		,				-1							).contiguous()
        a__     : List[Any]  =input_mask.unsqueeze(1							).expand(-1		,				self.num_choices		,				-1							).contiguous()
        a__     : Dict  =model(
            lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__		,				labels=lowerCAmelCase__		,				)
        self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_choices)							)







      def        _lowercase  (   self							)       ->						Tuple:



        '''simple docstring'''




        a__     : Optional[Any]  =self.prepare_config_and_inputs()
        (
            (
            a__
        ) ,							(
            a__
        ) ,							(
            a__
        ) ,							(
            a__
        ) ,							(
            a__
        ) ,							(
            a__
        ) ,							(
            a__
        ) ,							
        )     : List[str]  =config_and_inputs
        a__     : str  ={"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
        return config, inputs_dict







@require_torch
class     __lowerCAmelCase	(   UpperCamelCase__						,    UpperCamelCase__						,    unittest.TestCase):
      _lowercase    : int            =      (
          (
              NystromformerModel,
              NystromformerForMaskedLM,
              NystromformerForMultipleChoice,
              NystromformerForQuestionAnswering,
              NystromformerForSequenceClassification,
              NystromformerForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      _lowercase    : Union[str, Any]            =      (
          {
              """feature-extraction""": NystromformerModel,
              """fill-mask""": NystromformerForMaskedLM,
              """question-answering""": NystromformerForQuestionAnswering,
              """text-classification""": NystromformerForSequenceClassification,
              """token-classification""": NystromformerForTokenClassification,
              """zero-shot""": NystromformerForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _lowercase    : Union[str, Any]            =      False
      _lowercase    : Union[str, Any]            =      False







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



        '''simple docstring'''




        a__     : Optional[int]  =NystromformerModelTester(self							)
        a__     : Optional[int]  =ConfigTester(self		,				config_class=lowerCAmelCase__		,				hidden_size=3_7							)







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        self.config_tester.run_common_tests()







      def        _lowercase  (   self							)       ->						Dict:



        '''simple docstring'''




        a__     : List[str]  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowerCAmelCase__							)







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



        '''simple docstring'''




        a__     : Tuple  =self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
          a__     : int  =type
          self.model_tester.create_and_check_model(*lowerCAmelCase__							)







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



        '''simple docstring'''




        a__     : Optional[Any]  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Dict:



        '''simple docstring'''




        a__     : str  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*lowerCAmelCase__							)







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



        '''simple docstring'''




        a__     : Tuple  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase__							)







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



        '''simple docstring'''




        a__     : Tuple  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*lowerCAmelCase__							)







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



        '''simple docstring'''




        a__     : int  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase__							)







      @slow
      def        _lowercase  (   self							)       ->						str:



        '''simple docstring'''




        for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          a__     : int  =NystromformerModel.from_pretrained(lowerCAmelCase__							)
          self.assertIsNotNone(lowerCAmelCase__							)







@require_torch
class     __lowerCAmelCase	(   unittest.TestCase):







      @slow
      def        _lowercase  (   self							)       ->						str:



        '''simple docstring'''




        a__     : str  =NystromformerModel.from_pretrained("uw-madison/nystromformer-512"							)
        a__     : int  =torch.tensor([[0, 1, 2, 3, 4, 5]]							)

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

        a__     : List[str]  =torch.Size((1, 6, 7_6_8)							)
        self.assertEqual(output.shape		,				lowerCAmelCase__							)

        a__     : int  =torch.tensor(
            [[[-0.45_32, -0.09_36, 0.51_37], [-0.26_76, 0.06_28, 0.61_86], [-0.36_29, -0.17_26, 0.47_16]]]							)

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







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



        '''simple docstring'''




        a__     : Union[str, Any]  ="the [MASK] of Belgium is Brussels"

        a__     : str  =AutoTokenizer.from_pretrained("uw-madison/nystromformer-512"							)
        a__     : int  =NystromformerForMaskedLM.from_pretrained("uw-madison/nystromformer-512"							)

        a__     : List[Any]  =tokenizer(lowerCAmelCase__		,				return_tensors="pt"							)

        with torch.no_grad():
          a__     : str  =model(encoding.input_ids							).logits

        a__     : List[str]  =token_logits[:, 2, :].argmax(-1							)[0]

        self.assertEqual(tokenizer.decode(lowerCAmelCase__							)		,				"capital"							)






