





import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import UperNetForSemanticSegmentation
		from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor

class       __A     :


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_=13   , UpperCAmelCase_=32   , UpperCAmelCase_=3   , UpperCAmelCase_=4   , UpperCAmelCase_=[10, 20, 30, 40]   , UpperCAmelCase_=[2, 2, 3, 2]   , UpperCAmelCase_=True   , UpperCAmelCase_=True   , UpperCAmelCase_=37   , UpperCAmelCase_="gelu"   , UpperCAmelCase_=10   , UpperCAmelCase_=0.0_2   , UpperCAmelCase_=["stage2", "stage3", "stage4"]   , UpperCAmelCase_=3   , UpperCAmelCase_=None   , ):
								lowerCamelCase	 =parent
								lowerCamelCase	 =batch_size
								lowerCamelCase	 =image_size
								lowerCamelCase	 =num_channels
								lowerCamelCase	 =num_stages
								lowerCamelCase	 =hidden_sizes
								lowerCamelCase	 =depths
								lowerCamelCase	 =is_training
								lowerCamelCase	 =use_labels
								lowerCamelCase	 =intermediate_size
								lowerCamelCase	 =hidden_act
								lowerCamelCase	 =type_sequence_label_size
								lowerCamelCase	 =initializer_range
								lowerCamelCase	 =out_features
								lowerCamelCase	 =num_labels
								lowerCamelCase	 =scope
								lowerCamelCase	 =num_stages


				def _snake_case       (		self							):
								lowerCamelCase	 =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

								lowerCamelCase	 =None
								if self.use_labels:
												lowerCamelCase	 =ids_tensor([self.batch_size]   , self.type_sequence_label_size							)

								lowerCamelCase	 =self.get_config()

								return config, pixel_values, labels


				def _snake_case       (		self							):
								return ConvNextConfig(
								    num_channels=self.num_channels   , num_stages=self.num_stages   , hidden_sizes=self.hidden_sizes   , depths=self.depths   , is_training=self.is_training   , intermediate_size=self.intermediate_size   , hidden_act=self.hidden_act   , out_features=self.out_features   , )


				def _snake_case       (		self							):
								return UperNetConfig(
								    backbone_config=self.get_backbone_config()   , hidden_size=512   , pool_scales=[1, 2, 3, 6]   , use_auxiliary_head=UpperCamelCase_   , auxiliary_loss_weight=0.4   , auxiliary_in_channels=40   , auxiliary_channels=256   , auxiliary_num_convs=1   , auxiliary_concat_input=UpperCamelCase_   , loss_ignore_index=255   , num_labels=self.num_labels   , )


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							):
								lowerCamelCase	 =UperNetForSemanticSegmentation(config=UpperCamelCase_							)
								model.to(UpperCamelCase_							)
								model.eval()
								lowerCamelCase	 =model(UpperCamelCase_							)
								self.parent.assertEqual(
								    result.logits.shape   , (self.batch_size, self.num_labels, self.image_size, self.image_size)							)


				def _snake_case       (		self							):
								lowerCamelCase	 =self.prepare_config_and_inputs()
								(
								    lowerCamelCase    
								)	 =config_and_inputs
								lowerCamelCase	 ={'pixel_values': pixel_values}
								return config, inputs_dict




