






import argparse
import json
import os
import pickle
import shutil

import numpy as np
import torch
from distiller import Distiller
from lm_seqs_dataset import LmSeqsDataset

from transformers import (
    BertConfig,
    BertForMaskedLM,
    BertTokenizer,
    DistilBertConfig,
    DistilBertForMaskedLM,
    DistilBertTokenizer,
    GPTaConfig,
    GPTaLMHeadModel,
    GPTaTokenizer,
    RobertaConfig,
    RobertaForMaskedLM,
    RobertaTokenizer,
)
from utils import git_log, init_gpu_params, logger, set_seed


__lowerCamelCase   :     List[Any] 		= {
    """distilbert""": (DistilBertConfig, DistilBertForMaskedLM, DistilBertTokenizer),
    """roberta""": (RobertaConfig, RobertaForMaskedLM, RobertaTokenizer),
    """bert""": (BertConfig, BertForMaskedLM, BertTokenizer),
    """gpt2""": (GPTaConfig, GPTaLMHeadModel, GPTaTokenizer),
}







def 	SCREAMING_SNAKE_CASE							(					snake_case_   :				Union[str, Any]  ):
	assert (args.mlm and args.alpha_mlm > 0.0) or (not args.mlm and args.alpha_mlm == 0.0)
	assert (args.alpha_mlm > 0.0 and args.alpha_clm == 0.0) or (args.alpha_mlm == 0.0 and args.alpha_clm > 0.0)
	if args.mlm:
		assert os.path.isfile(args.token_counts  )
		assert (args.student_type in ["roberta", "distilbert"]) and (args.teacher_type in ["roberta", "bert"])
	else:
		assert (args.student_type in ["gpt2"]) and (args.teacher_type in ["gpt2"])

	assert args.teacher_type == args.student_type or (
	    args.student_type == "distilbert" and args.teacher_type == "bert"
	)
	assert os.path.isfile(args.student_config  )
	if args.student_pretrained_weights is not None:
		assert os.path.isfile(args.student_pretrained_weights  )

	if args.freeze_token_type_embds:
		assert args.student_type in ["roberta"]

	assert args.alpha_ce >= 0.0
	assert args.alpha_mlm >= 0.0
	assert args.alpha_clm >= 0.0
	assert args.alpha_mse >= 0.0
	assert args.alpha_cos >= 0.0
	assert args.alpha_ce + args.alpha_mlm + args.alpha_clm + args.alpha_mse + args.alpha_cos > 0.0







def 	SCREAMING_SNAKE_CASE							(					snake_case_   :				List[str]	,    snake_case_   :				Any  ):
	if args.student_type == "roberta":
		snake_case__					: Dict  =    False
	elif args.student_type == "gpt2":
		snake_case__					: Optional[Any]  =    False







def 	SCREAMING_SNAKE_CASE							(					snake_case_   :				Tuple	,    snake_case_   :				List[str]  ):
	if args.student_type == "roberta":
		snake_case__					: Tuple  =    False







