




import argparse
import collections

import numpy as np
import torch
from flax import traverse_util
from tax import checkpoints

from transformers import MTaConfig, UMTaEncoderModel, UMTaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()




def     _UpperCamelCase       (				A	:      Union[str, Any],						A	:      Tuple,						A	:      int     ):






			'''simple docstring'''







			return params[F"""{prefix}/{prefix}/relpos_bias/rel_embedding"""][:, i, :]




def     _UpperCamelCase       (				A	:      Union[str, Any],						A	:      List[Any],						A	:      Dict,						A	:      Optional[int]="attention"     ):






			'''simple docstring'''







			a					=    a					=    np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/key/kernel"""][:, i, :, :]     )
			a					=    k_tmp.reshape(k_tmp.shape[0],						k_tmp.shape[1] * k_tmp.shape[2]     )
			a					=    np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/out/kernel"""][:, i, :, :]     )
			a					=    o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1],						o_tmp.shape[2]     )
			a					=    np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/query/kernel"""][:, i, :, :]     )
			a					=    q_tmp.reshape(q_tmp.shape[0],						q_tmp.shape[1] * q_tmp.shape[2]     )
			a					=    np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/value/kernel"""][:, i, :, :]     )
			a					=    v_tmp.reshape(v_tmp.shape[0],						v_tmp.shape[1] * v_tmp.shape[2]     )
			return k, o, q, v




def     _UpperCamelCase       (				A	:      Tuple,						A	:      Any,						A	:      int,						A	:      Optional[int]=False     ):






			'''simple docstring'''







			if split_mlp_wi:
						a					=    params[F"""{prefix}/{prefix}/mlp/wi_0/kernel"""][:, i, :]
						a					=    params[F"""{prefix}/{prefix}/mlp/wi_1/kernel"""][:, i, :]
						a					=    (wi_a, wi_a)
			else:
						a					=    params[F"""{prefix}/{prefix}/mlp/wi/kernel"""][:, i, :]

			a					=    params[F"""{prefix}/{prefix}/mlp/wo/kernel"""][:, i, :]
			return wi, wo




def     _UpperCamelCase       (				A	:      Tuple,						A	:      List[str],						A	:      Any,						A	:      Any     ):






			'''simple docstring'''







			return params[F"""{prefix}/{prefix}/{layer_name}/scale"""][:, i]




def     _UpperCamelCase       (				A	:      List[str],						*, A	:      int,						A	:      Optional[Any],						A	:      List[Any] = False     ):






			'''simple docstring'''







			a					=    traverse_util.flatten_dict(variables["target"]     )
			a					=    {"/".join(a__     ): v for k, v in old.items()}

			# v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
			a					=    "encoder/encoder/mlp/wi_0/kernel" in old
			print("Split MLP:",						a__     )

			a					=    collections.OrderedDict()

			# Shared embeddings.
			a					=    old["token_embedder/embedding"]

			# Encoder.
			for i in range(a__     ):
						# Block i, layer 0 (Self Attention).
						a					=    tax_layer_norm_lookup(a__,						a__,						"encoder",						"pre_attention_layer_norm"     )
						a					,	a					,	a					,	a					=    tax_attention_lookup(a__,						a__,						"encoder",						"attention"     )
						a					=    layer_norm
						a					=    k.T
						a					=    o.T
						a					=    q.T
						a					=    v.T

						# Block i, layer 1 (MLP).
						a					=    tax_layer_norm_lookup(a__,						a__,						"encoder",						"pre_mlp_layer_norm"     )
						a					,	a					=    tax_mlp_lookup(a__,						a__,						"encoder",						a__     )
						a					=    layer_norm
						if split_mlp_wi:
									a					=    wi[0].T
									a					=    wi[1].T
						else:
									a					=    wi.T
						a					=    wo.T
						if scalable_attention:
									# convert the rel_embedding of each layer
									a					=    tax_relpos_bias_lookup(
									    a__,						a__,						"encoder"     ).T

			a					=    old["encoder/encoder_norm/scale"]

			if not scalable_attention:
						a					=    tax_relpos_bias_lookup(
						    a__,						0,						"encoder"     ).T
						a					=    tax_relpos_bias_lookup(
						    a__,						0,						"decoder"     ).T

			if not is_encoder_only:
						# Decoder.
						for i in range(a__     ):
									# Block i, layer 0 (Self Attention).
									a					=    tax_layer_norm_lookup(a__,						a__,						"decoder",						"pre_self_attention_layer_norm"     )
									a					,	a					,	a					,	a					=    tax_attention_lookup(a__,						a__,						"decoder",						"self_attention"     )
									a					=    layer_norm
									a					=    k.T
									a					=    o.T
									a					=    q.T
									a					=    v.T

									# Block i, layer 1 (Cross Attention).
									a					=    tax_layer_norm_lookup(a__,						a__,						"decoder",						"pre_cross_attention_layer_norm"     )
									a					,	a					,	a					,	a					=    tax_attention_lookup(a__,						a__,						"decoder",						"encoder_decoder_attention"     )
									a					=    layer_norm
									a					=    k.T
									a					=    o.T
									a					=    q.T
									a					=    v.T

									# Block i, layer 2 (MLP).
									a					=    tax_layer_norm_lookup(a__,						a__,						"decoder",						"pre_mlp_layer_norm"     )
									a					,	a					=    tax_mlp_lookup(a__,						a__,						"decoder",						a__     )
									a					=    layer_norm
									if split_mlp_wi:
												a					=    wi[0].T
												a					=    wi[1].T
									else:
												a					=    wi.T
									a					=    wo.T

									if scalable_attention:
												# convert the rel_embedding of each layer
												a					=    tax_relpos_bias_lookup(a__,						a__,						"decoder"     ).T

						a					=    old["decoder/decoder_norm/scale"]

						# LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
						if "decoder/logits_dense/kernel" in old:
									a					=    old["decoder/logits_dense/kernel"].T

			return new