@require_torch
class       __A     (						a    ,					a    ,					unittest.TestCase    ):
				__A      =				(UperNetForSemanticSegmentation,) if is_torch_available() else ()
				__A      =				{"""image-segmentation""": UperNetForSemanticSegmentation} if is_torch_available() else {}
				__A      =				False
				__A      =				False
				__A      =				False
				__A      =				False
				__A      =				False
				__A      =				False


				def _snake_case       (		self							):
								lowerCamelCase	 =UperNetModelTester(self							)
								lowerCamelCase	 =ConfigTester(self   , config_class=UpperCamelCase_   , has_text_modality=UpperCamelCase_   , hidden_size=37							)


				def _snake_case       (		self							):
								self.create_and_test_config_common_properties()
								self.config_tester.create_and_test_config_to_json_string()
								self.config_tester.create_and_test_config_to_json_file()
								self.config_tester.create_and_test_config_from_and_save_pretrained()
								self.config_tester.create_and_test_config_with_num_labels()
								self.config_tester.check_config_can_be_init_without_params()
								self.config_tester.check_config_arguments_init()


				def _snake_case       (		self							):
								return


				def _snake_case       (		self							):
								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												lowerCamelCase	 =model_class(UpperCamelCase_							)
												lowerCamelCase	 =inspect.signature(model.forward							)
												# signature.parameters is an OrderedDict => so arg_names order is deterministic
												lowerCamelCase	 =[*signature.parameters.keys()]

												lowerCamelCase	 =['pixel_values']
												self.assertListEqual(arg_names[:1]   , UpperCamelCase_							)


				def _snake_case       (		self							):
								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_semantic_segmentation(*UpperCamelCase_							)


				@unittest.skip(reason="""UperNet does not use inputs_embeds"""							)
				def _snake_case       (		self							):
								pass


				@unittest.skip(reason="""UperNet does not support input and output embeddings"""							)
				def _snake_case       (		self							):
								pass


				@unittest.skip(reason="""UperNet does not have a base model"""							)
				def _snake_case       (		self							):
								pass


				@unittest.skip(reason="""UperNet does not have a base model"""							)
				def _snake_case       (		self							):
								pass


				@require_torch_multi_gpu
				@unittest.skip(reason="""UperNet has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`"""							)
				def _snake_case       (		self							):
								pass


				@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""							)
				def _snake_case       (		self							):
								pass


				def _snake_case       (		self							):





								def check_hidden_states_output(UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							):
												lowerCamelCase	 =model_class(UpperCamelCase_							)
												model.to(UpperCamelCase_							)
												model.eval()

												with torch.no_grad():
																lowerCamelCase	 =model(**self._prepare_for_class(UpperCamelCase_   , UpperCamelCase_							)							)

												lowerCamelCase	 =outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

												lowerCamelCase	 =self.model_tester.num_stages
												self.assertEqual(len(UpperCamelCase_							)   , expected_num_stages + 1							)

												# ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
												self.assertListEqual(
												    list(hidden_states[0].shape[-2:]							)   , [self.model_tester.image_size // 4, self.model_tester.image_size // 4]   , )

								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												lowerCamelCase	 =True
												check_hidden_states_output(UpperCamelCase_   , UpperCamelCase_   , UpperCamelCase_							)

												# check that output_hidden_states also work using config
												del inputs_dict["output_hidden_states"]
												lowerCamelCase	 =True

												check_hidden_states_output(UpperCamelCase_   , UpperCamelCase_   , UpperCamelCase_							)


				def _snake_case       (		self							):
								lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()

								lowerCamelCase	 =_config_zero_init(UpperCamelCase_							)
								lowerCamelCase	 =_config_zero_init(configs_no_init.backbone_config							)
								for model_class in self.all_model_classes:
												lowerCamelCase	 =model_class(config=UpperCamelCase_							)
												for name, param in model.named_parameters():
																if param.requires_grad:
																				self.assertIn(
																				    ((param.data.mean() * 1E9).round() / 1E9).item()   , [0.0, 1.0]   , msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""   , )


				@unittest.skip(reason="""UperNet does not have tied weights"""							)
				def _snake_case       (		self							):
								pass


				@slow
				def _snake_case       (		self							):





								for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												lowerCamelCase	 =UperNetForSemanticSegmentation.from_pretrained(UpperCamelCase_							)
												self.assertIsNotNone(UpperCamelCase_							)





def      _lowercase					(						)						->  int:
				lowerCamelCase	 =hf_hub_download(
				    repo_id="""hf-internal-testing/fixtures_ade20k"""			,      repo_type="""dataset"""			,      filename="""ADE_val_00000001.jpg"""				)
				lowerCamelCase	 =Image.open(_lowercase				).convert("""RGB"""				)
				return image




@require_torch
@require_vision
@slow
class       __A     (						unittest.TestCase    ):


				def _snake_case       (		self							):
								lowerCamelCase	 =AutoImageProcessor.from_pretrained("""openmmlab/upernet-swin-tiny"""							)
								lowerCamelCase	 =UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-swin-tiny"""							).to(UpperCamelCase_							)

								lowerCamelCase	 =prepare_img()
								lowerCamelCase	 =processor(images=UpperCamelCase_   , return_tensors="""pt"""							).to(UpperCamelCase_							)

								with torch.no_grad():
												lowerCamelCase	 =model(**UpperCamelCase_							)

								lowerCamelCase	 =torch.Size((1, model.config.num_labels, 512, 512)							)
								self.assertEqual(outputs.logits.shape   , UpperCamelCase_							)

								lowerCamelCase	 =torch.tensor(
								    [[-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.4_7_9_7, -7.4_7_9_7, -7.3_0_6_8]]							).to(UpperCamelCase_							)
								self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]   , UpperCamelCase_   , atol=1E-4							)							)


				def _snake_case       (		self							):
								lowerCamelCase	 =AutoImageProcessor.from_pretrained("""openmmlab/upernet-convnext-tiny"""							)
								lowerCamelCase	 =UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-convnext-tiny"""							).to(UpperCamelCase_							)

								lowerCamelCase	 =prepare_img()
								lowerCamelCase	 =processor(images=UpperCamelCase_   , return_tensors="""pt"""							).to(UpperCamelCase_							)

								with torch.no_grad():
												lowerCamelCase	 =model(**UpperCamelCase_							)

								lowerCamelCase	 =torch.Size((1, model.config.num_labels, 512, 512)							)
								self.assertEqual(outputs.logits.shape   , UpperCamelCase_							)

								lowerCamelCase	 =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]]							).to(UpperCamelCase_							)
								self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]   , UpperCamelCase_   , atol=1E-4							)							)












