




'''simple docstring'''

import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    FalconForCausalLM,
						    FalconForQuestionAnswering,
						    FalconForSequenceClassification,
						    FalconForTokenClassification,
						    FalconModel,
						)



class 		SCREAMING_SNAKE_CASE__  :





			def __init__(      self						:    str      ,		a_						:    List[str]      ,		a_						:    Tuple=3      ,		a_						:    Any=7      ,		a_						:    Any=True      ,		a_						:    Union[str, Any]=True      ,		a_						:    Tuple=False      ,		a_						:    Optional[int]=True      ,		a_						:    Any=99      ,		a_						:    Dict=32      ,		a_						:    Dict=5      ,		a_						:    List[Any]=4      ,		a_						:    Any=37      ,		a_						:    Any="gelu"      ,		a_						:    List[str]=0.1      ,		a_						:    Dict=0.1      ,		a_						:    Optional[Any]=512      ,		a_						:    List[Any]=16      ,		a_						:    Any=2      ,		a_						:    str=0.02      ,		a_						:    Any=3      ,		a_						:    List[Any]=4      ,		a_						:    List[str]=None      ,		):



						"""simple docstring"""




						__snake_case			=		parent
						__snake_case			=		batch_size
						__snake_case			=		seq_length
						__snake_case			=		is_training
						__snake_case			=		use_input_mask
						__snake_case			=		use_token_type_ids
						__snake_case			=		use_labels
						__snake_case			=		vocab_size
						__snake_case			=		hidden_size
						__snake_case			=		num_hidden_layers
						__snake_case			=		num_attention_heads
						__snake_case			=		intermediate_size
						__snake_case			=		hidden_act
						__snake_case			=		hidden_dropout_prob
						__snake_case			=		attention_probs_dropout_prob
						__snake_case			=		max_position_embeddings
						__snake_case			=		type_vocab_size
						__snake_case			=		type_sequence_label_size
						__snake_case			=		initializer_range
						__snake_case			=		num_labels
						__snake_case			=		num_choices
						__snake_case			=		scope





			def 							A (      self						:    Any							):



						"""simple docstring"""




						__snake_case			=		ids_tensor([self.batch_size, self.seq_length]      ,		self.vocab_size							)

						__snake_case			=		None
						if self.use_input_mask:
									__snake_case			=		random_attention_mask([self.batch_size, self.seq_length]							)

						__snake_case			=		None

						__snake_case			=		None
						__snake_case			=		None
						__snake_case			=		None
						if self.use_labels:
									__snake_case			=		ids_tensor([self.batch_size]      ,		self.type_sequence_label_size							)
									__snake_case			=		ids_tensor([self.batch_size, self.seq_length]      ,		self.num_labels							)
									__snake_case			=		ids_tensor([self.batch_size]      ,		self.num_choices							)

						__snake_case			=		self.get_config()

						return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





			def 							A (      self						:    Optional[int]							):



						"""simple docstring"""




						return FalconConfig(
						    vocab_size=self.vocab_size      ,		hidden_size=self.hidden_size      ,		num_hidden_layers=self.num_hidden_layers      ,		num_attention_heads=self.num_attention_heads      ,		intermediate_size=self.intermediate_size      ,		hidden_act=self.hidden_act      ,		hidden_dropout_prob=self.hidden_dropout_prob      ,		attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,		max_position_embeddings=self.max_position_embeddings      ,		type_vocab_size=self.type_vocab_size      ,		is_decoder=a_      ,		initializer_range=self.initializer_range      ,		pad_token_id=1      ,		new_decoder_architecture=a_      ,		)





			def 							A (      self						:    List[str]      ,		a_						:    Dict      ,		a_						:    Tuple      ,		a_						:    Optional[Any]      ,		a_						:    Dict      ,		a_						:    Dict      ,		a_						:    Dict      ,		a_						:    Union[str, Any]							):



						"""simple docstring"""




						__snake_case			=		FalconModel(config=a_							)
						model.to(a_							)
						model.eval()
						__snake_case			=		model(a_      ,		attention_mask=a_							)
						__snake_case			=		model(a_							)
						self.parent.assertEqual(result.last_hidden_state.shape      ,		(self.batch_size, self.seq_length, self.hidden_size)							)





			def 							A (      self						:    List[Any]      ,		a_						:    List[Any]      ,		a_						:    Union[str, Any]      ,		a_						:    Optional[Any]      ,		a_						:    Any      ,		a_						:    List[Any]      ,		a_						:    Optional[Any]      ,		a_						:    Union[str, Any]      ,		a_						:    Tuple      ,		a_						:    Optional[int]      ,		):



						"""simple docstring"""




						__snake_case			=		True
						__snake_case			=		FalconModel(a_							)
						model.to(a_							)
						model.eval()
						__snake_case			=		model(
						    a_      ,		attention_mask=a_      ,		encoder_hidden_states=a_      ,		encoder_attention_mask=a_      ,		)
						__snake_case			=		model(
						    a_      ,		attention_mask=a_      ,		encoder_hidden_states=a_      ,		)
						__snake_case			=		model(a_      ,		attention_mask=a_							)
						self.parent.assertEqual(result.last_hidden_state.shape      ,		(self.batch_size, self.seq_length, self.hidden_size)							)





			def 							A (      self						:    Optional[int]      ,		a_						:    int      ,		a_						:    int      ,		a_						:    List[Any]      ,		a_						:    str      ,		a_						:    List[str]      ,		a_						:    str      ,		a_						:    str      ,		a_						:    Union[str, Any]      ,		a_						:    Optional[int]      ,		):



						"""simple docstring"""




						__snake_case			=		FalconForCausalLM(config=a_							)
						model.to(a_							)
						model.eval()
						__snake_case			=		model(a_      ,		attention_mask=a_      ,		labels=a_							)
						self.parent.assertEqual(result.logits.shape      ,		(self.batch_size, self.seq_length, self.vocab_size)							)





			def 							A (      self						:    List[Any]      ,		a_						:    Optional[int]      ,		a_						:    Optional[Any]      ,		a_						:    str      ,		a_						:    Tuple      ,		a_						:    str      ,		a_						:    List[Any]      ,		a_						:    Optional[Any]      ,		a_						:    Any      ,		a_						:    Dict      ,		):



						"""simple docstring"""




						__snake_case			=		True
						__snake_case			=		True
						__snake_case			=		FalconForCausalLM(config=a_							)
						model.to(a_							)
						model.eval()

						# first forward pass
						__snake_case			=		model(
						    a_      ,		attention_mask=a_      ,		encoder_hidden_states=a_      ,		encoder_attention_mask=a_      ,		use_cache=a_      ,		)
						__snake_case			=		outputs.past_key_values

						# create hypothetical multiple next token and extent to next_input_ids
						__snake_case			=		ids_tensor((self.batch_size, 3)      ,		config.vocab_size							)
						__snake_case			=		ids_tensor((self.batch_size, 3)      ,		vocab_size=2							)

						# append to next input_ids and
						__snake_case			=		torch.cat([input_ids, next_tokens]      ,		dim=-1							)
						__snake_case			=		torch.cat([input_mask, next_mask]      ,		dim=-1							)

						__snake_case			=		model(
						    a_      ,		attention_mask=a_      ,		encoder_hidden_states=a_      ,		encoder_attention_mask=a_      ,		output_hidden_states=a_      ,		)["hidden_states"][0]
						__snake_case			=		model(
						    a_      ,		attention_mask=a_      ,		encoder_hidden_states=a_      ,		encoder_attention_mask=a_      ,		past_key_values=a_      ,		output_hidden_states=a_      ,		)["hidden_states"][0]

						# select random slice
						__snake_case			=		ids_tensor((1,)      ,		output_from_past.shape[-1]							).item()
						__snake_case			=		output_from_no_past[:, -3:, random_slice_idx].detach()
						__snake_case			=		output_from_past[:, :, random_slice_idx].detach()

						self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]							)

						# test that outputs are equal for slice
						self.parent.assertTrue(torch.allclose(a_      ,		a_      ,		atol=1e-3							)							)






			def 							A (      self						:    Optional[Any]							):



						"""simple docstring"""




						__snake_case			=		self.prepare_config_and_inputs()
						(
						    (
						    __snake_case
						)		,     (
						    __snake_case
						)		,     (
						    __snake_case
						)		,     (
						    __snake_case
						)		,     (
						    __snake_case
						)		,     (
						    __snake_case
						)		,     (
						    __snake_case
						)		,     
						)			=		config_and_inputs
						__snake_case			=		{"input_ids": input_ids, "attention_mask": input_mask}
						return config, inputs_dict







