





'''simple docstring'''






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


__UpperCamelCase					    =	["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif"]




class  _A (					__lowercase		):

		def __init__(   self					:   Tuple  ,      __magic_name__					:   Dict  ,      __magic_name__					:   Tuple  ,      __magic_name__					:   Tuple=None  ,      __magic_name__					:   Optional[Any]=1			)					->    Optional[Any]:


					"""simple docstring"""

					__snake_case					:      Any							      =     tokenizer
					__snake_case					:      Dict							      =     dataset
					__snake_case					:      Optional[Any]							      =     len(__magic_name__			) if n_tasks is None else n_tasks
					__snake_case					:      List[Any]							      =     n_copies

		def __iter__(   self					:   int			)					->    Optional[Any]:


					"""simple docstring"""

					__snake_case					:      Tuple							      =     []
					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(__magic_name__  ,      padding=__magic_name__  ,      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 (					__lowercase		):

		def __init__(   self					:   Any  ,      __magic_name__					:   Optional[Any]  ,      __magic_name__					:   Optional[int]  ,      __magic_name__					:   str			)					->    Tuple:


					"""simple docstring"""

					__snake_case					:      int							      =     start_length
					__snake_case					:      Union[str, Any]							      =     eof_strings
					__snake_case					:      Dict							      =     tokenizer

		def __call__(   self					:   str  ,      __magic_name__					:   str  ,      __magic_name__					:   Union[str, Any]  ,      **__magic_name__					:   List[Any]			)					->    List[Any]:


					"""simple docstring"""

					__snake_case					:      Optional[Any]							      =     self.tokenizer.batch_decode(input_ids[:, self.start_length :]			)
					__snake_case					:      int							      =     []
					for decoded_generation in decoded_generations:
								done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings			)			)
					return all(__magic_name__			)







def 					_a						(       _lowerCamelCase ) ->				Any:



			"""simple docstring"""


			__snake_case					:      Optional[int]							      =     re.split("""(%s)""" % """|""".join(_lowerCamelCase )    ,	_lowerCamelCase )
			# last string should be ""
			return "".join(string_list[:-2] )







def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase=20    ,	**_lowerCamelCase ) ->				int:



			"""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					:      int							      =     batch["""ids"""].shape[-1]
									__snake_case					:      Dict							      =     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					:      int							      =     batch["""task_id"""].repeat(_lowerCamelCase )
									__snake_case					:      Tuple							      =     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					:      List[str]							      =     generated_tokens.cpu().numpy()
									__snake_case					:      Any							      =     generated_tasks.cpu().numpy()

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

			__snake_case					:      List[str]							      =     [[] 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 					_a						(       ) ->				List[Any]:



			"""simple docstring"""


			__snake_case					:      int							      =     HfArgumentParser(_lowerCamelCase )
			__snake_case					:      Dict							      =     parser.parse_args()

			transformers.logging.set_verbosity_error()
			# enables code execution in code_eval metric
			__snake_case					:      Optional[int]							      =     args.HF_ALLOW_CODE_EVAL
			# make sure tokenizer plays nice with multiprocessing
			__snake_case					:      int							      =     """false"""

			if args.num_workers is None:
						__snake_case					:      Dict							      =     multiprocessing.cpu_count()

			# Use dataset load to feed to accelerate
			__snake_case					:      str							      =     Accelerator()
			set_seed(args.seed    ,	device_specific=_lowerCamelCase )

			# Load model and tokenizer
			__snake_case					:      Any							      =     AutoTokenizer.from_pretrained(args.model_ckpt )
			__snake_case					:      Dict							      =     tokenizer.eos_token
			__snake_case					:      str							      =     AutoModelForCausalLM.from_pretrained(args.model_ckpt )

			# Generation settings
			__snake_case					:      Optional[int]							      =     {
			    """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					:      Dict							      =     load_dataset("""openai_humaneval""" )
			__snake_case					:      int							      =     load_metric("""code_eval""" )

			__snake_case					:      int							      =     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					:      List[str]							      =     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					:      Dict							      =     DataLoader(_lowerCamelCase    ,	batch_size=1 )

			# Run a quick test to see if code evaluation is enabled
			try:
						__snake_case					:      int							      =     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					:      List[str]							      =     accelerator.prepare(_lowerCamelCase    ,	_lowerCamelCase )

			__snake_case					:      Optional[int]							      =     complete_code(
			    _lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase    ,	n_tasks=_lowerCamelCase    ,	batch_size=args.batch_size    ,	**_lowerCamelCase    ,	)

			if accelerator.is_main_process:
						__snake_case					:      str							      =     []

						for task in tqdm(range(_lowerCamelCase ) ):
									__snake_case					:      Union[str, Any]							      =     human_eval["""test"""][task]["""test"""]
									__snake_case					:      List[str]							      =     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					:      Optional[int]							      =     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()









