






"""simple docstring"""





from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from ... import AutoBackbone
from ...modeling_outputs import SemanticSegmenterOutput
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
from ...utils.backbone_utils import BackboneMixin
from .configuration_upernet import UperNetConfig


UpperCamelCase_	:				List[str]										= [
    """openmmlab/upernet-convnext-tiny""",
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
UpperCamelCase_	:				List[str]										= """UperNetConfig"""


class      _lowercase (     nn.Module						):


      def __init__(    self				:	Any							,   a				:	int							,   a				:	int							,   a				:	Union[int, Tuple[int, int]]							,   a				:	Union[int, Tuple[int, int], str] = 0							,   a				:	bool = False							,   a				:	Union[int, Tuple[int, int]] = 1							,   ):





            """simple docstring"""







            super().__init__()
            __snake_case      :    int  			=nn.Convad(
                in_channels=a							,   out_channels=a							,   kernel_size=a							,   padding=a							,   bias=a							,   dilation=a							,   )
            __snake_case      :    Union[str, Any]  			=nn.BatchNormad(a						)
            __snake_case      :    List[str]  			=nn.ReLU()







      def _UpperCamelCase						(    self				:	str							,   a				:	torch.Tensor						):





            """simple docstring"""







            __snake_case      :    Union[str, Any]  			=self.conv(a						)
            __snake_case      :    Any  			=self.batch_norm(a						)
            __snake_case      :    Union[str, Any]  			=self.activation(a						)

            return output



class      _lowercase (     nn.Module						):


      def __init__(    self				:	Dict							,   a				:	int							,   a				:	int							,   a				:	int						):





            """simple docstring"""







            super().__init__()
            __snake_case      :    str  			=[
                nn.AdaptiveAvgPoolad(a						),
                UperNetConvModule(a							,   a							,   kernel_size=1						),
            ]
            for i, layer in enumerate(self.layers						):
                  self.add_module(str(a						)							,   a						)







      def _UpperCamelCase						(    self				:	Dict							,   a				:	torch.Tensor						):





            """simple docstring"""







            __snake_case      :    int  			=input
            for layer in self.layers:
                  __snake_case      :    int  			=layer(a						)
            return hidden_state



class      _lowercase (     nn.Module						):


      def __init__(    self				:	int							,   a				:	Tuple[int, ...]							,   a				:	int							,   a				:	int							,   a				:	bool						):





            """simple docstring"""







            super().__init__()
            __snake_case      :    Tuple  			=pool_scales
            __snake_case      :    Optional[Any]  			=align_corners
            __snake_case      :    List[str]  			=in_channels
            __snake_case      :    Any  			=channels
            __snake_case      :    str  			=[]
            for i, pool_scale in enumerate(a						):
                  __snake_case      :    Optional[int]  			=UperNetPyramidPoolingBlock(pool_scale=a							,   in_channels=a							,   channels=a						)
                  self.blocks.append(a						)
                  self.add_module(str(a						)							,   a						)







      def _UpperCamelCase						(    self				:	int							,   a				:	torch.Tensor						):





            """simple docstring"""







            __snake_case      :    List[str]  			=[]
            for ppm in self.blocks:
                  __snake_case      :    List[str]  			=ppm(a						)
                  __snake_case      :    Tuple  			=nn.functional.interpolate(
                      a							,   size=x.size()[2:]							,   mode='''bilinear'''							,   align_corners=self.align_corners						)
                  ppm_outs.append(a						)
            return ppm_outs



class      _lowercase (     nn.Module						):


      def __init__(    self				:	Optional[Any]							,   a				:	str							,   a				:	int						):





            """simple docstring"""







            super().__init__()

            __snake_case      :    Optional[int]  			=config
            __snake_case      :    Dict  			=config.pool_scales  # e.g. (1, 2, 3, 6)
            __snake_case      :    int  			=in_channels
            __snake_case      :    int  			=config.hidden_size
            __snake_case      :    Optional[int]  			=False
            __snake_case      :    List[Any]  			=nn.Convad(self.channels							,   config.num_labels							,   kernel_size=1						)

            # PSP Module
            __snake_case      :    Optional[int]  			=UperNetPyramidPoolingModule(
                self.pool_scales							,   self.in_channels[-1]							,   self.channels							,   align_corners=self.align_corners							,   )
            __snake_case      :    Any  			=UperNetConvModule(
                self.in_channels[-1] + len(self.pool_scales						) * self.channels							,   self.channels							,   kernel_size=3							,   padding=1							,   )
            # FPN Module
            __snake_case      :    Tuple  			=nn.ModuleList()
            __snake_case      :    Optional[int]  			=nn.ModuleList()
            for in_channels in self.in_channels[:-1]:  # skip the top layer
                  __snake_case      :    Optional[int]  			=UperNetConvModule(a							,   self.channels							,   kernel_size=1						)
                  __snake_case      :    List[str]  			=UperNetConvModule(self.channels							,   self.channels							,   kernel_size=3							,   padding=1						)
                  self.lateral_convs.append(a						)
                  self.fpn_convs.append(a						)

            __snake_case      :    str  			=UperNetConvModule(
                len(self.in_channels						) * self.channels							,   self.channels							,   kernel_size=3							,   padding=1							,   )


      def _UpperCamelCase						(    self				:	Any						):





            """simple docstring"""







            self.apply(self._init_weights						)


      def _UpperCamelCase						(    self				:	List[str]							,   a				:	List[Any]						):





            """simple docstring"""







            if isinstance(a							,   nn.Convad						):
                  module.weight.data.normal_(mean=0.0							,   std=self.config.initializer_range						)
                  if module.bias is not None:
                        module.bias.data.zero_()


      def _UpperCamelCase						(    self				:	Tuple							,   a				:	List[str]						):





            """simple docstring"""







            __snake_case      :    Any  			=inputs[-1]
            __snake_case      :    str  			=[x]
            psp_outs.extend(self.psp_modules(a						)						)
            __snake_case      :    str  			=torch.cat(a							,   dim=1						)
            __snake_case      :    str  			=self.bottleneck(a						)

            return output







      def _UpperCamelCase						(    self				:	Optional[Any]							,   a				:	torch.Tensor						):





            """simple docstring"""







            __snake_case      :    Optional[Any]  			=[lateral_conv(encoder_hidden_states[i]						) for i, lateral_conv in enumerate(self.lateral_convs						)]

            laterals.append(self.psp_forward(a						)						)

            # build top-down path
            __snake_case      :    Optional[int]  			=len(a						)
            for i in range(used_backbone_levels - 1							,   0							,   -1						):
                  __snake_case      :    List[Any]  			=laterals[i - 1].shape[2:]
                  __snake_case      :    Optional[int]  			=laterals[i - 1] + nn.functional.interpolate(
                      laterals[i]							,   size=a							,   mode='''bilinear'''							,   align_corners=self.align_corners						)

            # build outputs
            __snake_case      :    List[Any]  			=[self.fpn_convs[i](laterals[i]						) for i in range(used_backbone_levels - 1						)]
            # append psp feature
            fpn_outs.append(laterals[-1]						)

            for i in range(used_backbone_levels - 1							,   0							,   -1						):
                  __snake_case      :    Optional[Any]  			=nn.functional.interpolate(
                      fpn_outs[i]							,   size=fpn_outs[0].shape[2:]							,   mode='''bilinear'''							,   align_corners=self.align_corners						)
            __snake_case      :    Dict  			=torch.cat(a							,   dim=1						)
            __snake_case      :    List[Any]  			=self.fpn_bottleneck(a						)
            __snake_case      :    int  			=self.classifier(a						)

            return output



class      _lowercase (     nn.Module						):


      def __init__(    self				:	Tuple							,   a				:	List[str]							,   a				:	int = 2							,   a				:	int = 3							,   a				:	Union[int, Tuple[int, int]] = 1						):





            """simple docstring"""







            super().__init__()

            __snake_case      :    Optional[Any]  			=config
            __snake_case      :    str  			=config.auxiliary_in_channels
            __snake_case      :    Any  			=config.auxiliary_channels
            __snake_case      :    int  			=config.auxiliary_num_convs
            __snake_case      :    Optional[int]  			=config.auxiliary_concat_input
            __snake_case      :    int  			=in_index

            __snake_case      :    Union[str, Any]  			=(kernel_size // 2) * dilation
            __snake_case      :    int  			=[]
            convs.append(
                UperNetConvModule(
                    self.in_channels							,   self.channels							,   kernel_size=a							,   padding=a							,   dilation=a						)						)
            for i in range(self.num_convs - 1						):
                  convs.append(
                      UperNetConvModule(
                          self.channels							,   self.channels							,   kernel_size=a							,   padding=a							,   dilation=a						)						)
            if self.num_convs == 0:
                  __snake_case      :    Any  			=nn.Identity()
            else:
                  __snake_case      :    Dict  			=nn.Sequential(*a						)
            if self.concat_input:
                  __snake_case      :    List[Any]  			=UperNetConvModule(
                      self.in_channels + self.channels							,   self.channels							,   kernel_size=a							,   padding=kernel_size // 2						)

            __snake_case      :    Optional[int]  			=nn.Convad(self.channels							,   config.num_labels							,   kernel_size=1						)


      def _UpperCamelCase						(    self				:	Any						):





            """simple docstring"""







            self.apply(self._init_weights						)


      def _UpperCamelCase						(    self				:	int							,   a				:	int						):





            """simple docstring"""







            if isinstance(a							,   nn.Convad						):
                  module.weight.data.normal_(mean=0.0							,   std=self.config.initializer_range						)
                  if module.bias is not None:
                        module.bias.data.zero_()







      def _UpperCamelCase						(    self				:	str							,   a				:	torch.Tensor						):





            """simple docstring"""







            __snake_case      :    Any  			=encoder_hidden_states[self.in_index]
            __snake_case      :    Dict  			=self.convs(a						)
            if self.concat_input:
                  __snake_case      :    Dict  			=self.conv_cat(torch.cat([hidden_states, output]							,   dim=1						)						)
            __snake_case      :    Union[str, Any]  			=self.classifier(a						)
            return output



class      _lowercase (     lowerCAmelCase						):
      _a   :     Dict								=   UperNetConfig
      _a   :     List[Any]								=   '''pixel_values'''
      _a   :     Optional[int]								=   True


      def _UpperCamelCase						(    self				:	Any							,   a				:	str						):





            """simple docstring"""







            if isinstance(a							,   a						):
                  module.backbone.init_weights()
                  module.decode_head.init_weights()
                  module.auxiliary_head.init_weights()


      def _UpperCamelCase						(    self				:	str						):





            """simple docstring"""







            self.backbone.init_weights()
            self.decode_head.init_weights()
            self.auxiliary_head.init_weights()







      def _UpperCamelCase						(    self				:	Tuple							,   a				:	Union[str, Any]							,   a				:	List[str]=False						):





            """simple docstring"""







            if isinstance(a							,   a						):
                  __snake_case      :    List[str]  			=value



UpperCamelCase_	:				Optional[int]										= r"""
    Parameters:
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.
        config ([`UperNetConfig`]): 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_	:				str										= r"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See
            `attentions` under returned tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers of the backbone. See `hidden_states` under
            returned tensors for more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""


@add_start_docstrings(
    '''UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes.'''		,  lowerCAmelCase		,  )
class      _lowercase (     lowerCAmelCase						):


      def __init__(    self				:	Optional[Any]							,   a				:	str						):





            """simple docstring"""







            super().__init__(a						)

            __snake_case      :    List[Any]  			=AutoBackbone.from_config(config.backbone_config						)

            # Semantic segmentation head(s)
            __snake_case      :    Optional[Any]  			=UperNetHead(a							,   in_channels=self.backbone.channels						)
            __snake_case      :    str  			=UperNetFCNHead(a						) if config.use_auxiliary_head else None

            # Initialize weights and apply final processing
            self.post_init()







      @add_start_docstrings_to_model_forward(UPERNET_INPUTS_DOCSTRING.format('''batch_size, sequence_length'''						)						)
      @replace_return_docstrings(output_type=a							,   config_class=_CONFIG_FOR_DOC						)
      def _UpperCamelCase						(    self				:	str							,   a				:	Optional[torch.Tensor] = None							,   a				:	Optional[bool] = None							,   a				:	Optional[bool] = None							,   a				:	Optional[torch.Tensor] = None							,   a				:	Optional[bool] = None							,   ):





            """simple docstring"""







            __snake_case      :    Tuple  			=return_dict if return_dict is not None else self.config.use_return_dict
            __snake_case      :    Optional[Any]  			=(
                output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
            )
            __snake_case      :    Dict  			=output_attentions if output_attentions is not None else self.config.output_attentions

            __snake_case      :    Union[str, Any]  			=self.backbone.forward_with_filtered_kwargs(
                a							,   output_hidden_states=a							,   output_attentions=a						)
            __snake_case      :    Tuple  			=outputs.feature_maps

            __snake_case      :    List[str]  			=self.decode_head(a						)
            __snake_case      :    Tuple  			=nn.functional.interpolate(a							,   size=pixel_values.shape[2:]							,   mode='''bilinear'''							,   align_corners=a						)

            __snake_case      :    List[Any]  			=None
            if self.auxiliary_head is not None:
                  __snake_case      :    List[str]  			=self.auxiliary_head(a						)
                  __snake_case      :    int  			=nn.functional.interpolate(
                      a							,   size=pixel_values.shape[2:]							,   mode='''bilinear'''							,   align_corners=a						)

            __snake_case      :    Union[str, Any]  			=None
            if labels is not None:
                  if self.config.num_labels == 1:
                        raise ValueError('''The number of labels should be greater than one'''						)
                  else:
                        # compute weighted loss
                        __snake_case      :    Any  			=CrossEntropyLoss(ignore_index=self.config.loss_ignore_index						)
                        __snake_case      :    int  			=loss_fct(a							,   a						)
                        __snake_case      :    Optional[Any]  			=loss_fct(a							,   a						)
                        __snake_case      :    Any  			=main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

            if not return_dict:
                  if output_hidden_states:
                        __snake_case      :    Any  			=(logits,) + outputs[1:]
                  else:
                        __snake_case      :    str  			=(logits,) + outputs[2:]
                  return ((loss,) + output) if loss is not None else output

            return SemanticSegmenterOutput(
                loss=a							,   logits=a							,   hidden_states=outputs.hidden_states							,   attentions=outputs.attentions							,   )










"""simple docstring"""





from typing import TYPE_CHECKING

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


UpperCamelCase_	:				Any										= {"""configuration_reformer""": ["""REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ReformerConfig"""]}

try:
    if not is_sentencepiece_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    UpperCamelCase_	:				List[str]										= ["""ReformerTokenizer"""]

try:
    if not is_tokenizers_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    UpperCamelCase_	:				List[str]										= ["""ReformerTokenizerFast"""]

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    UpperCamelCase_	:				Optional[int]										= [
        """REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """ReformerAttention""",
        """ReformerForMaskedLM""",
        """ReformerForQuestionAnswering""",
        """ReformerForSequenceClassification""",
        """ReformerLayer""",
        """ReformerModel""",
        """ReformerModelWithLMHead""",
        """ReformerPreTrainedModel""",
    ]


if TYPE_CHECKING:
    from .configuration_reformer import REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, ReformerConfig

    try:
        if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_reformer import ReformerTokenizer

    try:
        if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_reformer_fast import ReformerTokenizerFast

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_reformer import (
            REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
            ReformerAttention,
            ReformerForMaskedLM,
            ReformerForQuestionAnswering,
            ReformerForSequenceClassification,
            ReformerLayer,
            ReformerModel,
            ReformerModelWithLMHead,
            ReformerPreTrainedModel,
        )

else:
    import sys

    UpperCamelCase_	:				Tuple										= _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