@require_torch
class 		SCREAMING_SNAKE_CASE__  (			_UpperCamelCase				,     _UpperCamelCase				,     _UpperCamelCase				,     unittest.TestCase	):
			__SCREAMING_SNAKE_CASE    =			(
			    (
			        FalconModel,
			        FalconForCausalLM,
			        FalconForSequenceClassification,
			        FalconForTokenClassification,
			        FalconForQuestionAnswering,
			    )
			    if is_torch_available()
			    else ()
			)
			__SCREAMING_SNAKE_CASE    =			(FalconForCausalLM,) if is_torch_available() else ()
			__SCREAMING_SNAKE_CASE    =			(
			    {
			        """feature-extraction""": FalconModel,
			        """text-classification""": FalconForSequenceClassification,
			        """text-generation""": FalconForCausalLM,
			        """question-answering""": FalconForQuestionAnswering,
			        """token-classification""": FalconForTokenClassification,
			        """zero-shot""": FalconForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			__SCREAMING_SNAKE_CASE    =			False
			__SCREAMING_SNAKE_CASE    =			False





			def 							A (      self						:    Optional[Any]							):



						"""simple docstring"""




						__snake_case			=		FalconModelTester(self							)
						__snake_case			=		ConfigTester(self      ,		config_class=a_      ,		hidden_size=37							)





			def 							A (      self						:    Optional[Any]							):



						"""simple docstring"""




						self.config_tester.run_common_tests()





			def 							A (      self						:    List[Any]							):



						"""simple docstring"""




						__snake_case			=		self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*a_							)





			def 							A (      self						:    List[str]							):



						"""simple docstring"""




						__snake_case		,     *__snake_case			=		self.model_tester.prepare_config_and_inputs()
						for alibi in [True, False]:
									__snake_case			=		alibi
									self.model_tester.create_and_check_model(a_      ,		*a_							)





			def 							A (      self						:    Tuple							):



						"""simple docstring"""




						__snake_case		,     __snake_case			=		self.model_tester.prepare_config_and_inputs_for_common()
						__snake_case			=		3
						__snake_case			=		input_dict["input_ids"]
						__snake_case			=		input_ids.ne(1							).to(a_							)
						__snake_case			=		ids_tensor([self.model_tester.batch_size]      ,		self.model_tester.type_sequence_label_size							)
						__snake_case			=		FalconForSequenceClassification(a_							)
						model.to(a_							)
						model.eval()
						__snake_case			=		model(a_      ,		attention_mask=a_      ,		labels=a_							)
						self.assertEqual(result.logits.shape      ,		(self.model_tester.batch_size, self.model_tester.num_labels)							)





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



						"""simple docstring"""




						__snake_case		,     __snake_case			=		self.model_tester.prepare_config_and_inputs_for_common()
						__snake_case			=		3
						__snake_case			=		"single_label_classification"
						__snake_case			=		input_dict["input_ids"]
						__snake_case			=		input_ids.ne(1							).to(a_							)
						__snake_case			=		ids_tensor([self.model_tester.batch_size]      ,		self.model_tester.type_sequence_label_size							)
						__snake_case			=		FalconForSequenceClassification(a_							)
						model.to(a_							)
						model.eval()
						__snake_case			=		model(a_      ,		attention_mask=a_      ,		labels=a_							)
						self.assertEqual(result.logits.shape      ,		(self.model_tester.batch_size, self.model_tester.num_labels)							)





			def 							A (      self						:    Optional[Any]							):



						"""simple docstring"""




						__snake_case		,     __snake_case			=		self.model_tester.prepare_config_and_inputs_for_common()
						__snake_case			=		input_dict["input_ids"]
						__snake_case			=		FalconForCausalLM(a_							)
						model.to(a_							)
						model.eval()
						__snake_case			=		model(a_      ,		use_cache=a_							)
						__snake_case			=		input_ids.shape[0]
						__snake_case			=		model._convert_to_rw_cache(result.past_key_values							)
						__snake_case			=		model._convert_cache_to_standard_format(a_      ,		a_							)
						for layer in range(len(a_							)							):
									for tensor_idx in range(2							):
												self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3							)
												self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4							)
												self.assertTrue(
												    torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]							)							)





			def 							A (      self						:    Optional[Any]							):



						"""simple docstring"""




						__snake_case		,     __snake_case			=		self.model_tester.prepare_config_and_inputs_for_common()
						__snake_case			=		3
						__snake_case			=		"multi_label_classification"
						__snake_case			=		input_dict["input_ids"]
						__snake_case			=		input_ids.ne(1							).to(a_							)
						__snake_case			=		ids_tensor(
						    [self.model_tester.batch_size, config.num_labels]      ,		self.model_tester.type_sequence_label_size							).to(torch.float							)
						__snake_case			=		FalconForSequenceClassification(a_							)
						model.to(a_							)
						model.eval()
						__snake_case			=		model(a_      ,		attention_mask=a_      ,		labels=a_							)
						self.assertEqual(result.logits.shape      ,		(self.model_tester.batch_size, self.model_tester.num_labels)							)






			def 							A (      self						:    Dict							):



						"""simple docstring"""




						for model_class in self.all_generative_model_classes:
									__snake_case		,     __snake_case			=		self.model_tester.prepare_config_and_inputs_for_common()

									# If it doesn't support cache, pass the test
									if not hasattr(a_      ,		"use_cache"							):
												return

									__snake_case			=		model_class(a_							).to(a_							)
									if "use_cache" not in inputs:
												__snake_case			=		True
									__snake_case			=		model(**a_							)

									# If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
									if "past_key_values" not in outputs:
												return

									__snake_case			=		(
									    getattr(a_      ,		"decoder_layers"      ,		a_							)
									    or getattr(a_      ,		"num_decoder_layers"      ,		a_							)
									    or config.num_hidden_layers
									)
									__snake_case			=		getattr(a_      ,		"num_kv_heads"      ,		config.num_attention_heads							)
									__snake_case			=		getattr(a_      ,		"d_model"      ,		config.hidden_size							)
									__snake_case			=		embed_dim // num_attention_heads

									__snake_case			=		outputs["past_key_values"]
									self.assertEqual(len(a_							)      ,		a_							)

									__snake_case		,     __snake_case			=		inputs["input_ids"].shape
									for i in range(a_							):
												if config.new_decoder_architecture:
															__snake_case			=		config.num_attention_heads
												elif config.multi_query:
															__snake_case			=		1
												self.assertEqual(len(past_kv[0]							)      ,		2							)  # K V for the decoder = 2
												self.assertEqual(
												    past_kv[i][0].shape      ,		(batch_size, num_attention_heads, seq_length, per_head_embed_dim)							)
												self.assertEqual(
												    past_kv[i][1].shape      ,		(batch_size, num_attention_heads, seq_length, per_head_embed_dim)							)