'''simple docstring'''






import argparse
import json
import os

import torch
from torch import nn

from transformers import NllbMoeConfig, NllbMoeModel
from transformers.modeling_utils import dtype_byte_size
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME







def 					_a						(       _lowerCamelCase ) ->				List[Any]:



			"""simple docstring"""


			__snake_case					:      Union[str, Any]							      =     [
			    """encoder.version""",
			    """decoder.version""",
			    """model.encoder.version""",
			    """model.decoder.version""",
			    """decoder.output_projection.weight""",
			    """_float_tensor""",
			    """encoder.embed_positions._float_tensor""",
			    """decoder.embed_positions._float_tensor""",
			]
			for k in ignore_keys:
						state_dict.pop(_lowerCamelCase    ,	_lowerCamelCase )







def 					_a						(       _lowerCamelCase ) ->				List[str]:



			"""simple docstring"""


			__snake_case      ,     __snake_case					:      Dict							      =     emb.weight.shape
			__snake_case					:      Optional[int]							      =     nn.Linear(_lowerCamelCase    ,	_lowerCamelCase    ,	bias=_lowerCamelCase )
			__snake_case					:      Union[str, Any]							      =     emb.weight.data
			return lin_layer







def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase=None ) ->				Union[str, Any]:



			"""simple docstring"""


			__snake_case					:      Any							      =     {}
			for old_key in state_dict.keys():
						__snake_case					:      Union[str, Any]							      =     old_key
						if "moe_layer.experts." in key:
									if expert_idx is not None:
												__snake_case					:      Tuple							      =     key.replace("""moe_layer.experts.0"""    ,	F'''ffn.experts.expert_{expert_idx}''' )
									else:
												__snake_case					:      Optional[int]							      =     key.replace("""moe_layer.experts."""    ,	"""ffn.experts.expert_""" )
						if "gate" in key:
									__snake_case					:      Dict							      =     key.replace(""".moe_layer.gate.wg"""    ,	""".ffn.router.classifier""" )
						if "fc2" and "experts" not in key:
									__snake_case					:      Union[str, Any]							      =     key.replace(""".fc2."""    ,	""".ffn.fc2.""" )
						if "fc1" and "experts" not in key:
									__snake_case					:      Optional[int]							      =     key.replace(""".fc1."""    ,	""".ffn.fc1.""" )
						if ".encoder_attn." in key:
									__snake_case					:      Tuple							      =     key.replace(""".encoder_attn."""    ,	""".cross_attention.""" )
						if "encoder_attn_layer_norm" in key:
									__snake_case					:      Union[str, Any]							      =     key.replace("""encoder_attn_layer_norm"""    ,	"""cross_attention_layer_norm""" )
						if "final_layer_norm" in key:
									__snake_case					:      str							      =     key.replace("""final_layer_norm"""    ,	"""ff_layer_norm""" )
						__snake_case					:      str							      =     state_dict[old_key]
			return new_dict







