

import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

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


_lowerCAmelCase  :      str			 =logging.get_logger(__name__)

_lowerCAmelCase  :      Tuple			 ='''▁'''

_lowerCAmelCase  :      int			 ={'''vocab_file''': '''sentencepiece.bpe.model'''}

_lowerCAmelCase  :      int			 ={
    '''vocab_file''': {
        '''xlm-roberta-base''': '''https://huggingface.co/xlm-roberta-base/resolve/main/sentencepiece.bpe.model''',
        '''xlm-roberta-large''': '''https://huggingface.co/xlm-roberta-large/resolve/main/sentencepiece.bpe.model''',
        '''xlm-roberta-large-finetuned-conll02-dutch''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-dutch/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll02-spanish''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll03-english''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-english/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll03-german''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-german/resolve/main/sentencepiece.bpe.model'''
        ),
    }
}

_lowerCAmelCase  :      Optional[Any]			 ={
    '''xlm-roberta-base''': 5_12,
    '''xlm-roberta-large''': 5_12,
    '''xlm-roberta-large-finetuned-conll02-dutch''': 5_12,
    '''xlm-roberta-large-finetuned-conll02-spanish''': 5_12,
    '''xlm-roberta-large-finetuned-conll03-english''': 5_12,
    '''xlm-roberta-large-finetuned-conll03-german''': 5_12,
}

class __UpperCamelCase     (				_a		):






					'''simple docstring'''



					__magic_name__        =   VOCAB_FILES_NAMES
					__magic_name__        =   PRETRAINED_VOCAB_FILES_MAP
					__magic_name__        =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					__magic_name__        =   ['''input_ids''', '''attention_mask''']





					def __init__(				self				,      lowerCamelCase__				,      lowerCamelCase__="<s>"				,      lowerCamelCase__="</s>"				,      lowerCamelCase__="</s>"				,      lowerCamelCase__="<s>"				,      lowerCamelCase__="<unk>"				,      lowerCamelCase__="<pad>"				,      lowerCamelCase__="<mask>"				,      lowerCamelCase__ = None				,      **lowerCamelCase__				,      ):
						UpperCAmelCase__:	int									= AddedToken(__lowerCAmelCase				,      lstrip=__lowerCAmelCase				,      rstrip=__lowerCAmelCase						) if isinstance(__lowerCAmelCase				,      __lowerCAmelCase						) else mask_token

						UpperCAmelCase__:	Tuple									= {} if sp_model_kwargs is None else sp_model_kwargs

						super().__init__(
						    bos_token=__lowerCAmelCase				,      eos_token=__lowerCAmelCase				,      unk_token=__lowerCAmelCase				,      sep_token=__lowerCAmelCase				,      cls_token=__lowerCAmelCase				,      pad_token=__lowerCAmelCase				,      mask_token=__lowerCAmelCase				,      sp_model_kwargs=self.sp_model_kwargs				,      **__lowerCAmelCase				,      )

						UpperCAmelCase__:	List[str]									= spm.SentencePieceProcessor(**self.sp_model_kwargs						)
						self.sp_model.Load(str(__lowerCAmelCase						)						)
						UpperCAmelCase__:	Tuple									= vocab_file

						# Original fairseq vocab and spm vocab must be "aligned":
						# Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
						# -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
						# fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
						# spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

						# Mimic fairseq token-to-id alignment for the first 4 token
						UpperCAmelCase__:	Tuple									= {"<s>": 0, "<pad>": 1, "</s>": 2, "<unk>": 3}

						# The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
						UpperCAmelCase__:	Dict									= 1

						UpperCAmelCase__:	int									= len(self.sp_model						) + self.fairseq_offset
						UpperCAmelCase__:	Tuple									= {v: k for k, v in self.fairseq_tokens_to_ids.items()}





					def __getstate__(				self						):
						UpperCAmelCase__:	Dict									= self.__dict__.copy()
						UpperCAmelCase__:	Any									= None
						UpperCAmelCase__:	Union[str, Any]									= self.sp_model.serialized_model_proto()
						return state





					def __setstate__(				self				,      lowerCamelCase__						):
						UpperCAmelCase__:	Any									= d

						# for backward compatibility
						if not hasattr(self				,      "sp_model_kwargs"						):
							UpperCAmelCase__:	Tuple									= {}

						UpperCAmelCase__:	Tuple									= spm.SentencePieceProcessor(**self.sp_model_kwargs						)
						self.sp_model.LoadFromSerializedProto(self.sp_model_proto						)





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__				,      lowerCamelCase__ = None						):






						if token_ids_a is None:
							return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
						UpperCAmelCase__:	Union[str, Any]									= [self.cls_token_id]
						UpperCAmelCase__:	List[Any]									= [self.sep_token_id]
						return cls + token_ids_a + sep + sep + token_ids_a + sep





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__				,      lowerCamelCase__ = None				,      lowerCamelCase__ = False						):






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

						if token_ids_a is None:
							return [1] + ([0] * len(__lowerCAmelCase						)) + [1]
						return [1] + ([0] * len(__lowerCAmelCase						)) + [1, 1] + ([0] * len(__lowerCAmelCase						)) + [1]





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__				,      lowerCamelCase__ = None						):
						UpperCAmelCase__:	Optional[Any]									= [self.sep_token_id]
						UpperCAmelCase__:	int									= [self.cls_token_id]

						if token_ids_a is None:
							return len(cls + token_ids_a + sep						) * [0]
						return len(cls + token_ids_a + sep + sep + token_ids_a + sep						) * [0]





					@property
					def    _UpperCAmelCase    (				self						):
						return len(self.sp_model						) + self.fairseq_offset + 1  # Add the <mask> token





					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	Optional[int]									= {self.convert_ids_to_tokens(__lowerCAmelCase						): i for i in range(self.vocab_size						)}
						vocab.update(self.added_tokens_encoder						)
						return vocab





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__						):
						return self.sp_model.encode(__lowerCAmelCase				,      out_type=__lowerCAmelCase						)





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__						):






						if token in self.fairseq_tokens_to_ids:
							return self.fairseq_tokens_to_ids[token]
						UpperCAmelCase__:	List[Any]									= self.sp_model.PieceToId(__lowerCAmelCase						)

						# Need to return unknown token if the SP model returned 0
						return spm_id + self.fairseq_offset if spm_id else self.unk_token_id





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__						):






						if index in self.fairseq_ids_to_tokens:
							return self.fairseq_ids_to_tokens[index]
						return self.sp_model.IdToPiece(index - self.fairseq_offset						)





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__						):
						UpperCAmelCase__:	List[str]									= "".join(__lowerCAmelCase						).replace(__lowerCAmelCase				,      " "						).strip()
						return out_string





					def    _UpperCAmelCase    (				self				,      lowerCamelCase__				,      lowerCamelCase__ = None						):






						if not os.path.isdir(__lowerCAmelCase						):
							logger.error(F"Vocabulary path ({save_directory}) should be a directory"						)
							return
						UpperCAmelCase__:	Dict									= os.path.join(
						    __lowerCAmelCase				,      (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]						)

						if os.path.abspath(self.vocab_file						) != os.path.abspath(__lowerCAmelCase						) and os.path.isfile(self.vocab_file						):
							copyfile(self.vocab_file				,      __lowerCAmelCase						)
						elif not os.path.isfile(self.vocab_file						):
							with open(__lowerCAmelCase				,      "wb"						) as fi:
								UpperCAmelCase__:	Tuple									= self.sp_model.serialized_model_proto()
								fi.write(__lowerCAmelCase						)

						return (out_vocab_file,)






