

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


snake_case_        =			[
    '''openmmlab/upernet-convnext-tiny''',
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
snake_case_        =			'''UperNetConfig'''


class 				SCREAMING_SNAKE_CASE__ (nn.Module					):





    def __init__(				self						,       a						,       a						,       a						,       a = 0						,       a = False						,       a = 1						,       ):
         super().__init__()
         lowercase__   :       Dict				=	nn.Convad(
             in_channels=__UpperCamelCase						,       out_channels=__UpperCamelCase						,       kernel_size=__UpperCamelCase						,       padding=__UpperCamelCase						,       bias=__UpperCamelCase						,       dilation=__UpperCamelCase						,       )
         lowercase__   :       Any				=	nn.BatchNormad(__UpperCamelCase)
         lowercase__   :       List[str]				=	nn.ReLU()





    def 				snake_case_    (				self						,       a):
         lowercase__   :       Optional[Any]				=	self.conv(__UpperCamelCase)
         lowercase__   :       Dict				=	self.batch_norm(__UpperCamelCase)
         lowercase__   :       Optional[int]				=	self.activation(__UpperCamelCase)

         return output



class 				SCREAMING_SNAKE_CASE__ (nn.Module					):





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





    def 				snake_case_    (				self						,       a):
         lowercase__   :       Union[str, Any]				=	input
         for layer in self.layers:
              lowercase__   :       Union[str, Any]				=	layer(__UpperCamelCase)
         return hidden_state



class 				SCREAMING_SNAKE_CASE__ (nn.Module					):





    def __init__(				self						,       a						,       a						,       a						,       a):
         super().__init__()
         lowercase__   :       Tuple				=	pool_scales
         lowercase__   :       Dict				=	align_corners
         lowercase__   :       Dict				=	in_channels
         lowercase__   :       Tuple				=	channels
         lowercase__   :       Any				=	[]
         for i, pool_scale in enumerate(__UpperCamelCase):
              lowercase__   :       Any				=	UperNetPyramidPoolingBlock(pool_scale=__UpperCamelCase						,       in_channels=__UpperCamelCase						,       channels=__UpperCamelCase)
              self.blocks.append(__UpperCamelCase)
              self.add_module(str(__UpperCamelCase)						,       __UpperCamelCase)





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



class 				SCREAMING_SNAKE_CASE__ (nn.Module					):





    def __init__(				self						,       a						,       a):
         super().__init__()

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

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

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





    def 				snake_case_    (				self):
         self.apply(self._init_weights)





    def 				snake_case_    (				self						,       a):
         if isinstance(__UpperCamelCase						,       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 				snake_case_    (				self						,       a):
         lowercase__   :       Optional[Any]				=	inputs[-1]
         lowercase__   :       List[Any]				=	[x]
         psp_outs.extend(self.psp_modules(__UpperCamelCase))
         lowercase__   :       Optional[int]				=	torch.cat(__UpperCamelCase						,       dim=1)
         lowercase__   :       Any				=	self.bottleneck(__UpperCamelCase)

         return output





    def 				snake_case_    (				self						,       a):
         # build laterals
         lowercase__   :       int				=	[lateral_conv(encoder_hidden_states[i]) for i, lateral_conv in enumerate(self.lateral_convs)]

         laterals.append(self.psp_forward(__UpperCamelCase))

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

         # build outputs
         lowercase__   :       Dict				=	[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):
              lowercase__   :       int				=	nn.functional.interpolate(
                  fpn_outs[i]						,       size=fpn_outs[0].shape[2:]						,       mode='bilinear'						,       align_corners=self.align_corners)
         lowercase__   :       Optional[int]				=	torch.cat(__UpperCamelCase						,       dim=1)
         lowercase__   :       int				=	self.fpn_bottleneck(__UpperCamelCase)
         lowercase__   :       List[str]				=	self.classifier(__UpperCamelCase)

         return output



class 				SCREAMING_SNAKE_CASE__ (nn.Module					):





    def __init__(				self						,       a						,       a = 2						,       a = 3						,       a = 1):
         super().__init__()

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

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

         lowercase__   :       str				=	nn.Convad(self.channels						,       config.num_labels						,       kernel_size=1)





    def 				snake_case_    (				self):
         self.apply(self._init_weights)





    def 				snake_case_    (				self						,       a):
         if isinstance(__UpperCamelCase						,       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 				snake_case_    (				self						,       a):
         # just take the relevant feature maps
         lowercase__   :       Dict				=	encoder_hidden_states[self.in_index]
         lowercase__   :       Dict				=	self.convs(__UpperCamelCase)
         if self.concat_input:
              lowercase__   :       Optional[Any]				=	self.conv_cat(torch.cat([hidden_states, output]						,       dim=1))
         lowercase__   :       Union[str, Any]				=	self.classifier(__UpperCamelCase)
         return output



class 				SCREAMING_SNAKE_CASE__ (_lowercase					):
    __lowerCamelCase						:	Optional[int]														=					UperNetConfig
    __lowerCamelCase						:	Optional[int]														=					"""pixel_values"""
    __lowerCamelCase						:	int														=					True





    def 				snake_case_    (				self						,       a):
         if isinstance(__UpperCamelCase						,       __UpperCamelCase):
              module.backbone.init_weights()
              module.decode_head.init_weights()
              module.auxiliary_head.init_weights()





    def 				snake_case_    (				self):
         self.backbone.init_weights()
         self.decode_head.init_weights()
         self.auxiliary_head.init_weights()





    def 				snake_case_    (				self						,       a						,       a=False):
         if isinstance(__UpperCamelCase						,       __UpperCamelCase):
              lowercase__   :       Optional[int]				=	value



snake_case_        =			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.
'''

snake_case_        =			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."""		,	_lowercase		,	)
class 				SCREAMING_SNAKE_CASE__ (_lowercase					):





    def __init__(				self						,       a):
         super().__init__(__UpperCamelCase)

         lowercase__   :       Tuple				=	AutoBackbone.from_config(config.backbone_config)

         # Semantic segmentation head(s)
         lowercase__   :       str				=	UperNetHead(__UpperCamelCase						,       in_channels=self.backbone.channels)
         lowercase__   :       Dict				=	UperNetFCNHead(__UpperCamelCase) 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=__UpperCamelCase						,       config_class=_CONFIG_FOR_DOC)
    def 				snake_case_    (				self						,       a = None						,       a = None						,       a = None						,       a = None						,       a = None						,       ):
         lowercase__   :       Optional[int]				=	return_dict if return_dict is not None else self.config.use_return_dict
         lowercase__   :       int				=	(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         lowercase__   :       Any				=	output_attentions if output_attentions is not None else self.config.output_attentions

         lowercase__   :       List[Any]				=	self.backbone.forward_with_filtered_kwargs(
             __UpperCamelCase						,       output_hidden_states=__UpperCamelCase						,       output_attentions=__UpperCamelCase)
         lowercase__   :       Union[str, Any]				=	outputs.feature_maps

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

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

         lowercase__   :       Optional[int]				=	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
                   lowercase__   :       Dict				=	CrossEntropyLoss(ignore_index=self.config.loss_ignore_index)
                   lowercase__   :       Dict				=	loss_fct(__UpperCamelCase						,       __UpperCamelCase)
                   lowercase__   :       List[Any]				=	loss_fct(__UpperCamelCase						,       __UpperCamelCase)
                   lowercase__   :       Optional[int]				=	main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

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

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




import json
import os
from typing import Dict, List, Optional, Tuple

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


snake_case_        =			logging.get_logger(__name__)


snake_case_        =			{
    '''vocab_file''': '''vocab.json''',
    '''tokenizer_config_file''': '''tokenizer_config.json''',
    '''merges_file''': '''merges.txt''',
}

snake_case_        =			{
    '''vocab_file''': {
        '''facebook/s2t-wav2vec2-large-en-de''': (
            '''https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/vocab.json'''
        ),
    },
    '''tokenizer_config_file''': {
        '''facebook/s2t-wav2vec2-large-en-de''': (
            '''https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/tokenizer_config.json'''
        ),
    },
    '''merges_file''': {
        '''facebook/s2t-wav2vec2-large-en-de''': (
            '''https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/merges.txt'''
        ),
    },
}

snake_case_        =			'''</w>'''
snake_case_        =			'''@@ '''

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   Optional[Any]      ):



     '''simple docstring'''





     lowercase__   :       Optional[Any]				=	set()
     lowercase__   :       Optional[Any]				=	word[0]
     for char in word[1:]:
          pairs.add((prev_char, char)      )
          lowercase__   :       Optional[int]				=	char
     return pairs


# Speech2Text2 has no max input length
snake_case_        =			{'''facebook/s2t-wav2vec2-large-en-de''': 1_024}


class 				SCREAMING_SNAKE_CASE__ (__snake_case					):
    __lowerCamelCase						:	List[str]														=					VOCAB_FILES_NAMES
    __lowerCamelCase						:	int														=					PRETRAINED_VOCAB_FILES_MAP
    __lowerCamelCase						:	Dict														=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    __lowerCamelCase						:	Union[str, Any]														=					["""input_ids""", """attention_mask"""]





    def __init__(				self						,       a						,       a="<s>"						,       a="<pad>"						,       a="</s>"						,       a="<unk>"						,       a=False						,       a=None						,       **a						,       ):
         super().__init__(
             unk_token=a						,       bos_token=a						,       eos_token=a						,       pad_token=a						,       do_lower_case=a						,       **a						,       )

         lowercase__   :       str				=	do_lower_case

         with open(a						,       encoding='utf-8') as vocab_handle:
              lowercase__   :       Tuple				=	json.load(a)
         lowercase__   :       Union[str, Any]				=	{v: k for k, v in self.encoder.items()}

         if merges_file is None:
              logger.info(f"""No merges files provided. {self.__class__.__name__} can only be used for decoding.""")

              lowercase__   :       int				=	None
              lowercase__   :       List[Any]				=	None
         else:
              with open(a						,       encoding='utf-8') as merges_handle:
                   lowercase__   :       List[Any]				=	merges_handle.read().split('\n')[:-1]

              lowercase__   :       Optional[int]				=	[tuple(merge.split()[:2]) for merge in merges]
              lowercase__   :       Tuple				=	dict(zip(a						,       range(len(a))))
              lowercase__   :       List[str]				=	{}





    @property
    def 				snake_case_    (				self):
         return len(self.decoder)





    def 				snake_case_    (				self):
         return dict(self.encoder						,       **self.added_tokens_encoder)





    def 				snake_case_    (				self						,       a):
         lowercase__   :       int				=	tuple(token[:-1]) + (token[-1] + BPE_TOKEN_MERGES,)
         if token in self.cache:
              return self.cache[token]
         lowercase__   :       Any				=	get_pairs(a)

         if not pairs:
              return token

         while True:
              lowercase__   :       List[str]				=	min(a						,       key=lambda a: self.bpe_ranks.get(a						,       float('inf')))
              if bigram not in self.bpe_ranks:
                   break
              lowercase__				,			lowercase__   :       Dict				=	bigram
              lowercase__   :       Union[str, Any]				=	[]
              lowercase__   :       int				=	0
              while i < len(a):
                   try:
                        lowercase__   :       Dict				=	word.index(a						,       a)
                   except ValueError:
                        new_word.extend(word[i:])
                        break
                   else:
                        new_word.extend(word[i:j])
                        lowercase__   :       Optional[int]				=	j

                   if word[i] == first and i < len(a) - 1 and word[i + 1] == second:
                        new_word.append(first + second)
                        i += 2
                   else:
                        new_word.append(word[i])
                        i += 1
              lowercase__   :       str				=	tuple(a)
              lowercase__   :       Union[str, Any]				=	new_word
              if len(a) == 1:
                   break
              else:
                   lowercase__   :       Optional[Any]				=	get_pairs(a)
         lowercase__   :       List[str]				=	' '.join(a)
         if word == "\n  " + BPE_TOKEN_MERGES:
              lowercase__   :       Optional[int]				=	'\n' + BPE_TOKEN_MERGES

         if word.endswith(a):
              lowercase__   :       Dict				=	word.replace(a						,       '')

         lowercase__   :       int				=	word.replace(' '						,       a)
         lowercase__   :       List[str]				=	word
         return word





    def 				snake_case_    (				self						,       a):

         if self.bpe_ranks is None:
              raise ValueError(
                  'This tokenizer was instantiated without a `merges.txt` file, so'
                  ' that it can only be used for decoding, not for encoding.'
                  'Make sure to provide `merges.txt` file at instantiation to enable '
                  'encoding.')

         if self.do_lower_case:
              lowercase__   :       int				=	text.lower()

         lowercase__   :       Optional[int]				=	text.split()

         lowercase__   :       Optional[int]				=	[]
         for token in text:
              if token:
                   split_tokens.extend(list(self.bpe(a).split(' ')))

         return split_tokens





    def 				snake_case_    (				self						,       a):
         return self.encoder.get(a						,       self.encoder.get(self.unk_token))





    def 				snake_case_    (				self						,       a):
         lowercase__   :       Union[str, Any]				=	self.decoder.get(a						,       self.unk_token)
         return result





    def 				snake_case_    (				self						,       a):
         lowercase__   :       Union[str, Any]				=	' '.join(a)

         # make sure @@ tokens are concatenated
         lowercase__   :       Optional[int]				=	''.join(string.split(a))

         return string





    def 				snake_case_    (				self						,       a						,       a = None):
         if not os.path.isdir(a):
              logger.error(f"""Vocabulary path ({save_directory}) should be a directory""")
              return
         lowercase__   :       Optional[Any]				=	os.path.join(
             a						,       (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'])
         lowercase__   :       List[Any]				=	os.path.join(
             a						,       (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['merges_file'])

         with open(a						,       'w'						,       encoding='utf-8') as f:
              f.write(json.dumps(self.encoder						,       indent=2						,       sort_keys=a						,       ensure_ascii=a) + '\n')

         lowercase__   :       Optional[Any]				=	0
         if self.bpe_ranks is None:
              return (vocab_file,)

         with open(a						,       'w'						,       encoding='utf-8') as writer:
              for bpe_tokens, token_index in sorted(self.bpe_ranks.items()						,       key=lambda a: kv[1]):
                   if index != token_index:
                        logger.warning(
                            f"""Saving vocabulary to {merges_file}: BPE merge indices are not consecutive."""
                            ' Please check that the tokenizer is not corrupted!')
                        lowercase__   :       Dict				=	token_index
                   writer.write(' '.join(a) + '\n')
                   index += 1

         return (vocab_file, merges_file)