def 	SCREAMING_SNAKE_CASE							(					):
	snake_case__					: str  =    argparse.ArgumentParser(description="Training"  )
	parser.add_argument("--force"	,    action="store_true"	,    help="Overwrite dump_path if it already exists."  )

	parser.add_argument(
	    "--dump_path"	,    type=snake_case_	,    required=snake_case_	,    help="The output directory (log, checkpoints, parameters, etc.)"  )
	parser.add_argument(
	    "--data_file"	,    type=snake_case_	,    required=snake_case_	,    help="The binarized file (tokenized + tokens_to_ids) and grouped by sequence."	,    )

	parser.add_argument(
	    "--student_type"	,    type=snake_case_	,    choices=["distilbert", "roberta", "gpt2"]	,    required=snake_case_	,    help="The student type (DistilBERT, RoBERTa)."	,    )
	parser.add_argument("--student_config"	,    type=snake_case_	,    required=snake_case_	,    help="Path to the student configuration."  )
	parser.add_argument(
	    "--student_pretrained_weights"	,    default=snake_case_	,    type=snake_case_	,    help="Load student initialization checkpoint."  )

	parser.add_argument(
	    "--teacher_type"	,    choices=["bert", "roberta", "gpt2"]	,    required=snake_case_	,    help="Teacher type (BERT, RoBERTa)."  )
	parser.add_argument("--teacher_name"	,    type=snake_case_	,    required=snake_case_	,    help="The teacher model."  )

	parser.add_argument("--temperature"	,    default=2.0	,    type=snake_case_	,    help="Temperature for the softmax temperature."  )
	parser.add_argument(
	    "--alpha_ce"	,    default=0.5	,    type=snake_case_	,    help="Linear weight for the distillation loss. Must be >=0."  )
	parser.add_argument(
	    "--alpha_mlm"	,    default=0.0	,    type=snake_case_	,    help="Linear weight for the MLM loss. Must be >=0. Should be used in conjunction with `mlm` flag."	,    )
	parser.add_argument("--alpha_clm"	,    default=0.5	,    type=snake_case_	,    help="Linear weight for the CLM loss. Must be >=0."  )
	parser.add_argument("--alpha_mse"	,    default=0.0	,    type=snake_case_	,    help="Linear weight of the MSE loss. Must be >=0."  )
	parser.add_argument(
	    "--alpha_cos"	,    default=0.0	,    type=snake_case_	,    help="Linear weight of the cosine embedding loss. Must be >=0."  )

	parser.add_argument(
	    "--mlm"	,    action="store_true"	,    help="The LM step: MLM or CLM. If `mlm` is True, the MLM is used over CLM."  )
	parser.add_argument(
	    "--mlm_mask_prop"	,    default=0.15	,    type=snake_case_	,    help="Proportion of tokens for which we need to make a prediction."	,    )
	parser.add_argument("--word_mask"	,    default=0.8	,    type=snake_case_	,    help="Proportion of tokens to mask out."  )
	parser.add_argument("--word_keep"	,    default=0.1	,    type=snake_case_	,    help="Proportion of tokens to keep."  )
	parser.add_argument("--word_rand"	,    default=0.1	,    type=snake_case_	,    help="Proportion of tokens to randomly replace."  )
	parser.add_argument(
	    "--mlm_smoothing"	,    default=0.7	,    type=snake_case_	,    help="Smoothing parameter to emphasize more rare tokens (see XLM, similar to word2vec)."	,    )
	parser.add_argument("--token_counts"	,    type=snake_case_	,    help="The token counts in the data_file for MLM."  )

	parser.add_argument(
	    "--restrict_ce_to_mask"	,    action="store_true"	,    help="If true, compute the distillation loss only the [MLM] prediction distribution."	,    )
	parser.add_argument(
	    "--freeze_pos_embs"	,    action="store_true"	,    help="Freeze positional embeddings during distillation. For student_type in [\'roberta\', \'gpt2\'] only."	,    )
	parser.add_argument(
	    "--freeze_token_type_embds"	,    action="store_true"	,    help="Freeze token type embeddings during distillation if existent. For student_type in [\'roberta\'] only."	,    )

	parser.add_argument("--n_epoch"	,    type=snake_case_	,    default=3	,    help="Number of pass on the whole dataset."  )
	parser.add_argument("--batch_size"	,    type=snake_case_	,    default=5	,    help="Batch size (for each process)."  )
	parser.add_argument(
	    "--group_by_size"	,    action="store_false"	,    help="If true, group sequences that have similar length into the same batch. Default is true."	,    )

	parser.add_argument(
	    "--gradient_accumulation_steps"	,    type=snake_case_	,    default=50	,    help="Gradient accumulation for larger training batches."	,    )
	parser.add_argument("--warmup_prop"	,    default=0.05	,    type=snake_case_	,    help="Linear warmup proportion."  )
	parser.add_argument("--weight_decay"	,    default=0.0	,    type=snake_case_	,    help="Weight decay if we apply some."  )
	parser.add_argument("--learning_rate"	,    default=5E-4	,    type=snake_case_	,    help="The initial learning rate for Adam."  )
	parser.add_argument("--adam_epsilon"	,    default=1E-6	,    type=snake_case_	,    help="Epsilon for Adam optimizer."  )
	parser.add_argument("--max_grad_norm"	,    default=5.0	,    type=snake_case_	,    help="Max gradient norm."  )
	parser.add_argument("--initializer_range"	,    default=0.02	,    type=snake_case_	,    help="Random initialization range."  )

	parser.add_argument(
	    "--fp16"	,    action="store_true"	,    help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"	,    )
	parser.add_argument(
	    "--fp16_opt_level"	,    type=snake_case_	,    default="O1"	,    help=(
	        "For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\']."
	        "See details at https://nvidia.github.io/apex/amp.html"
	    )	,    )
	parser.add_argument("--n_gpu"	,    type=snake_case_	,    default=1	,    help="Number of GPUs in the node."  )
	parser.add_argument("--local_rank"	,    type=snake_case_	,    default=-1	,    help="Distributed training - Local rank"  )
	parser.add_argument("--seed"	,    type=snake_case_	,    default=56	,    help="Random seed"  )

	parser.add_argument("--log_interval"	,    type=snake_case_	,    default=500	,    help="Tensorboard logging interval."  )
	parser.add_argument("--checkpoint_interval"	,    type=snake_case_	,    default=4000	,    help="Checkpoint interval."  )
	snake_case__					: Any  =    parser.parse_args()
	sanity_checks(snake_case_  )

	# ARGS #
	init_gpu_params(snake_case_  )
	set_seed(snake_case_  )
	if args.is_master:
		if os.path.exists(args.dump_path  ):
			if not args.force:
				raise ValueError(
				    F'''Serialization dir {args.dump_path} already exists, but you have not precised wheter to overwrite'''
				    " itUse `--force` if you want to overwrite it"  )
			else:
				shutil.rmtree(args.dump_path  )

		if not os.path.exists(args.dump_path  ):
			os.makedirs(args.dump_path  )
		logger.info(F'''Experiment will be dumped and logged in {args.dump_path}'''  )

		# SAVE PARAMS #
		logger.info(F'''Param: {args}'''  )
		with open(os.path.join(args.dump_path	,    "parameters.json"  )	,    "w"  ) as f:
			json.dump(vars(snake_case_  )	,    snake_case_	,    indent=4  )
		git_log(args.dump_path  )

	snake_case__,		snake_case__,		snake_case__					: int  =    MODEL_CLASSES[args.student_type]
	snake_case__,		snake_case__,		snake_case__					: Tuple  =    MODEL_CLASSES[args.teacher_type]

	# TOKENIZER #
	snake_case__					: Optional[Any]  =    teacher_tokenizer_class.from_pretrained(args.teacher_name  )
	snake_case__					: List[str]  =    {}
	for tok_name, tok_symbol in tokenizer.special_tokens_map.items():
		snake_case__					: int  =    tokenizer.all_special_tokens.index(snake_case_  )
		snake_case__					: Dict  =    tokenizer.all_special_ids[idx]
	logger.info(F'''Special tokens {special_tok_ids}'''  )
	snake_case__					: int  =    special_tok_ids
	snake_case__					: Any  =    tokenizer.max_model_input_sizes[args.teacher_name]

	# DATA LOADER #
	logger.info(F'''Loading data from {args.data_file}'''  )
	with open(args.data_file	,    "rb"  ) as fp:
		snake_case__					: Tuple  =    pickle.load(snake_case_  )

	if args.mlm:
		logger.info(F'''Loading token counts from {args.token_counts} (already pre-computed)'''  )
		with open(args.token_counts	,    "rb"  ) as fp:
			snake_case__					: Optional[int]  =    pickle.load(snake_case_  )

		snake_case__					: Tuple  =    np.maximum(snake_case_	,    1  ) ** -args.mlm_smoothing
		for idx in special_tok_ids.values():
			snake_case__					: Optional[int]  =    0.0  # do not predict special tokens
		snake_case__					: Dict  =    torch.from_numpy(snake_case_  )
	else:
		snake_case__					: Optional[Any]  =    None

	snake_case__					: List[Any]  =    LmSeqsDataset(params=snake_case_	,    data=snake_case_  )
	logger.info("Data loader created."  )

	# STUDENT #
	logger.info(F'''Loading student config from {args.student_config}'''  )
	snake_case__					: Tuple  =    student_config_class.from_pretrained(args.student_config  )
	snake_case__					: Union[str, Any]  =    True

	if args.student_pretrained_weights is not None:
		logger.info(F'''Loading pretrained weights from {args.student_pretrained_weights}'''  )
		snake_case__					: List[Any]  =    student_model_class.from_pretrained(args.student_pretrained_weights	,    config=snake_case_  )
	else:
		snake_case__					: Union[str, Any]  =    student_model_class(snake_case_  )

	if args.n_gpu > 0:
		student.to(F'''cuda:{args.local_rank}'''  )
	logger.info("Student loaded."  )

	# TEACHER #
	snake_case__					: Optional[int]  =    teacher_model_class.from_pretrained(args.teacher_name	,    output_hidden_states=snake_case_  )
	if args.n_gpu > 0:
		teacher.to(F'''cuda:{args.local_rank}'''  )
	logger.info(F'''Teacher loaded from {args.teacher_name}.'''  )

	# FREEZING #
	if args.freeze_pos_embs:
		freeze_pos_embeddings(snake_case_	,    snake_case_  )
	if args.freeze_token_type_embds:
		freeze_token_type_embeddings(snake_case_	,    snake_case_  )

	# SANITY CHECKS #
	assert student.config.vocab_size == teacher.config.vocab_size
	assert student.config.hidden_size == teacher.config.hidden_size
	assert student.config.max_position_embeddings == teacher.config.max_position_embeddings
	if args.mlm:
		assert token_probs.size(0  ) == stu_architecture_config.vocab_size

	# DISTILLER #
	torch.cuda.empty_cache()
	snake_case__					: Tuple  =    Distiller(
	    params=snake_case_	,    dataset=snake_case_	,    token_probs=snake_case_	,    student=snake_case_	,    teacher=snake_case_  )
	distiller.train()
	logger.info("Let\'s go get some drinks."  )


if __name__ == "__main__":
	main()




import argparse
import json

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

from transformers import SegformerImageProcessor, SwinConfig, UperNetConfig, UperNetForSemanticSegmentation


def 	lowerCamelCase					(	UpperCamelCase     :				List[str]							)   ->    Optional[Any]:
		_lowerCamelCase				   =  3_84
		_lowerCamelCase				   =  7
		if "tiny" in model_name:
				_lowerCamelCase				   =  96
				_lowerCamelCase				   =  (2, 2, 6, 2)
				_lowerCamelCase				   =  (3, 6, 12, 24)
		elif "small" in model_name:
				_lowerCamelCase				   =  96
				_lowerCamelCase				   =  (2, 2, 18, 2)
				_lowerCamelCase				   =  (3, 6, 12, 24)
		elif "base" in model_name:
				_lowerCamelCase				   =  1_28
				_lowerCamelCase				   =  (2, 2, 18, 2)
				_lowerCamelCase				   =  (4, 8, 16, 32)
				_lowerCamelCase				   =  12
				_lowerCamelCase				   =  5_12
		elif "large" in model_name:
				_lowerCamelCase				   =  1_92
				_lowerCamelCase				   =  (2, 2, 18, 2)
				_lowerCamelCase				   =  (6, 12, 24, 48)
				_lowerCamelCase				   =  12
				_lowerCamelCase				   =  7_68

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

		_lowerCamelCase				   =  SwinConfig(
		    embed_dim=UpperCamelCase			,       depths=UpperCamelCase			,       num_heads=UpperCamelCase			,       window_size=UpperCamelCase			,       out_features=['stage1', 'stage2', 'stage3', 'stage4']			,       )
		_lowerCamelCase				   =  UperNetConfig(
		    backbone_config=UpperCamelCase			,       auxiliary_in_channels=UpperCamelCase			,       num_labels=UpperCamelCase			,       idalabel=UpperCamelCase			,       labelaid=UpperCamelCase			,       )

		return config


def 	lowerCamelCase					(	UpperCamelCase     :				Union[str, Any]							)   ->    int:
		_lowerCamelCase				   =  []

		# fmt: off
		# stem
		rename_keys.append(('backbone.patch_embed.projection.weight', 'backbone.embeddings.patch_embeddings.projection.weight')							)
		rename_keys.append(('backbone.patch_embed.projection.bias', 'backbone.embeddings.patch_embeddings.projection.bias')							)
		rename_keys.append(('backbone.patch_embed.norm.weight', 'backbone.embeddings.norm.weight')							)
		rename_keys.append(('backbone.patch_embed.norm.bias', 'backbone.embeddings.norm.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}.blocks.{j}.norm1.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.weight""")							)
						rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.norm1.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.bias""")							)
						rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_bias_table""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table""")							)
						rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_index""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index""")							)
						rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight""")							)
						rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias""")							)
						rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.norm2.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.weight""")							)
						rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.norm2.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.bias""")							)
						rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight""")							)
						rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias""")							)
						rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.1.weight""", F"""backbone.encoder.layers.{i}.blocks.{j}.output.dense.weight""")							)
						rename_keys.append((F"""backbone.stages.{i}.blocks.{j}.ffn.layers.1.bias""", F"""backbone.encoder.layers.{i}.blocks.{j}.output.dense.bias""")							)

				if i < 3:
						rename_keys.append((F"""backbone.stages.{i}.downsample.reduction.weight""", F"""backbone.encoder.layers.{i}.downsample.reduction.weight""")							)
						rename_keys.append((F"""backbone.stages.{i}.downsample.norm.weight""", F"""backbone.encoder.layers.{i}.downsample.norm.weight""")							)
						rename_keys.append((F"""backbone.stages.{i}.downsample.norm.bias""", F"""backbone.encoder.layers.{i}.downsample.norm.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					(	UpperCamelCase     :				str			,       UpperCamelCase     :				Any			,       UpperCamelCase     :				Any							)   ->    str:
		_lowerCamelCase				   =  dct.pop(UpperCamelCase							)
		_lowerCamelCase				   =  val


def 	lowerCamelCase					(	UpperCamelCase     :				Tuple			,       UpperCamelCase     :				Tuple							)   ->    str:
		_lowerCamelCase				   =  [int(backbone_config.embed_dim * 2**i							) for i in range(len(backbone_config.depths							)							)]
		for i in range(len(backbone_config.depths							)							):
				_lowerCamelCase				   =  num_features[i]
				for j in range(backbone_config.depths[i]							):
						# fmt: off
						# read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
						_lowerCamelCase				   =  state_dict.pop(F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight"""							)
						_lowerCamelCase				   =  state_dict.pop(F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.bias"""							)
						# next, add query, keys and values (in that order) to the state dict
						_lowerCamelCase				   =  in_proj_weight[:dim, :]
						_lowerCamelCase				   =  in_proj_bias[: dim]
						_lowerCamelCase				   =  in_proj_weight[
						    dim : dim * 2, :
						]
						_lowerCamelCase				   =  in_proj_bias[
						    dim : dim * 2
						]
						_lowerCamelCase				   =  in_proj_weight[
						    -dim :, :
						]
						_lowerCamelCase				   =  in_proj_bias[-dim :]
						# fmt: on


def 	lowerCamelCase					(	UpperCamelCase     :				Any							)   ->    Optional[int]:
		_lowerCamelCase   ,							_lowerCamelCase				   =  x.shape
		_lowerCamelCase				   =  x.reshape(UpperCamelCase			,       4			,       in_channel // 4							)
		_lowerCamelCase				   =  x[:, [0, 2, 1, 3], :].transpose(1			,       2							).reshape(UpperCamelCase			,       UpperCamelCase							)
		return x


def 	lowerCamelCase					(	UpperCamelCase     :				Optional[Any]							)   ->    int:
		_lowerCamelCase   ,							_lowerCamelCase				   =  x.shape
		_lowerCamelCase				   =  x.reshape(UpperCamelCase			,       in_channel // 4			,       4							)
		_lowerCamelCase				   =  x[:, :, [0, 2, 1, 3]].transpose(1			,       2							).reshape(UpperCamelCase			,       UpperCamelCase							)

		return x


def 	lowerCamelCase					(	UpperCamelCase     :				Tuple							)   ->    List[str]:
		_lowerCamelCase				   =  x.shape[0]
		_lowerCamelCase				   =  x.reshape(4			,       in_channel // 4							)
		_lowerCamelCase				   =  x[[0, 2, 1, 3], :].transpose(0			,       1							).reshape(UpperCamelCase							)
		return x


def 	lowerCamelCase					(	UpperCamelCase     :				Tuple							)   ->    Union[str, Any]:
		_lowerCamelCase				   =  x.shape[0]
		_lowerCamelCase				   =  x.reshape(in_channel // 4			,       4							)
		_lowerCamelCase				   =  x[:, [0, 2, 1, 3]].transpose(0			,       1							).reshape(UpperCamelCase							)
		return x


def 	lowerCamelCase					(	UpperCamelCase     :				List[Any]			,       UpperCamelCase     :				int			,       UpperCamelCase     :				Dict							)   ->    Union[str, Any]:
		_lowerCamelCase				   =  {
		    'upernet-swin-tiny': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210531_112542-e380ad3e.pth',
		    'upernet-swin-small': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210526_192015-ee2fff1c.pth',
		    'upernet-swin-base': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K_20210531_125459-429057bf.pth',
		    'upernet-swin-large': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k_20220318_091743-9ba68901.pth',
		}
		_lowerCamelCase				   =  model_name_to_url[model_name]
		_lowerCamelCase				   =  torch.hub.load_state_dict_from_url(UpperCamelCase			,       map_location='cpu'			,       file_name=UpperCamelCase							)[
		    'state_dict'
		]

		for name, param in state_dict.items():
				print(UpperCamelCase			,       param.shape							)

		_lowerCamelCase				   =  get_upernet_config(UpperCamelCase							)
		_lowerCamelCase				   =  UperNetForSemanticSegmentation(UpperCamelCase							)
		model.eval()

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

		# rename keys
		_lowerCamelCase				   =  create_rename_keys(UpperCamelCase							)
		for src, dest in rename_keys:
				rename_key(UpperCamelCase			,       UpperCamelCase			,       UpperCamelCase							)
		read_in_q_k_v(UpperCamelCase			,       config.backbone_config							)

		# fix downsample parameters
		for key, value in state_dict.items():
				if "downsample" in key:
						if "reduction" in key:
								_lowerCamelCase				   =  reverse_correct_unfold_reduction_order(UpperCamelCase							)
						if "norm" in key:
								_lowerCamelCase				   =  reverse_correct_unfold_norm_order(UpperCamelCase							)

		model.load_state_dict(UpperCamelCase							)

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

		_lowerCamelCase				   =  SegformerImageProcessor()
		_lowerCamelCase				   =  processor(UpperCamelCase			,       return_tensors='pt'							).pixel_values

		with torch.no_grad():
				_lowerCamelCase				   =  model(UpperCamelCase							)
				_lowerCamelCase				   =  outputs.logits

		print(logits.shape							)
		print('First values of logits:'			,       logits[0, 0, :3, :3]							)
		# assert values
		if model_name == "upernet-swin-tiny":
				_lowerCamelCase				   =  torch.tensor(
				    [[-7.5_958, -7.5_958, -7.4_302], [-7.5_958, -7.5_958, -7.4_302], [-7.4_797, -7.4_797, -7.3_068]]							)
		elif model_name == "upernet-swin-small":
				_lowerCamelCase				   =  torch.tensor(
				    [[-7.1_921, -7.1_921, -6.9_532], [-7.1_921, -7.1_921, -6.9_532], [-7.0_908, -7.0_908, -6.8_534]]							)
		elif model_name == "upernet-swin-base":
				_lowerCamelCase				   =  torch.tensor(
				    [[-6.5_851, -6.5_851, -6.4_330], [-6.5_851, -6.5_851, -6.4_330], [-6.4_763, -6.4_763, -6.3_254]]							)
		elif model_name == "upernet-swin-large":
				_lowerCamelCase				   =  torch.tensor(
				    [[-7.5_297, -7.5_297, -7.3_802], [-7.5_297, -7.5_297, -7.3_802], [-7.4_044, -7.4_044, -7.2_586]]							)
		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__":
				A						=			argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    '--model_name',
				    default='upernet-swin-tiny',
				    type=str,
				    choices=[F'''upernet-swin-{size}''' for size in ['tiny', 'small', 'base', 'large']],
				    help='Name of the Swin + 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						=			parser.parse_args()
				convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)