



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       (					__lowerCAmelCase      )   ->					Tuple:


				"""simple docstring"""

				snake_case__	:      Union[str, Any]     		=	384
				if "tiny" in model_name:
								snake_case__	:      Any     		=	[3, 3, 9, 3]
								snake_case__	:      str     		=	[96, 192, 384, 768]
				if "small" in model_name:
								snake_case__	:      str     		=	[3, 3, 27, 3]
								snake_case__	:      Dict     		=	[96, 192, 384, 768]
				if "base" in model_name:
								snake_case__	:      Union[str, Any]     		=	[3, 3, 27, 3]
								snake_case__	:      Optional[Any]     		=	[128, 256, 512, 1024]
								snake_case__	:      str     		=	512
				if "large" in model_name:
								snake_case__	:      Any     		=	[3, 3, 27, 3]
								snake_case__	:      Dict     		=	[192, 384, 768, 1536]
								snake_case__	:      Dict     		=	768
				if "xlarge" in model_name:
								snake_case__	:      Any     		=	[3, 3, 27, 3]
								snake_case__	:      Tuple     		=	[256, 512, 1024, 2048]
								snake_case__	:      int     		=	1024

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

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

				return config



def 			_lowerCAmelCase       (					__lowerCAmelCase      )   ->					str:


				"""simple docstring"""

				snake_case__	:      List[Any]     		=	[]

				# 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       (					__lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase      )   ->					Any:


				"""simple docstring"""

				snake_case__	:      str     		=	dct.pop(__lowerCAmelCase      )
				snake_case__	:      Tuple     		=	val



def 			_lowerCAmelCase       (					__lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase      )   ->					Tuple:


				"""simple docstring"""

				snake_case__	:      Optional[int]     		=	{
				    '''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''',
				}
				snake_case__	:      Dict     		=	model_name_to_url[model_name]
				snake_case__	:      Tuple     		=	torch.hub.load_state_dict_from_url(__lowerCAmelCase  ,  map_location='''cpu'''      )['''state_dict''']

				snake_case__	:      Union[str, Any]     		=	get_upernet_config(__lowerCAmelCase      )
				snake_case__	:      List[str]     		=	UperNetForSemanticSegmentation(__lowerCAmelCase      )
				model.eval()

				# replace "bn" => "batch_norm"
				for key in state_dict.copy().keys():
								snake_case__	:      str     		=	state_dict.pop(__lowerCAmelCase      )
								if "bn" in key:
												snake_case__	:      Any     		=	key.replace('''bn'''  ,  '''batch_norm'''      )
								snake_case__	:      List[str]     		=	val

				# rename keys
				snake_case__	:      List[str]     		=	create_rename_keys(__lowerCAmelCase      )
				for src, dest in rename_keys:
								rename_key(__lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase      )

				model.load_state_dict(__lowerCAmelCase      )

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

				snake_case__	:      Any     		=	SegformerImageProcessor()
				snake_case__	:      List[Any]     		=	processor(__lowerCAmelCase  ,  return_tensors='''pt'''      ).pixel_values

				with torch.no_grad():
								snake_case__	:      Union[str, Any]     		=	model(__lowerCAmelCase      )

				if model_name == "upernet-convnext-tiny":
								snake_case__	:      Union[str, Any]     		=	torch.tensor(
								    [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]      )
				elif model_name == "upernet-convnext-small":
								snake_case__	:      Tuple     		=	torch.tensor(
								    [[-8.8_236, -8.8_236, -8.6_771], [-8.8_236, -8.8_236, -8.6_771], [-8.7_638, -8.7_638, -8.6_240]]      )
				elif model_name == "upernet-convnext-base":
								snake_case__	:      Dict     		=	torch.tensor(
								    [[-8.8_558, -8.8_558, -8.6_905], [-8.8_558, -8.8_558, -8.6_905], [-8.7_669, -8.7_669, -8.6_021]]      )
				elif model_name == "upernet-convnext-large":
								snake_case__	:      str     		=	torch.tensor(
								    [[-8.6_660, -8.6_660, -8.6_210], [-8.6_660, -8.6_660, -8.6_210], [-8.6_310, -8.6_310, -8.5_964]]      )
				elif model_name == "upernet-convnext-xlarge":
								snake_case__	:      Optional[Any]     		=	torch.tensor(
								    [[-8.4_980, -8.4_980, -8.3_977], [-8.4_980, -8.4_980, -8.3_977], [-8.4_379, -8.4_379, -8.3_412]]      )
				print('''Logits:'''  ,  outputs.logits[0, 0, :3, :3]      )
				assert torch.allclose(outputs.logits[0, 0, :3, :3]  ,  __lowerCAmelCase  ,  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(__lowerCAmelCase      )
								print(f"""Saving processor to {pytorch_dump_folder_path}"""      )
								processor.save_pretrained(__lowerCAmelCase      )

				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-convnext-tiny''',
				    type=str,
				    choices=[f"""upernet-convnext-{size}""" for size in ['''tiny''', '''small''', '''base''', '''large''', '''xlarge''']],
				    help='''Name of the ConvNext UperNet model you\'d like to convert.''',
				)
				parser.add_argument(
				    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
				)
				parser.add_argument(
				    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
				)

				A__								=						parser.parse_args()
				convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