import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextConfig, SegformerImageProcessor, UperNetConfig, UperNetForSemanticSegmentation







def   __lowerCamelCase   (				__a  :int       )						->							int:






		"""simple docstring"""

		A__     =  3_8_4
		if "tiny" in model_name:
				A__     =  [3, 3, 9, 3]
				A__     =  [9_6, 1_9_2, 3_8_4, 7_6_8]
		if "small" in model_name:
				A__     =  [3, 3, 2_7, 3]
				A__     =  [9_6, 1_9_2, 3_8_4, 7_6_8]
		if "base" in model_name:
				A__     =  [3, 3, 2_7, 3]
				A__     =  [1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
				A__     =  5_1_2
		if "large" in model_name:
				A__     =  [3, 3, 2_7, 3]
				A__     =  [1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
				A__     =  7_6_8
		if "xlarge" in model_name:
				A__     =  [3, 3, 2_7, 3]
				A__     =  [2_5_6, 5_1_2, 1_0_2_4, 2_0_4_8]
				A__     =  1_0_2_4

		# set label information
		A__     =  1_5_0
		A__     =  """huggingface/label-files"""
		A__     =  """ade20k-id2label.json"""
		A__     =  json.load(open(hf_hub_download(__a       ,	__a       ,	repo_type="""dataset"""       )       ,	"""r"""       )       )
		A__     =  {int(__a       ): v for k, v in idalabel.items()}
		A__     =  {v: k for k, v in idalabel.items()}

		A__     =  ConvNextConfig(
		    depths=__a       ,	hidden_sizes=__a       ,	out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]       )
		A__     =  UperNetConfig(
		    backbone_config=__a       ,	auxiliary_in_channels=__a       ,	num_labels=__a       ,	idalabel=__a       ,	labelaid=__a       ,	)

		return config







def   __lowerCamelCase   (				__a  :Union[str, Any]       )						->							List[Any]:






		"""simple docstring"""

		A__     =  []

		# fmt: off
		# stem
		rename_keys.append(("""backbone.downsample_layers.0.0.weight""", """backbone.embeddings.patch_embeddings.weight""")       )
		rename_keys.append(("""backbone.downsample_layers.0.0.bias""", """backbone.embeddings.patch_embeddings.bias""")       )
		rename_keys.append(("""backbone.downsample_layers.0.1.weight""", """backbone.embeddings.layernorm.weight""")       )
		rename_keys.append(("""backbone.downsample_layers.0.1.bias""", """backbone.embeddings.layernorm.bias""")       )
		# stages
		for i in range(len(config.backbone_config.depths       )       ):
				for j in range(config.backbone_config.depths[i]       ):
						rename_keys.append((F'backbone.stages.{i}.{j}.gamma', F'backbone.encoder.stages.{i}.layers.{j}.layer_scale_parameter')       )
						rename_keys.append((F'backbone.stages.{i}.{j}.depthwise_conv.weight', F'backbone.encoder.stages.{i}.layers.{j}.dwconv.weight')       )
						rename_keys.append((F'backbone.stages.{i}.{j}.depthwise_conv.bias', F'backbone.encoder.stages.{i}.layers.{j}.dwconv.bias')       )
						rename_keys.append((F'backbone.stages.{i}.{j}.norm.weight', F'backbone.encoder.stages.{i}.layers.{j}.layernorm.weight')       )
						rename_keys.append((F'backbone.stages.{i}.{j}.norm.bias', F'backbone.encoder.stages.{i}.layers.{j}.layernorm.bias')       )
						rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv1.weight', F'backbone.encoder.stages.{i}.layers.{j}.pwconv1.weight')       )
						rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv1.bias', F'backbone.encoder.stages.{i}.layers.{j}.pwconv1.bias')       )
						rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv2.weight', F'backbone.encoder.stages.{i}.layers.{j}.pwconv2.weight')       )
						rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv2.bias', F'backbone.encoder.stages.{i}.layers.{j}.pwconv2.bias')       )
				if i > 0:
						rename_keys.append((F'backbone.downsample_layers.{i}.0.weight', F'backbone.encoder.stages.{i}.downsampling_layer.0.weight')       )
						rename_keys.append((F'backbone.downsample_layers.{i}.0.bias', F'backbone.encoder.stages.{i}.downsampling_layer.0.bias')       )
						rename_keys.append((F'backbone.downsample_layers.{i}.1.weight', F'backbone.encoder.stages.{i}.downsampling_layer.1.weight')       )
						rename_keys.append((F'backbone.downsample_layers.{i}.1.bias', F'backbone.encoder.stages.{i}.downsampling_layer.1.bias')       )

				rename_keys.append((F'backbone.norm{i}.weight', F'backbone.hidden_states_norms.stage{i+1}.weight')       )
				rename_keys.append((F'backbone.norm{i}.bias', F'backbone.hidden_states_norms.stage{i+1}.bias')       )

		# decode head
		rename_keys.extend(
		    [
		        ("""decode_head.conv_seg.weight""", """decode_head.classifier.weight"""),
		        ("""decode_head.conv_seg.bias""", """decode_head.classifier.bias"""),
		        ("""auxiliary_head.conv_seg.weight""", """auxiliary_head.classifier.weight"""),
		        ("""auxiliary_head.conv_seg.bias""", """auxiliary_head.classifier.bias"""),
		    ]       )
		# fmt: on

		return rename_keys







def   __lowerCamelCase   (				__a  :Union[str, Any]       ,	__a  :List[str]       ,	__a  :str       )						->							str:






		"""simple docstring"""

		A__     =  dct.pop(__a       )
		A__     =  val







def   __lowerCamelCase   (				__a  :Any       ,	__a  :int       ,	__a  :Any       )						->							List[Any]:






		"""simple docstring"""

		A__     =  {
		    """upernet-convnext-tiny""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_tiny_fp16_512x512_160k_ade20k/upernet_convnext_tiny_fp16_512x512_160k_ade20k_20220227_124553-cad485de.pth""",
		    """upernet-convnext-small""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_small_fp16_512x512_160k_ade20k/upernet_convnext_small_fp16_512x512_160k_ade20k_20220227_131208-1b1e394f.pth""",
		    """upernet-convnext-base""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_base_fp16_512x512_160k_ade20k/upernet_convnext_base_fp16_512x512_160k_ade20k_20220227_181227-02a24fc6.pth""",
		    """upernet-convnext-large""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_large_fp16_640x640_160k_ade20k/upernet_convnext_large_fp16_640x640_160k_ade20k_20220226_040532-e57aa54d.pth""",
		    """upernet-convnext-xlarge""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_xlarge_fp16_640x640_160k_ade20k/upernet_convnext_xlarge_fp16_640x640_160k_ade20k_20220226_080344-95fc38c2.pth""",
		}
		A__     =  model_name_to_url[model_name]
		A__     =  torch.hub.load_state_dict_from_url(__a       ,	map_location="""cpu"""       )["""state_dict"""]

		A__     =  get_upernet_config(__a       )
		A__     =  UperNetForSemanticSegmentation(__a       )
		model.eval()

		# replace "bn" => "batch_norm"
		for key in state_dict.copy().keys():
				A__     =  state_dict.pop(__a       )
				if "bn" in key:
						A__     =  key.replace("""bn"""       ,	"""batch_norm"""       )
				A__     =  val

		# rename keys
		A__     =  create_rename_keys(__a       )
		for src, dest in rename_keys:
				rename_key(__a       ,	__a       ,	__a       )

		model.load_state_dict(__a       )

		# verify on image
		A__     =  """https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"""
		A__     =  Image.open(requests.get(__a       ,	stream=__a       ).raw       ).convert("""RGB"""       )

		A__     =  SegformerImageProcessor()
		A__     =  processor(__a       ,	return_tensors="""pt"""       ).pixel_values

		with torch.no_grad():
				A__     =  model(__a       )

		if model_name == "upernet-convnext-tiny":
				A__     =  torch.tensor(
				    [[-8.8110, -8.8110, -8.6521], [-8.8110, -8.8110, -8.6521], [-8.7746, -8.7746, -8.6130]]       )
		elif model_name == "upernet-convnext-small":
				A__     =  torch.tensor(
				    [[-8.8236, -8.8236, -8.6771], [-8.8236, -8.8236, -8.6771], [-8.7638, -8.7638, -8.6240]]       )
		elif model_name == "upernet-convnext-base":
				A__     =  torch.tensor(
				    [[-8.8558, -8.8558, -8.6905], [-8.8558, -8.8558, -8.6905], [-8.7669, -8.7669, -8.6021]]       )
		elif model_name == "upernet-convnext-large":
				A__     =  torch.tensor(
				    [[-8.6660, -8.6660, -8.6210], [-8.6660, -8.6660, -8.6210], [-8.6310, -8.6310, -8.5964]]       )
		elif model_name == "upernet-convnext-xlarge":
				A__     =  torch.tensor(
				    [[-8.4980, -8.4980, -8.3977], [-8.4980, -8.4980, -8.3977], [-8.4379, -8.4379, -8.3412]]       )
		print("""Logits:"""       ,	outputs.logits[0, 0, :3, :3]       )
		assert torch.allclose(outputs.logits[0, 0, :3, :3]       ,	__a       ,	atol=1E-4       )
		print("""Looks ok!"""       )

		if pytorch_dump_folder_path is not None:
				print(F'Saving model {model_name} to {pytorch_dump_folder_path}'       )
				model.save_pretrained(__a       )
				print(F'Saving processor to {pytorch_dump_folder_path}'       )
				processor.save_pretrained(__a       )

		if push_to_hub:
				print(F'Pushing model and processor for {model_name} to hub'       )
				model.push_to_hub(F'openmmlab/{model_name}'       )
				processor.push_to_hub(F'openmmlab/{model_name}'       )


if __name__ == "__main__":
					A       :   Optional[int]					= argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '''--model_name''',
					    default='''upernet-convnext-tiny''',
					    type=str,
					    choices=[F'''upernet-convnext-{size}''' for size in ['''tiny''', '''small''', '''base''', '''large''', '''xlarge''']],
					    help='''Name of the ConvNext UperNet model you\'d like to convert.''',
					)
					parser.add_argument(
					    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
					)
					parser.add_argument(
					    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
					)

					A       :   str					= parser.parse_args()
					convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)