def     _UpperCamelCase       (				A	:      List[str],						A	:      Dict     ):






			'''simple docstring'''







			a					=    collections.OrderedDict([(k, torch.from_numpy(v.copy()     )) for (k, v) in converted_params.items()]     )

			# Add what is missing.
			if "encoder.embed_tokens.weight" not in state_dict:
						a					=    state_dict["shared.weight"]

			if not is_encoder_only:
						if "decoder.embed_tokens.weight" not in state_dict:
									a					=    state_dict["shared.weight"]

						if "lm_head.weight" not in state_dict:  # For old 1.0 models.
									print("Using shared word embeddings as lm_head."     )
									a					=    state_dict["shared.weight"]

			return state_dict




def     _UpperCamelCase       (				A	:      Any,						A	:      Union[str, Any],						A	:      Any,						A	:      int,						A	:      List[str]     ):






			'''simple docstring'''







			a					=    checkpoints.load_tax_checkpoint(a__     )
			a					=    convert_tax_to_pytorch(
			    a__,						num_layers=config.num_layers,						is_encoder_only=a__,						scalable_attention=a__     )
			a					=    make_state_dict(a__,						a__     )
			model.load_state_dict(a__,						strict=a__     )




def     _UpperCamelCase       (				A	:      Optional[Any],						A	:      str,						A	:      Optional[Any],						A	:      Optional[int] = False,						A	:      int = False,						):






			'''simple docstring'''







			a					=    MTaConfig.from_json_file(a__     )
			print(F"""Building PyTorch model from configuration: {config}"""     )
			# Non-v1.1 checkpoints could also use T5Model, but this works for all.
			# The v1.0 checkpoints will simply have an LM head that is the word embeddings.
			if is_encoder_only:
						a					=    UMTaEncoderModel(a__     )
			else:
						a					=    UMTaForConditionalGeneration(a__     )

			# Load weights from tf checkpoint
			load_tax_weights_in_ta(a__,						a__,						a__,						a__,						a__     )

			# Save pytorch-model
			print(F"""Save PyTorch model to {pytorch_dump_path}"""     )
			model.save_pretrained(a__     )

			# Verify that we can load the checkpoint.
			model.from_pretrained(a__     )
			print("Done"     )


if __name__ == "__main__":
		__lowerCAmelCase	:       Union[str, Any]     =	argparse.ArgumentParser(description='Converts a native T5X checkpoint into a PyTorch checkpoint.')
		# Required parameters
		parser.add_argument(
		    '--t5x_checkpoint_path', default=None, type=str, required=True, help='Path to the T5X checkpoint.'
		)
		parser.add_argument(
		    '--config_file',
		    default=None,
		    type=str,
		    required=True,
		    help='The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.',
		)
		parser.add_argument(
		    '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
		)
		parser.add_argument(
		    '--is_encoder_only', action='store_true', help='Check if the model is encoder-decoder model', default=False
		)
		parser.add_argument(
		    '--scalable_attention',
		    action='store_true',
		    help='Whether the model uses scaled attention (umt5 model)',
		    default=False,
		)
		__lowerCAmelCase	:       str     =	parser.parse_args()
		convert_tax_checkpoint_to_pytorch(
		    args.tax_checkpoint_path,
		    args.config_file,
		    args.pytorch_dump_path,
		    args.is_encoder_only,
		    args.scalable_attention,
		)












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