import json
import multiprocessing
import os
import re
from collections import defaultdict

import torch
from accelerate import Accelerator
from accelerate.utils import set_seed
from arguments import HumanEvalArguments
from datasets import load_dataset, load_metric
from torch.utils.data import IterableDataset
from torch.utils.data.dataloader import DataLoader
from tqdm import tqdm

import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, StoppingCriteria, StoppingCriteriaList


A__								=						['''\nclass''', '''\ndef''', '''\n#''', '''\n@''', '''\nprint''', '''\nif''']


class       a							(			__lowerCamelCase				):



		def __init__(     self     :str   ,__lowercase     :Dict   ,__lowercase     :str   ,__lowercase     :Any=None   ,__lowercase     :Union[str, Any]=1  ):
						snake_case__	:      Optional[int]     		=	tokenizer
						snake_case__	:      List[Any]     		=	dataset
						snake_case__	:      Any     		=	len(__lowercase  ) if n_tasks is None else n_tasks
						snake_case__	:      List[Any]     		=	n_copies






		def __iter__(     self     :Tuple  ):
						snake_case__	:      Union[str, Any]     		=	[]
						for task in range(self.n_tasks  ):
										# without strip, the model generate commented codes ...
										prompts.append(self.tokenizer.eos_token + self.dataset[task]['''prompt'''].strip()  )
						snake_case__	:      Optional[int]     		=	self.tokenizer(__lowercase   ,padding=__lowercase   ,return_tensors='''pt'''  )
						for task in range(self.n_tasks  ):
										for _ in range(self.n_copies  ):
														yield {
														    "ids": outputs.input_ids[task],
														    "task_id": task,
														    "input_len": outputs.attention_mask[task].sum(),
														}


class       a							(			__lowerCamelCase				):



		def __init__(     self     :Union[str, Any]   ,__lowercase     :List[Any]   ,__lowercase     :int   ,__lowercase     :List[str]  ):
						snake_case__	:      List[str]     		=	start_length
						snake_case__	:      Any     		=	eof_strings
						snake_case__	:      Optional[Any]     		=	tokenizer






		def __call__(     self     :str   ,__lowercase     :str   ,__lowercase     :Any   ,**__lowercase     :Optional[int]  ):
						snake_case__	:      Tuple     		=	self.tokenizer.batch_decode(input_ids[:, self.start_length :]  )
						snake_case__	:      Union[str, Any]     		=	[]
						for decoded_generation in decoded_generations:
										done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings  )  )
						return all(__lowercase  )



def 			_lowerCAmelCase       (					__lowerCAmelCase      )   ->					List[Any]:


				"""simple docstring"""

				snake_case__	:      List[Any]     		=	re.split('''(%s)''' % '''|'''.join(__lowerCAmelCase      )  ,  __lowerCAmelCase      )
				# last string should be ""
				return "".join(string_list[:-2]      )



def 			_lowerCAmelCase       (					__lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase=20  ,  **__lowerCAmelCase      )   ->					List[Any]:


				"""simple docstring"""

				snake_case__	:      Optional[int]     		=	defaultdict(__lowerCAmelCase      )  # dict of list of generated tokens
				for step, batch in tqdm(enumerate(__lowerCAmelCase      )      ):
								with torch.no_grad():
												snake_case__	:      List[Any]     		=	batch['''ids'''].shape[-1]
												snake_case__	:      str     		=	accelerator.unwrap_model(__lowerCAmelCase      ).generate(
												    input_ids=batch['''ids'''][:, : batch['''input_len''']]  ,  num_return_sequences=__lowerCAmelCase  ,  **__lowerCAmelCase      )
												# each task is generated batch_size times
												snake_case__	:      str     		=	batch['''task_id'''].repeat(__lowerCAmelCase      )
												snake_case__	:      Optional[int]     		=	accelerator.pad_across_processes(
												    __lowerCAmelCase  ,  dim=1  ,  pad_index=tokenizer.pad_token_id      )

												snake_case__		, snake_case__	:      Any     		=	accelerator.gather((generated_tokens, generated_tasks)      )
												snake_case__	:      Union[str, Any]     		=	generated_tokens.cpu().numpy()
												snake_case__	:      Optional[Any]     		=	generated_tasks.cpu().numpy()

												for task, generated_tokens in zip(__lowerCAmelCase  ,  __lowerCAmelCase      ):
																gen_token_dict[task].append(__lowerCAmelCase      )

				snake_case__	:      List[Any]     		=	[[] for _ in range(__lowerCAmelCase      )]
				for task, generated_tokens in gen_token_dict.items():
								for s in generated_tokens:
												snake_case__	:      int     		=	tokenizer.decode(__lowerCAmelCase  ,  skip_special_tokens=__lowerCAmelCase  ,  clean_up_tokenization_spaces=__lowerCAmelCase      )
												code_gens[task].append(remove_last_block(__lowerCAmelCase      )      )
				return code_gens