def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase = WEIGHTS_NAME ) ->				Dict:



			"""simple docstring"""


			__snake_case					:      Optional[int]							      =     []
			__snake_case					:      Dict							      =     0
			os.makedirs(_lowerCamelCase    ,	exist_ok=_lowerCamelCase )

			for expert in range(_lowerCamelCase ):
						__snake_case					:      Tuple							      =     switch_checkpoint_path + F'''-rank-{expert}.pt'''
						if os.path.isfile(_lowerCamelCase ):
									__snake_case					:      Dict							      =     torch.load(_lowerCamelCase )["""model"""]
									remove_ignore_keys_(_lowerCamelCase )
									__snake_case					:      Optional[Any]							      =     rename_fairseq_keys(_lowerCamelCase    ,	_lowerCamelCase )
									__snake_case					:      List[Any]							      =     os.path.join(
									    _lowerCamelCase    ,	weights_name.replace(""".bin"""    ,	F'''-{len(_lowerCamelCase )+1:05d}-of-???.bin''' ) )
									torch.save(_lowerCamelCase    ,	_lowerCamelCase )
									sharded_state_dicts.append(expert_state.keys() )
									total_size += sum([value.numel() for key, value in expert_state.items()] ) * dtype_byte_size(
									    expert_state[list(_lowerCamelCase )[0]].dtype )

    # Add the last block
			__snake_case					:      Optional[Any]							      =     os.path.join(_lowerCamelCase    ,	weights_name.replace(""".bin"""    ,	F'''-{len(_lowerCamelCase )+1:05d}-of-???.bin''' ) )
			__snake_case					:      str							      =     torch.load(switch_checkpoint_path + """-shared.pt""" )["""model"""]
			remove_ignore_keys_(_lowerCamelCase )
			__snake_case					:      Optional[Any]							      =     rename_fairseq_keys(_lowerCamelCase    ,	_lowerCamelCase )
			__snake_case					:      List[str]							      =     shared_weights["""decoder.embed_tokens.weight"""]
			sharded_state_dicts.append(shared_weights.keys() )

			# If we only have the shared weights (dummy model/experts saved on the same file)
			if len(_lowerCamelCase ) == 1:
						__snake_case					:      Optional[Any]							      =     os.path.join(_lowerCamelCase    ,	_lowerCamelCase )
						torch.save(_lowerCamelCase    ,	_lowerCamelCase )
						return {weights_name: sharded_state_dicts[0]}, None
			else:
						torch.save(_lowerCamelCase    ,	_lowerCamelCase )
			# Otherwise, let's build the index
			__snake_case					:      Tuple							      =     {}
			for idx, shard in enumerate(_lowerCamelCase ):
						__snake_case					:      Any							      =     weights_name.replace(""".bin"""    ,	F'''-{idx+1:05d}-of-{len(_lowerCamelCase ):05d}.bin''' )
						__snake_case					:      int							      =     os.path.join(_lowerCamelCase    ,	weights_name.replace(""".bin"""    ,	F'''-{idx+1:05d}-of-???.bin''' ) )
						os.rename(_lowerCamelCase    ,	os.path.join(_lowerCamelCase    ,	_lowerCamelCase ) )
						for key in shard:
									__snake_case					:      str							      =     shard_file

    # Add the metadata
			__snake_case					:      Optional[Any]							      =     {"""total_size""": total_size}
			__snake_case					:      int							      =     {"""metadata""": metadata, """weight_map""": weight_map}

			with open(os.path.join(_lowerCamelCase    ,	_lowerCamelCase )    ,	"""w"""    ,	encoding="""utf-8""" ) as f:
						__snake_case					:      Union[str, Any]							      =     json.dumps(_lowerCamelCase    ,	indent=2    ,	sort_keys=_lowerCamelCase ) + """\n"""
						f.write(_lowerCamelCase )

			return metadata, index


if __name__ == "__main__":
	__UpperCamelCase					    =	argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    "--nllb_moe_checkpoint_path",
	    default="/home/arthur_huggingface_co/fairseq/weights/checkpoints/model_moe_54b/checkpoint_2_300000",
	    type=str,
	    required=False,
	    help="Path to a directory containing a folder per layer. Follows the original Google format.",
	)
	parser.add_argument("--dtype", default="float32", type=str, required=False, help="dtype of the saved model")
	parser.add_argument(
	    "--pytorch_dump_folder_path",
	    default="/home/arthur_huggingface_co/fairseq/weights/checkpoints/hf-converted-moe-54b",
	    type=str,
	    required=False,
	    help="Path to the output pytorch model.",
	)
	__UpperCamelCase					    =	parser.parse_args()
	__UpperCamelCase    ,		__UpperCamelCase					    =	shard_on_the_fly(
	    args.nllb_moe_checkpoint_path,
	    args.pytorch_dump_folder_path,
	    128,
	    args.dtype,
	)

	__UpperCamelCase					    =	NllbMoeConfig.from_pretrained(
	    "facebook/nllb-200-3.3B", encoder_sparse_step=4, decoder_sparse_step=4, num_experts=128
	)
	config.save_pretrained(args.pytorch_dump_folder_path)
	__UpperCamelCase					    =	NllbMoeModel.from_pretrained(args.pytorch_dump_folder_path)
	print("Done")
	model.save_pretrained(args.pytorch_dump_folder_path)