import argparse
import gc
import json
import os
import re

import torch
from huggingface_hub import hf_hub_download

from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast, RwkvConfig
from transformers.modeling_utils import WEIGHTS_INDEX_NAME, shard_checkpoint


UpperCAmelCase__						:			Optional[int]									={
    '''169M''': 12,
    '''430M''': 24,
    '''1B5''': 24,
    '''3B''': 32,
    '''7B''': 32,
    '''14B''': 40,
}

UpperCAmelCase__						:			str									={
    '''169M''': 7_68,
    '''430M''': 10_24,
    '''1B5''': 20_48,
    '''3B''': 25_60,
    '''7B''': 40_96,
    '''14B''': 51_20,
}





def      _lowercase					(						_UpperCAmelCase				)						->  Tuple:
				lowerCamelCase	 =list(state_dict.keys()				)
				for name in state_dict_keys:
								lowerCamelCase	 =state_dict.pop(_UpperCAmelCase				)
								# emb -> embedding
								if name.startswith("""emb."""				):
												lowerCamelCase	 =name.replace("""emb."""			,      """embeddings."""				)
								# ln_0 -> pre_ln (only present at block 0)
								if name.startswith("""blocks.0.ln0"""				):
												lowerCamelCase	 =name.replace("""blocks.0.ln0"""			,      """blocks.0.pre_ln"""				)
								# att -> attention
								lowerCamelCase	 =re.sub(r"""blocks\.(\d+)\.att"""			,      r"""blocks.\1.attention"""			,      _UpperCAmelCase				)
								# ffn -> feed_forward
								lowerCamelCase	 =re.sub(r"""blocks\.(\d+)\.ffn"""			,      r"""blocks.\1.feed_forward"""			,      _UpperCAmelCase				)
								# time_mix_k -> time_mix_key and reshape
								if name.endswith(""".time_mix_k"""				):
												lowerCamelCase	 =name.replace(""".time_mix_k"""			,      """.time_mix_key"""				)
								# time_mix_v -> time_mix_value and reshape
								if name.endswith(""".time_mix_v"""				):
												lowerCamelCase	 =name.replace(""".time_mix_v"""			,      """.time_mix_value"""				)
								# time_mix_r -> time_mix_key and reshape
								if name.endswith(""".time_mix_r"""				):
												lowerCamelCase	 =name.replace(""".time_mix_r"""			,      """.time_mix_receptance"""				)

								if name != "head.weight":
												lowerCamelCase	 ="""rwkv.""" + name

								lowerCamelCase	 =weight
				return state_dict





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase=None			,      _UpperCAmelCase=None			,      _UpperCAmelCase=False			,      _UpperCAmelCase=None				)						->  Tuple:
				# 1. If possible, build the tokenizer.
				if tokenizer_file is None:
								print("""No `--tokenizer_file` provided, we will use the default tokenizer."""				)
								lowerCamelCase	 =5_02_77
								lowerCamelCase	 =AutoTokenizer.from_pretrained("""EleutherAI/gpt-neox-20b"""				)
				else:
								lowerCamelCase	 =PreTrainedTokenizerFast(tokenizer_file=_UpperCAmelCase				)
								lowerCamelCase	 =len(_UpperCAmelCase				)
				tokenizer.save_pretrained(_UpperCAmelCase				)

				# 2. Build the config
				lowerCamelCase	 =list(NUM_HIDDEN_LAYERS_MAPPING.keys()				)
				if size is None:
								# Try to infer size from the checkpoint name
								for candidate in possible_sizes:
												if candidate in checkpoint_file:
																lowerCamelCase	 =candidate
																break
								if size is None:
												raise ValueError("""Could not infer the size, please provide it with the `--size` argument."""				)
				if size not in possible_sizes:
								raise ValueError(F"""`size` should be one of {possible_sizes}, got {size}."""				)

				lowerCamelCase	 =RwkvConfig(
				    vocab_size=_UpperCAmelCase			,      num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size]			,      hidden_size=HIDEN_SIZE_MAPPING[size]			,      )
				config.save_pretrained(_UpperCAmelCase				)

				# 3. Download model file then convert state_dict
				lowerCamelCase	 =hf_hub_download(_UpperCAmelCase			,      _UpperCAmelCase				)
				lowerCamelCase	 =torch.load(_UpperCAmelCase			,      map_location="""cpu"""				)
				lowerCamelCase	 =convert_state_dict(_UpperCAmelCase				)

				# 4. Split in shards and save
				lowerCamelCase ,  lowerCamelCase	 =shard_checkpoint(_UpperCAmelCase				)
				for shard_file, shard in shards.items():
								torch.save(_UpperCAmelCase			,      os.path.join(_UpperCAmelCase			,      _UpperCAmelCase				)				)

				if index is not None:
								lowerCamelCase	 =os.path.join(_UpperCAmelCase			,      _UpperCAmelCase				)
								# Save the index as well
								with open(_UpperCAmelCase			,      """w"""			,      encoding="""utf-8"""				) as f:
												lowerCamelCase	 =json.dumps(_UpperCAmelCase			,      indent=2			,      sort_keys=_UpperCAmelCase				) + """\n"""
												f.write(_UpperCAmelCase				)

								# 5. Clean up shards (for some reason the file PyTorch saves take the same space as the whole state_dict
								print(
								    """Cleaning up shards. This may error with an OOM error, it this is the case don't worry you still have converted the model."""				)
								lowerCamelCase	 =list(shards.keys()				)

								del state_dict
								del shards
								gc.collect()

								for shard_file in shard_files:
												lowerCamelCase	 =torch.load(os.path.join(_UpperCAmelCase			,      _UpperCAmelCase				)				)
												torch.save({k: v.cpu().clone() for k, v in state_dict.items()}			,      os.path.join(_UpperCAmelCase			,      _UpperCAmelCase				)				)

				del state_dict
				gc.collect()

				if push_to_hub:
								if model_name is None:
												raise ValueError("""Please provide a `model_name` to push the model to the Hub."""				)
								lowerCamelCase	 =AutoModelForCausalLM.from_pretrained(_UpperCAmelCase				)
								model.push_to_hub(_UpperCAmelCase			,      max_shard_size="""2GB"""				)
								tokenizer.push_to_hub(_UpperCAmelCase				)