def 			_lowerCAmelCase       (					)   ->					Optional[int]:


				"""simple docstring"""

				snake_case__	:      List[str]     		=	HfArgumentParser(__lowerCAmelCase      )
				snake_case__	:      Optional[Any]     		=	parser.parse_args()

				transformers.logging.set_verbosity_error()
				# enables code execution in code_eval metric
				snake_case__	:      Union[str, Any]     		=	args.HF_ALLOW_CODE_EVAL
				# make sure tokenizer plays nice with multiprocessing
				snake_case__	:      Union[str, Any]     		=	'''false'''

				if args.num_workers is None:
								snake_case__	:      Union[str, Any]     		=	multiprocessing.cpu_count()

				# Use dataset load to feed to accelerate
				snake_case__	:      Dict     		=	Accelerator()
				set_seed(args.seed  ,  device_specific=__lowerCAmelCase      )

				# Load model and tokenizer
				snake_case__	:      str     		=	AutoTokenizer.from_pretrained(args.model_ckpt      )
				snake_case__	:      List[str]     		=	tokenizer.eos_token
				snake_case__	:      Optional[Any]     		=	AutoModelForCausalLM.from_pretrained(args.model_ckpt      )

				# Generation settings
				snake_case__	:      List[Any]     		=	{
				    '''do_sample''': args.do_sample,
				    '''temperature''': args.temperature,
				    '''max_new_tokens''': args.max_new_tokens,
				    '''top_p''': args.top_p,
				    '''top_k''': args.top_k,
				    '''stopping_criteria''': StoppingCriteriaList([EndOfFunctionCriteria(0  ,  __lowerCAmelCase  ,  __lowerCAmelCase      )]      ),
				}

				# Load evaluation dataset and metric
				snake_case__	:      List[str]     		=	load_dataset('''openai_humaneval'''      )
				snake_case__	:      str     		=	load_metric('''code_eval'''      )

				snake_case__	:      List[Any]     		=	args.num_tasks if args.num_tasks is not None else len(human_eval['''test''']      )
				snake_case__	:      str     		=	args.n_samples // args.batch_size

				snake_case__	:      Tuple     		=	TokenizedDataset(__lowerCAmelCase  ,  human_eval['''test''']  ,  n_copies=__lowerCAmelCase  ,  n_tasks=__lowerCAmelCase      )
				# do not confuse args.batch_size, which is actually the num_return_sequences
				snake_case__	:      Tuple     		=	DataLoader(__lowerCAmelCase  ,  batch_size=1      )

				# Run a quick test to see if code evaluation is enabled
				try:
								snake_case__	:      Optional[Any]     		=	code_eval_metric.compute(references=['''''']  ,  predictions=[['''''']]      )
				except ValueError as exception:
								print(
								    '''Code evaluation not enabled. Read the warning below carefully and then use `--HF_ALLOW_CODE_EVAL="1"`'''
								    ''' flag to enable code evaluation.'''      )
								raise exception

				snake_case__		, snake_case__	:      Optional[Any]     		=	accelerator.prepare(__lowerCAmelCase  ,  __lowerCAmelCase      )

				snake_case__	:      Dict     		=	complete_code(
				    __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  n_tasks=__lowerCAmelCase  ,  batch_size=args.batch_size  ,  **__lowerCAmelCase  ,  )

				if accelerator.is_main_process:
								snake_case__	:      Optional[Any]     		=	[]

								for task in tqdm(range(__lowerCAmelCase      )      ):
												snake_case__	:      str     		=	human_eval['''test'''][task]['''test''']
												snake_case__	:      int     		=	f"""check({human_eval['test'][task]['entry_point']})"""
												references.append('''\n''' + test_func + '''\n''' + entry_point      )

								# Evaluate completions with "code_eval" metric
								snake_case__		, snake_case__	:      Any     		=	code_eval_metric.compute(
								    references=__lowerCAmelCase  ,  predictions=__lowerCAmelCase  ,  num_workers=args.num_workers      )
								print(f"""Results: {pass_at_k}"""      )

								# Save results to json file
								with open(args.output_file  ,  '''w'''      ) as fp:
												json.dump(__lowerCAmelCase  ,  __lowerCAmelCase      )


# For some reason the folliwng seems to be necessary sometimes for code_eval to work nice with multiprocessing
# https://stackoverflow.com/questions/60804599/python-multiprocessing-keeps-spawning-the-whole-script
if __name__ == "__main__":
				main()

