


'''simple docstring'''



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				(_SCREAMING_SNAKE_CASE :  Optional[Any]  ):
					__a     :			List[str]          =    384
					if "tiny" in model_name:
										__a     :			List[str]          =    [3, 3, 9, 3]
										__a     :			Optional[int]          =    [96, 192, 384, 768]
					if "small" in model_name:
										__a     :			Optional[int]          =    [3, 3, 27, 3]
										__a     :			Optional[int]          =    [96, 192, 384, 768]
					if "base" in model_name:
										__a     :			List[Any]          =    [3, 3, 27, 3]
										__a     :			Optional[int]          =    [128, 256, 512, 1_024]
										__a     :			List[str]          =    512
					if "large" in model_name:
										__a     :			Optional[Any]          =    [3, 3, 27, 3]
										__a     :			List[str]          =    [192, 384, 768, 1_536]
										__a     :			List[str]          =    768
					if "xlarge" in model_name:
										__a     :			str          =    [3, 3, 27, 3]
										__a     :			List[str]          =    [256, 512, 1_024, 2_048]
										__a     :			Any          =    1_024

					# set label information
					__a     :			int          =    150
					__a     :			List[Any]          =    'huggingface/label-files'
					__a     :			Tuple          =    'ade20k-id2label.json'
					__a     :			List[str]          =    json.load(open(hf_hub_download(UpperCamelCase__		, UpperCamelCase__		, repo_type='dataset'  )		, 'r'  )  )
					__a     :			List[Any]          =    {int(UpperCamelCase__  ): v for k, v in idalabel.items()}
					__a     :			Optional[Any]          =    {v: k for k, v in idalabel.items()}

					__a     :			List[str]          =    ConvNextConfig(
					    depths=UpperCamelCase__		, hidden_sizes=UpperCamelCase__		, out_features=['stage1', 'stage2', 'stage3', 'stage4']  )
					__a     :			Union[str, Any]          =    UperNetConfig(
					    backbone_config=UpperCamelCase__		, auxiliary_in_channels=UpperCamelCase__		, num_labels=UpperCamelCase__		, idalabel=UpperCamelCase__		, labelaid=UpperCamelCase__		, )

					return config






def     lowerCamelCase				(_SCREAMING_SNAKE_CASE :  int  ):
					__a     :			Dict          =    []

					# 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				(_SCREAMING_SNAKE_CASE :  Tuple		, _SCREAMING_SNAKE_CASE :  List[str]		, _SCREAMING_SNAKE_CASE :  int  ):
					__a     :			int          =    dct.pop(UpperCamelCase__  )
					__a     :			str          =    val