__lowerCAmelCase	:       Union[str, Any]     =	['\nclass', '\ndef', '\n#', '\n@', '\nprint', '\nif']




class 					snake_case__							(_UpperCamelCase     ):
				"""simple docstring"""



				def __init__(   self       :  Any      ,   __lowerCamelCase       :  List[str]      ,   __lowerCamelCase       :  Union[str, Any]      ,   __lowerCamelCase       :  Union[str, Any]=None      ,   __lowerCamelCase       :  Any=1							)      ->  Union[str, Any]:
							a					=    tokenizer
							a					=    dataset
							a					=    len(__lowerCamelCase							) if n_tasks is None else n_tasks
							a					=    n_copies





				def __iter__(   self       :  Tuple							)      ->  str:
							a					=    []
							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()							)
							a					=    self.tokenizer(__lowerCamelCase      ,   padding=__lowerCamelCase      ,   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 					snake_case__							(_UpperCamelCase     ):
				"""simple docstring"""



				def __init__(   self       :  int      ,   __lowerCamelCase       :  Dict      ,   __lowerCamelCase       :  Dict      ,   __lowerCamelCase       :  Tuple							)      ->  Optional[Any]:
							a					=    start_length
							a					=    eof_strings
							a					=    tokenizer





				def __call__(   self       :  int      ,   __lowerCamelCase       :  Dict      ,   __lowerCamelCase       :  List[str]      ,   **__lowerCamelCase       :  Optional[int]							)      ->  Optional[Any]:
							a					=    self.tokenizer.batch_decode(input_ids[:, self.start_length :]							)
							a					=    []
							for decoded_generation in decoded_generations:
										done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings							)							)
							return all(__lowerCamelCase							)




def     __magic_name__       (				A	:      List[Any]     ):






			'''simple docstring'''







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




def     __magic_name__       (				A	:      Union[str, Any],						A	:      Optional[Any],						A	:      List[Any],						A	:      Optional[Any],						A	:      List[str],						A	:      List[Any]=20,						**A	:      Union[str, Any]     ):






			'''simple docstring'''







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

									a					,	a					=    accelerator.gather((generated_tokens, generated_tasks)     )
									a					=    generated_tokens.cpu().numpy()
									a					=    generated_tasks.cpu().numpy()

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

			a					=    [[] for _ in range(A     )]
			for task, generated_tokens in gen_token_dict.items():
						for s in generated_tokens:
									a					=    tokenizer.decode(A,						skip_special_tokens=A,						clean_up_tokenization_spaces=A     )
									code_gens[task].append(remove_last_block(A     )     )
			return code_gens




def     __magic_name__       (				):






			'''simple docstring'''







			a					=    HfArgumentParser(A     )
			a					=    parser.parse_args()

			transformers.logging.set_verbosity_error()
			# enables code execution in code_eval metric
			a					=    args.HF_ALLOW_CODE_EVAL
			# make sure tokenizer plays nice with multiprocessing
			a					=    "false"

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

			# Use dataset load to feed to accelerate
			a					=    Accelerator()
			set_seed(args.seed,						device_specific=A     )

			# Load model and tokenizer
			a					=    AutoTokenizer.from_pretrained(args.model_ckpt     )
			a					=    tokenizer.eos_token
			a					=    AutoModelForCausalLM.from_pretrained(args.model_ckpt     )

			# Generation settings
			a					=    {
			    "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,						A,						A     )]     ),
			}

			# Load evaluation dataset and metric
			a					=    load_dataset("openai_humaneval"     )
			a					=    load_metric("code_eval"     )

			a					=    args.num_tasks if args.num_tasks is not None else len(human_eval["test"]     )
			a					=    args.n_samples // args.batch_size

			a					=    TokenizedDataset(A,						human_eval["test"],						n_copies=A,						n_tasks=A     )
			# do not confuse args.batch_size, which is actually the num_return_sequences
			a					=    DataLoader(A,						batch_size=1     )

			# Run a quick test to see if code evaluation is enabled
			try:
						a					=    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

			a					,	a					=    accelerator.prepare(A,						A     )

			a					=    complete_code(
			    A,						A,						A,						A,						n_tasks=A,						batch_size=args.batch_size,						**A,						)

			if accelerator.is_main_process:
						a					=    []

						for task in tqdm(range(A     )     ):
									a					=    human_eval["test"][task]["test"]
									a					=    F"""check({human_eval["test"][task]["entry_point"]})"""
									references.append("\n" + test_func + "\n" + entry_point     )

						# Evaluate completions with "code_eval" metric
						a					,	a					=    code_eval_metric.compute(
						    references=A,						predictions=A,						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(A,						A     )


# 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()