@require_torch
class 		SCREAMING_SNAKE_CASE__  (			unittest.TestCase	):





			@slow
			def 							A (      self						:    Any							):



						"""simple docstring"""




						__snake_case			=		AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b"							)
						__snake_case			=		FalconForCausalLM.from_pretrained("Rocketknight1/falcon-rw-1b"							)
						model.eval()
						model.to(a_							)
						__snake_case			=		tokenizer("My favorite food is"      ,		return_tensors="pt"							).to(a_							)

						__snake_case			=		(
						    "My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday."
						)

						__snake_case			=		model.generate(**a_      ,		do_sample=a_      ,		max_new_tokens=19							)
						__snake_case			=		tokenizer.batch_decode(a_							)[0]

						self.assertEqual(a_      ,		a_							)





			@slow
			def 							A (      self						:    Optional[int]							):



						"""simple docstring"""




						for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
									__snake_case			=		AutoTokenizer.from_pretrained(a_							)
									__snake_case			=		FalconForCausalLM.from_pretrained(a_							)
									model.eval()
									model.to(a_							)
									__snake_case			=		tokenizer("My favorite food is"      ,		return_tensors="pt"							).to(a_							)

									# We just test that these run without errors - the models are randomly initialized
									# and so the actual text outputs will be garbage
									model.generate(**a_      ,		do_sample=a_      ,		max_new_tokens=4							)
									model.generate(**a_      ,		do_sample=a_      ,		max_new_tokens=4							)
									model.generate(**a_      ,		num_beams=2      ,		max_new_tokens=4							)






			@slow
			def 							A (      self						:    Any							):



						"""simple docstring"""




						with torch.no_grad():
									for repo in [
									    "Rocketknight1/falcon-rw-1b",
									    "Rocketknight1/tiny-random-falcon-7b",
									    "Rocketknight1/tiny-random-falcon-40b",
									]:
												__snake_case			=		AutoTokenizer.from_pretrained(a_							)
												__snake_case			=		FalconForCausalLM.from_pretrained(a_							)
												model.eval()
												model.to(device=a_							)
												__snake_case			=		tokenizer("My favorite food is"      ,		return_tensors="pt"							).to(a_							)

												# Test results are the same with and without cache
												__snake_case			=		model.generate(**a_      ,		do_sample=a_      ,		max_new_tokens=20      ,		use_cache=a_							)
												__snake_case			=		model.generate(**a_      ,		do_sample=a_      ,		max_new_tokens=20      ,		use_cache=a_							)
												self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0							)








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 a(lowercase__						):





	'''simple docstring'''
	snake_case_   		=			384
	snake_case_   		=			7
	if "tiny" in model_name:
		snake_case_   		=			96
		snake_case_   		=			(2, 2, 6, 2)
		snake_case_   		=			(3, 6, 12, 24)
	elif "small" in model_name:
		snake_case_   		=			96
		snake_case_   		=			(2, 2, 18, 2)
		snake_case_   		=			(3, 6, 12, 24)
	elif "base" in model_name:
		snake_case_   		=			128
		snake_case_   		=			(2, 2, 18, 2)
		snake_case_   		=			(4, 8, 16, 32)
		snake_case_   		=			12
		snake_case_   		=			512
	elif "large" in model_name:
		snake_case_   		=			192
		snake_case_   		=			(2, 2, 18, 2)
		snake_case_   		=			(6, 12, 24, 48)
		snake_case_   		=			12
		snake_case_   		=			768

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

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

	return config