def     lowerCamelCase				(_SCREAMING_SNAKE_CASE :  Optional[int]		, _SCREAMING_SNAKE_CASE :  Any		, _SCREAMING_SNAKE_CASE :  int  ):
					__a     :			Union[str, Any]          =    {
					    '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     :			Dict          =    model_name_to_url[model_name]
					__a     :			Optional[int]          =    torch.hub.load_state_dict_from_url(UpperCamelCase__		, map_location='cpu'  )['state_dict']

					__a     :			Tuple          =    get_upernet_config(UpperCamelCase__  )
					__a     :			str          =    UperNetForSemanticSegmentation(UpperCamelCase__  )
					model.eval()

					# replace "bn" => "batch_norm"
					for key in state_dict.copy().keys():
										__a     :			str          =    state_dict.pop(UpperCamelCase__  )
										if "bn" in key:
															__a     :			Tuple          =    key.replace('bn'		, 'batch_norm'  )
										__a     :			Any          =    val

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

					model.load_state_dict(UpperCamelCase__  )

					# verify on image
					__a     :			Union[str, Any]          =    'https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg'
					__a     :			List[Any]          =    Image.open(requests.get(UpperCamelCase__		, stream=UpperCamelCase__  ).raw  ).convert('RGB'  )

					__a     :			int          =    SegformerImageProcessor()
					__a     :			Tuple          =    processor(UpperCamelCase__		, return_tensors='pt'  ).pixel_values

					with torch.no_grad():
										__a     :			List[Any]          =    model(UpperCamelCase__  )

					if model_name == "upernet-convnext-tiny":
										__a     :			Tuple          =    torch.tensor(
										    [[-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.7_7_4_6, -8.7_7_4_6, -8.6_1_3_0]]  )
					elif model_name == "upernet-convnext-small":
										__a     :			Tuple          =    torch.tensor(
										    [[-8.8_2_3_6, -8.8_2_3_6, -8.6_7_7_1], [-8.8_2_3_6, -8.8_2_3_6, -8.6_7_7_1], [-8.7_6_3_8, -8.7_6_3_8, -8.6_2_4_0]]  )
					elif model_name == "upernet-convnext-base":
										__a     :			Union[str, Any]          =    torch.tensor(
										    [[-8.8_5_5_8, -8.8_5_5_8, -8.6_9_0_5], [-8.8_5_5_8, -8.8_5_5_8, -8.6_9_0_5], [-8.7_6_6_9, -8.7_6_6_9, -8.6_0_2_1]]  )
					elif model_name == "upernet-convnext-large":
										__a     :			Dict          =    torch.tensor(
										    [[-8.6_6_6_0, -8.6_6_6_0, -8.6_2_1_0], [-8.6_6_6_0, -8.6_6_6_0, -8.6_2_1_0], [-8.6_3_1_0, -8.6_3_1_0, -8.5_9_6_4]]  )
					elif model_name == "upernet-convnext-xlarge":
										__a     :			Dict          =    torch.tensor(
										    [[-8.4_9_8_0, -8.4_9_8_0, -8.3_9_7_7], [-8.4_9_8_0, -8.4_9_8_0, -8.3_9_7_7], [-8.4_3_7_9, -8.4_3_7_9, -8.3_4_1_2]]  )
					print('Logits:'		, outputs.logits[0, 0, :3, :3]  )
					assert torch.allclose(outputs.logits[0, 0, :3, :3]		, UpperCamelCase__		, 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(UpperCamelCase__  )
										print(F"""Saving processor to {pytorch_dump_folder_path}"""  )
										processor.save_pretrained(UpperCamelCase__  )

					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__":
	__lowercase      :		List[str]			=      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.'
	)

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







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

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
				from .tokenization_albert import AlbertTokenizer
else:
				_lowercase							:   List[str]						=						None

_lowercase							:   List[str]						=						logging.get_logger(__name__)
_lowercase							:   List[Any]						=						{"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}

_lowercase							:   str						=						{
    "vocab_file": {
        "albert-base-v1": "https://huggingface.co/albert-base-v1/resolve/main/spiece.model",
        "albert-large-v1": "https://huggingface.co/albert-large-v1/resolve/main/spiece.model",
        "albert-xlarge-v1": "https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model",
        "albert-xxlarge-v1": "https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model",
        "albert-base-v2": "https://huggingface.co/albert-base-v2/resolve/main/spiece.model",
        "albert-large-v2": "https://huggingface.co/albert-large-v2/resolve/main/spiece.model",
        "albert-xlarge-v2": "https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model",
        "albert-xxlarge-v2": "https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model",
    },
    "tokenizer_file": {
        "albert-base-v1": "https://huggingface.co/albert-base-v1/resolve/main/tokenizer.json",
        "albert-large-v1": "https://huggingface.co/albert-large-v1/resolve/main/tokenizer.json",
        "albert-xlarge-v1": "https://huggingface.co/albert-xlarge-v1/resolve/main/tokenizer.json",
        "albert-xxlarge-v1": "https://huggingface.co/albert-xxlarge-v1/resolve/main/tokenizer.json",
        "albert-base-v2": "https://huggingface.co/albert-base-v2/resolve/main/tokenizer.json",
        "albert-large-v2": "https://huggingface.co/albert-large-v2/resolve/main/tokenizer.json",
        "albert-xlarge-v2": "https://huggingface.co/albert-xlarge-v2/resolve/main/tokenizer.json",
        "albert-xxlarge-v2": "https://huggingface.co/albert-xxlarge-v2/resolve/main/tokenizer.json",
    },
}

_lowercase							:   List[str]						=						{
    "albert-base-v1": 512,
    "albert-large-v1": 512,
    "albert-xlarge-v1": 512,
    "albert-xxlarge-v1": 512,
    "albert-base-v2": 512,
    "albert-large-v2": 512,
    "albert-xlarge-v2": 512,
    "albert-xxlarge-v2": 512,
}

_lowercase							:   Dict						=						"▁"






class    _UpperCamelCase     (       __snake_case					):

				"""simple docstring"""







				lowerCAmelCase					=			VOCAB_FILES_NAMES
				lowerCAmelCase					=			PRETRAINED_VOCAB_FILES_MAP
				lowerCAmelCase					=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				lowerCAmelCase					=			AlbertTokenizer

				def __init__(    self					, a__=None					, a__=None					, a__=True					, a__=True					, a__=False					, a__="[CLS]"					, a__="[SEP]"					, a__="<unk>"					, a__="[SEP]"					, a__="<pad>"					, a__="[CLS]"					, a__="[MASK]"					, **a__					, )						->					Dict:
					# Mask token behave like a normal word, i.e. include the space before it and
					# is included in the raw text, there should be a match in a non-normalized sentence.
					A             =					(
					    AddedToken(a__					, lstrip=a__					, rstrip=a__					, normalized=a__ )
					    if isinstance(a__					, a__ )
					    else mask_token
					)

					super().__init__(
					    a__					, tokenizer_file=a__					, do_lower_case=a__					, remove_space=a__					, keep_accents=a__					, bos_token=a__					, eos_token=a__					, unk_token=a__					, sep_token=a__					, pad_token=a__					, cls_token=a__					, mask_token=a__					, **a__					, )

					A             =					do_lower_case
					A             =					remove_space
					A             =					keep_accents
					A             =					vocab_file
					A             =					False if not self.vocab_file else True

				def   _UpperCAmelCase       (    self					, a__					, a__ = None )						->					List[int]:
					A             =					[self.sep_token_id]
					A             =					[self.cls_token_id]
					if token_ids_a is None:
						return cls + token_ids_a + sep
					return cls + token_ids_a + sep + token_ids_a + sep

				def   _UpperCAmelCase       (    self					, a__					, a__ = None )						->					List[int]:
					A             =					[self.sep_token_id]
					A             =					[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 ) * [0] + len(token_ids_a + sep ) * [1]

				def   _UpperCAmelCase       (    self					, a__					, a__ = None )						->					Tuple[str]:
					if not self.can_save_slow_tokenizer:
						raise ValueError(
						    """Your fast tokenizer does not have the necessary information to save the vocabulary for a slow """
						    """tokenizer.""" )

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

					if os.path.abspath(self.vocab_file ) != os.path.abspath(a__ ):
						copyfile(self.vocab_file					, a__ )

					return (out_vocab_file,)