if __name__ == "__main__":
							UpperCAmelCase__						:			Dict									=argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '''--repo_id''', default=None, type=str, required=True, help='''Repo ID from which to pull the checkpoint.'''
							)
							parser.add_argument(
							    '''--checkpoint_file''', default=None, type=str, required=True, help='''Name of the checkpoint file in the repo.'''
							)
							parser.add_argument(
							    '''--output_dir''', default=None, type=str, required=True, help='''Where to save the converted model.'''
							)
							parser.add_argument(
							    '''--tokenizer_file''',
							    default=None,
							    type=str,
							    help='''Path to the tokenizer file to use (if not provided, only the model is converted).''',
							)
							parser.add_argument(
							    '''--size''',
							    default=None,
							    type=str,
							    help='''Size of the model. Will be inferred from the `checkpoint_file` if not passed.''',
							)
							parser.add_argument(
							    '''--push_to_hub''',
							    action='''store_true''',
							    help='''Push to the Hub the converted model.''',
							)
							parser.add_argument(
							    '''--model_name''',
							    default=None,
							    type=str,
							    help='''Name of the pushed model on the Hub, including the username / organization.''',
							)

							UpperCAmelCase__						:			List[Any]									=parser.parse_args()
							convert_rmkv_checkpoint_to_hf_format(
							    args.repo_id,
							    args.checkpoint_file,
							    args.output_dir,
							    size=args.size,
							    tokenizer_file=args.tokenizer_file,
							    push_to_hub=args.push_to_hub,
							    model_name=args.model_name,
							)