def a(lowercase__						):





	'''simple docstring'''
	snake_case_   		=			[]

	# 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 a(lowercase__      ,			lowercase__      ,			lowercase__						):





	'''simple docstring'''
	snake_case_   		=			dct.pop(lowercase__						)
	snake_case_   		=			val





def a(lowercase__      ,			lowercase__						):





	'''simple docstring'''
	snake_case_   		=			[int(backbone_config.embed_dim * 2**i						) for i in range(len(backbone_config.depths						)						)]
	for i in range(len(backbone_config.depths						)						):
		snake_case_   		=			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)
			snake_case_   		=			state_dict.pop(f"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight"""						)
			snake_case_   		=			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
			snake_case_   		=			in_proj_weight[:dim, :]
			snake_case_   		=			in_proj_bias[: dim]
			snake_case_   		=			in_proj_weight[
			    dim : dim * 2, :
			]
			snake_case_   		=			in_proj_bias[
			    dim : dim * 2
			]
			snake_case_   		=			in_proj_weight[
			    -dim :, :
			]
			snake_case_   		=			in_proj_bias[-dim :]
			# fmt: on





def a(lowercase__						):





	'''simple docstring'''
	snake_case_		,							snake_case_   		=			x.shape
	snake_case_   		=			x.reshape(lowercase__      ,			4      ,			in_channel // 4						)
	snake_case_   		=			x[:, [0, 2, 1, 3], :].transpose(1      ,			2						).reshape(lowercase__      ,			lowercase__						)
	return x





def a(lowercase__						):





	'''simple docstring'''
	snake_case_		,							snake_case_   		=			x.shape
	snake_case_   		=			x.reshape(lowercase__      ,			in_channel // 4      ,			4						)
	snake_case_   		=			x[:, :, [0, 2, 1, 3]].transpose(1      ,			2						).reshape(lowercase__      ,			lowercase__						)

	return x





def a(lowercase__						):





	'''simple docstring'''
	snake_case_   		=			x.shape[0]
	snake_case_   		=			x.reshape(4      ,			in_channel // 4						)
	snake_case_   		=			x[[0, 2, 1, 3], :].transpose(0      ,			1						).reshape(lowercase__						)
	return x





def a(lowercase__						):





	'''simple docstring'''
	snake_case_   		=			x.shape[0]
	snake_case_   		=			x.reshape(in_channel // 4      ,			4						)
	snake_case_   		=			x[:, [0, 2, 1, 3]].transpose(0      ,			1						).reshape(lowercase__						)
	return x





def a(lowercase__      ,			lowercase__      ,			lowercase__						):





	'''simple docstring'''
	snake_case_   		=			{
	    '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',
	}
	snake_case_   		=			model_name_to_url[model_name]
	snake_case_   		=			torch.hub.load_state_dict_from_url(lowercase__      ,			map_location='cpu'      ,			file_name=lowercase__						)[
	    'state_dict'
	]

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

	snake_case_   		=			get_upernet_config(lowercase__						)
	snake_case_   		=			UperNetForSemanticSegmentation(lowercase__						)
	model.eval()

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

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

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

	model.load_state_dict(lowercase__						)

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

	snake_case_   		=			SegformerImageProcessor()
	snake_case_   		=			processor(lowercase__      ,			return_tensors='pt'						).pixel_values

	with torch.no_grad():
		snake_case_   		=			model(lowercase__						)
		snake_case_   		=			outputs.logits

	print(logits.shape						)
	print('First values of logits:'      ,			logits[0, 0, :3, :3]						)
	# assert values
	if model_name == "upernet-swin-tiny":
		snake_case_   		=			torch.tensor(
		    [[-7.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]]						)
	elif model_name == "upernet-swin-small":
		snake_case_   		=			torch.tensor(
		    [[-7.1921, -7.1921, -6.9532], [-7.1921, -7.1921, -6.9532], [-7.0908, -7.0908, -6.8534]]						)
	elif model_name == "upernet-swin-base":
		snake_case_   		=			torch.tensor(
		    [[-6.5851, -6.5851, -6.4330], [-6.5851, -6.5851, -6.4330], [-6.4763, -6.4763, -6.3254]]						)
	elif model_name == "upernet-swin-large":
		snake_case_   		=			torch.tensor(
		    [[-7.5297, -7.5297, -7.3802], [-7.5297, -7.5297, -7.3802], [-7.4044, -7.4044, -7.2586]]						)
	print('Logits:'      ,			outputs.logits[0, 0, :3, :3]						)
	assert torch.allclose(outputs.logits[0, 0, :3, :3]      ,			lowercase__      ,			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(lowercase__						)
		print(f"""Saving processor to {pytorch_dump_folder_path}"""						)
		processor.save_pretrained(lowercase__						)

	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)
