


"""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


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

# General docstring
lowercase_			     =  "UperNetConfig"
class 					__lowerCAmelCase					(   nn.Module       ):


		'''simple docstring'''



		def __init__(  self							, _a							, _a							, _a							, _a = 0							, _a = False							, _a = 1							, ):
				super().__init__()
				__a           =			nn.Convad(
				    in_channels=_a							, out_channels=_a							, kernel_size=_a							, padding=_a							, bias=_a							, dilation=_a							, )
				__a           =			nn.BatchNormad(_a      )
				__a           =			nn.ReLU()





		def 			__UpperCAmelCase       (  self							, _a      ):
				__a           =			self.conv(_a      )
				__a           =			self.batch_norm(_a      )
				__a           =			self.activation(_a      )

				return output







class 					__lowerCAmelCase					(   nn.Module       ):


		'''simple docstring'''



		def __init__(  self							, _a							, _a							, _a      ):
				super().__init__()
				__a           =			[
				    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							, _a      ):
				__a           =			input
				for layer in self.layers:
						__a           =			layer(_a      )
				return hidden_state







class 					__lowerCAmelCase					(   nn.Module       ):


		'''simple docstring'''



		def __init__(  self							, _a							, _a							, _a							, _a      ):
				super().__init__()
				__a           =			pool_scales
				__a           =			align_corners
				__a           =			in_channels
				__a           =			channels
				__a           =			[]
				for i, pool_scale in enumerate(_a      ):
						__a           =			UperNetPyramidPoolingBlock(pool_scale=_a							, in_channels=_a							, channels=_a      )
						self.blocks.append(_a      )
						self.add_module(str(_a      )							, _a      )





		def 			__UpperCAmelCase       (  self							, _a      ):
				__a           =			[]
				for ppm in self.blocks:
						__a           =			ppm(_a      )
						__a           =			nn.functional.interpolate(
						    _a							, size=x.size()[2:]							, mode='''bilinear'''							, align_corners=self.align_corners      )
						ppm_outs.append(_a      )
				return ppm_outs







class 					__lowerCAmelCase					(   nn.Module       ):


		'''simple docstring'''



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

				__a           =			config
				__a           =			config.pool_scales  # e.g. (1, 2, 3, 6)
				__a           =			in_channels
				__a           =			config.hidden_size
				__a           =			False
				__a           =			nn.Convad(self.channels							, config.num_labels							, kernel_size=1      )

				# PSP Module
				__a           =			UperNetPyramidPoolingModule(
				    self.pool_scales							, self.in_channels[-1]							, self.channels							, align_corners=self.align_corners							, )
				__a           =			UperNetConvModule(
				    self.in_channels[-1] + len(self.pool_scales      ) * self.channels							, self.channels							, kernel_size=3							, padding=1							, )
				# FPN Module
				__a           =			nn.ModuleList()
				__a           =			nn.ModuleList()
				for in_channels in self.in_channels[:-1]:  # skip the top layer
						__a           =			UperNetConvModule(_a							, self.channels							, kernel_size=1      )
						__a           =			UperNetConvModule(self.channels							, self.channels							, kernel_size=3							, padding=1      )
						self.lateral_convs.append(_a      )
						self.fpn_convs.append(_a      )

				__a           =			UperNetConvModule(
				    len(self.in_channels      ) * self.channels							, self.channels							, kernel_size=3							, padding=1							, )

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

		def 			__UpperCAmelCase       (  self							, _a      ):
				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							, _a      ):
				__a           =			inputs[-1]
				__a           =			[x]
				psp_outs.extend(self.psp_modules(_a      )      )
				__a           =			torch.cat(_a							, dim=1      )
				__a           =			self.bottleneck(_a      )

				return output





		def 			__UpperCAmelCase       (  self							, _a      ):
				# build laterals
				__a           =			[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
				__a           =			len(_a      )
				for i in range(used_backbone_levels - 1							, 0							, -1      ):
						__a           =			laterals[i - 1].shape[2:]
						__a           =			laterals[i - 1] + nn.functional.interpolate(
						    laterals[i]							, size=_a							, mode='''bilinear'''							, align_corners=self.align_corners      )

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

				return output







class 					__lowerCAmelCase					(   nn.Module       ):


		'''simple docstring'''



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

				__a           =			config
				__a           =			config.auxiliary_in_channels
				__a           =			config.auxiliary_channels
				__a           =			config.auxiliary_num_convs
				__a           =			config.auxiliary_concat_input
				__a           =			in_index

				__a           =			(kernel_size // 2) * dilation
				__a           =			[]
				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:
						__a           =			nn.Identity()
				else:
						__a           =			nn.Sequential(*_a      )
				if self.concat_input:
						__a           =			UperNetConvModule(
						    self.in_channels + self.channels							, self.channels							, kernel_size=_a							, padding=kernel_size // 2      )

				__a           =			nn.Convad(self.channels							, config.num_labels							, kernel_size=1      )

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

		def 			__UpperCAmelCase       (  self							, _a      ):
				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							, _a      ):
				# just take the relevant feature maps
				__a           =			encoder_hidden_states[self.in_index]
				__a           =			self.convs(_a      )
				if self.concat_input:
						__a           =			self.conv_cat(torch.cat([hidden_states, output]							, dim=1      )      )
				__a           =			self.classifier(_a      )
				return output







class 					__lowerCAmelCase					(   __SCREAMING_SNAKE_CASE       ):


		'''simple docstring'''



		__UpperCAmelCase		:			int										=				UperNetConfig
		__UpperCAmelCase		:			Union[str, Any]										=				'pixel_values'
		__UpperCAmelCase		:			Dict										=				True

		def 			__UpperCAmelCase       (  self							, _a      ):
				if isinstance(_a							, _a      ):
						module.backbone.init_weights()
						module.decode_head.init_weights()
						module.auxiliary_head.init_weights()

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





		def 			__UpperCAmelCase       (  self							, _a							, _a=False      ):
				if isinstance(_a							, _a      ):
						__a           =			value







lowercase_			     =  r"\n    Parameters:\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n        config ([`UperNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

lowercase_			     =  r"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using\n            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.\n        output_attentions (`bool`, *optional*):\n            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See\n            `attentions` under returned tensors for more detail.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers of the backbone. See `hidden_states` under\n            returned tensors for more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"
@add_start_docstrings(
    'UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes.'     ,			__SCREAMING_SNAKE_CASE     ,			)
class 					__lowerCAmelCase					(   __SCREAMING_SNAKE_CASE       ):


		'''simple docstring'''



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

				__a           =			AutoBackbone.from_config(config.backbone_config      )

				# Semantic segmentation head(s)
				__a           =			UperNetHead(_a							, in_channels=self.backbone.channels      )
				__a           =			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							, _a = None							, _a = None							, _a = None							, _a = None							, _a = None							, ):
				__a           =			return_dict if return_dict is not None else self.config.use_return_dict
				__a           =			(
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				__a           =			output_attentions if output_attentions is not None else self.config.output_attentions

				__a           =			self.backbone.forward_with_filtered_kwargs(
				    _a							, output_hidden_states=_a							, output_attentions=_a      )
				__a           =			outputs.feature_maps

				__a           =			self.decode_head(_a      )
				__a           =			nn.functional.interpolate(_a							, size=pixel_values.shape[2:]							, mode='''bilinear'''							, align_corners=_a      )

				__a           =			None
				if self.auxiliary_head is not None:
						__a           =			self.auxiliary_head(_a      )
						__a           =			nn.functional.interpolate(
						    _a							, size=pixel_values.shape[2:]							, mode='''bilinear'''							, align_corners=_a      )

				__a           =			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
								__a           =			CrossEntropyLoss(ignore_index=self.config.loss_ignore_index      )
								__a           =			loss_fct(_a							, _a      )
								__a           =			loss_fct(_a							, _a      )
								__a           =			main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

				if not return_dict:
						if output_hidden_states:
								__a           =			(logits,) + outputs[1:]
						else:
								__a           =			(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'''





import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

from ...tokenization_utils import BatchEncoding, PreTrainedTokenizer
from ...utils import logging


lowercase_		       =  logging.get_logger(__name__)

lowercase_		       =  "▁"

lowercase_		       =  {
    "vocab_file": "vocab.json",
    "spm_file": "sentencepiece.bpe.model",
    "tokenizer_config_file": "tokenizer_config.json",
}

lowercase_		       =  {
    "vocab_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/vocab.json",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/vocab.json",
    },
    "spm_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/sentencepiece.bpe.model",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/sentencepiece.bpe.model",
    },
    "tokenizer_config_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/tokenizer_config.json",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/tokenizer_config.json",
    },
}

lowercase_		       =  {
    "facebook/m2m100_418M": 1_024,
}

# fmt: off
lowercase_		       =  {
    "m2m100": ["af", "am", "ar", "ast", "az", "ba", "be", "bg", "bn", "br", "bs", "ca", "ceb", "cs", "cy", "da", "de", "el", "en", "es", "et", "fa", "ff", "fi", "fr", "fy", "ga", "gd", "gl", "gu", "ha", "he", "hi", "hr", "ht", "hu", "hy", "id", "ig", "ilo", "is", "it", "ja", "jv", "ka", "kk", "km", "kn", "ko", "lb", "lg", "ln", "lo", "lt", "lv", "mg", "mk", "ml", "mn", "mr", "ms", "my", "ne", "nl", "no", "ns", "oc", "or", "pa", "pl", "ps", "pt", "ro", "ru", "sd", "si", "sk", "sl", "so", "sq", "sr", "ss", "su", "sv", "sw", "ta", "th", "tl", "tn", "tr", "uk", "ur", "uz", "vi", "wo", "xh", "yi", "yo", "zh", "zu"],
    "wmt21": ["en", "ha", "is", "ja", "cs", "ru", "zh", "de"]
}






class __A			(  A						):



 '''simple docstring'''







 __lowerCamelCase    :       Union[str, Any]          =       VOCAB_FILES_NAMES
 __lowerCamelCase    :       Tuple          =       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCamelCase    :       Tuple          =       PRETRAINED_VOCAB_FILES_MAP
 __lowerCamelCase    :       Dict          =       ['input_ids', 'attention_mask']

 __lowerCamelCase    :       List[int]          =       []
 __lowerCamelCase    :       List[int]          =       []



 def __init__(self ,   A ,   A ,   A=None ,   A=None ,   A="<s>" ,   A="</s>" ,   A="</s>" ,   A="<pad>" ,   A="<unk>" ,   A="m2m100" ,   A = None ,   A=8 ,   **A ,   )    ->     None:


      """simple docstring"""



      _a      =  {} if sp_model_kwargs is None else sp_model_kwargs

      _a      =  language_codes
      _a      =  FAIRSEQ_LANGUAGE_CODES[language_codes]
      _a      =  {lang_code: f'''__{lang_code}__''' for lang_code in fairseq_language_code}

      _a      =  kwargs.get('''additional_special_tokens''' ,   []       )
      kwargs["additional_special_tokens"] += [
          self.get_lang_token(A       )
          for lang_code in fairseq_language_code
          if self.get_lang_token(A       ) not in kwargs["additional_special_tokens"]
      ]

      super().__init__(
          src_lang=A ,   tgt_lang=A ,   bos_token=A ,   eos_token=A ,   sep_token=A ,   unk_token=A ,   pad_token=A ,   language_codes=A ,   sp_model_kwargs=self.sp_model_kwargs ,   num_madeup_words=A ,   **A ,   )

      _a      =  vocab_file
      _a      =  load_json(A       )
      _a      =  {v: k for k, v in self.encoder.items()}
      _a      =  spm_file
      _a      =  load_spm(A ,   self.sp_model_kwargs       )

      _a      =  len(self.encoder       )

      _a      =  {
          self.get_lang_token(A       ): self.encoder_size + i for i, lang_code in enumerate(A       )
      }
      _a      =  {lang_code: self.encoder_size + i for i, lang_code in enumerate(A       )}
      _a      =  {v: k for k, v in self.lang_token_to_id.items()}

      _a      =  src_lang if src_lang is not None else '''en'''
      _a      =  tgt_lang
      _a      =  self.get_lang_id(self._src_lang       )
      self.set_src_lang_special_tokens(self._src_lang       )

      _a      =  num_madeup_words



 @property
 def 					a__  (self       )    ->     int:


      """simple docstring"""



      return len(self.encoder       ) + len(self.lang_token_to_id       )



 @property
 def 					a__  (self       )    ->     str:


      """simple docstring"""



      return self._src_lang



 @src_lang.setter
 def 					a__  (self ,   A       )    ->     None:


      """simple docstring"""



      _a      =  new_src_lang
      self.set_src_lang_special_tokens(self._src_lang       )



 def 					a__  (self ,   A       )    ->     List[str]:


      """simple docstring"""



      return self.sp_model.encode(A ,   out_type=A       )



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


      """simple docstring"""



      if token in self.lang_token_to_id:
           return self.lang_token_to_id[token]
      return self.encoder.get(A ,   self.encoder[self.unk_token]       )



 def 					a__  (self ,   A       )    ->     str:


      """simple docstring"""



      if index in self.id_to_lang_token:
           return self.id_to_lang_token[index]
      return self.decoder.get(A ,   self.unk_token       )



 def 					a__  (self ,   A       )    ->     Dict:


      """simple docstring"""



      _a      =  []
      _a      =  ''''''
      for token in tokens:
           # make sure that special tokens are not decoded using sentencepiece model
           if token in self.all_special_tokens:
                out_string += self.sp_model.decode(A       ) + token
                _a      =  []
           else:
                current_sub_tokens.append(A       )
      out_string += self.sp_model.decode(A       )
      return out_string.strip()



 def 					a__  (self ,   A ,   A = None ,   A = False       )    ->     List[int]:


      """simple docstring"""



      if already_has_special_tokens:
           return super().get_special_tokens_mask(
               token_ids_a=A ,   token_ids_a=A ,   already_has_special_tokens=A       )

      _a      =  [1] * len(self.prefix_tokens       )
      _a      =  [1] * len(self.suffix_tokens       )
      if token_ids_a is None:
           return prefix_ones + ([0] * len(A       )) + suffix_ones
      return prefix_ones + ([0] * len(A       )) + ([0] * len(A       )) + suffix_ones



 def 					a__  (self ,   A ,   A = None       )    ->     List[int]:


      """simple docstring"""



      if token_ids_a is None:
           return self.prefix_tokens + token_ids_a + self.suffix_tokens
      # We don't expect to process pairs, but leave the pair logic for API consistency
      return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens



 def 					a__  (self       )    ->     Dict:


      """simple docstring"""



      _a      =  {self.convert_ids_to_tokens(A       ): i for i in range(self.vocab_size       )}
      vocab.update(self.added_tokens_encoder       )
      return vocab



 def __getstate__(self       )    ->     Dict:


      """simple docstring"""



      _a      =  self.__dict__.copy()
      _a      =  None
      return state



 def __setstate__(self ,   A       )    ->     None:


      """simple docstring"""



      _a      =  d

      # for backward compatibility
      if not hasattr(self ,   '''sp_model_kwargs'''       ):
           _a      =  {}

      _a      =  load_spm(self.spm_file ,   self.sp_model_kwargs       )



 def 					a__  (self ,   A ,   A = None       )    ->     Tuple[str]:


      """simple docstring"""



      _a      =  Path(A       )
      if not save_dir.is_dir():
           raise OSError(f'''{save_directory} should be a directory'''       )
      _a      =  save_dir / (
          (filename_prefix + '''-''' if filename_prefix else '''''') + self.vocab_files_names['''vocab_file''']
      )
      _a      =  save_dir / (
          (filename_prefix + '''-''' if filename_prefix else '''''') + self.vocab_files_names['''spm_file''']
      )

      save_json(self.encoder ,   A       )

      if os.path.abspath(self.spm_file       ) != os.path.abspath(A       ) and os.path.isfile(self.spm_file       ):
           copyfile(self.spm_file ,   A       )
      elif not os.path.isfile(self.spm_file       ):
           with open(A ,   '''wb'''       ) as fi:
                _a      =  self.sp_model.serialized_model_proto()
                fi.write(A       )

      return (str(A       ), str(A       ))



 def 					a__  (self ,   A ,   A = "en" ,   A = None ,   A = "ro" ,   **A ,   )    ->     BatchEncoding:


      """simple docstring"""



      _a      =  src_lang
      _a      =  tgt_lang
      self.set_src_lang_special_tokens(self.src_lang       )
      return super().prepare_seqaseq_batch(A ,   A ,   **A       )



 def 					a__  (self ,   A ,   A ,   A ,   **A       )    ->     Union[str, Any]:


      """simple docstring"""



      if src_lang is None or tgt_lang is None:
           raise ValueError('''Translation requires a `src_lang` and a `tgt_lang` for this model'''       )
      _a      =  src_lang
      _a      =  self(A ,   add_special_tokens=A ,   **A       )
      _a      =  self.get_lang_id(A       )
      _a      =  tgt_lang_id
      return inputs



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


      """simple docstring"""



      self.set_src_lang_special_tokens(self.src_lang       )



 def 					a__  (self       )    ->     Tuple:


      """simple docstring"""



      self.set_tgt_lang_special_tokens(self.tgt_lang       )



 def 					a__  (self ,   A       )    ->     None:


      """simple docstring"""



      _a      =  self.get_lang_token(A       )
      _a      =  self.lang_token_to_id[lang_token]
      _a      =  [self.cur_lang_id]
      _a      =  [self.eos_token_id]



 def 					a__  (self ,   A       )    ->     None:


      """simple docstring"""



      _a      =  self.get_lang_token(A       )
      _a      =  self.lang_token_to_id[lang_token]
      _a      =  [self.cur_lang_id]
      _a      =  [self.eos_token_id]



 def 					a__  (self ,   A       )    ->     str:


      """simple docstring"""



      return self.lang_code_to_token[lang]



 def 					a__  (self ,   A       )    ->     int:


      """simple docstring"""



      _a      =  self.get_lang_token(A       )
      return self.lang_token_to_id[lang_token]





def lowerCAmelCase  (__A  ,	__A):



     """simple docstring"""



     _a      =  sentencepiece.SentencePieceProcessor(**__A)
     spm.Load(str(__A))
     return spm




def lowerCAmelCase  (__A):



     """simple docstring"""



     with open(__A  ,	'''r''') as f:
          return json.load(__A)




def lowerCAmelCase  (__A  ,	__A):



     """simple docstring"""



     with open(__A  ,	'''w''') as f:
          json.dump(__A  ,	__A  ,	indent=2)



