


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 : str		=				['''\nclass''', '''\ndef''', '''\n#''', '''\n@''', '''\nprint''', '''\nif''']
class 	lowerCamelCase_       (				_UpperCamelCase				):



					'''simple docstring'''







					def __init__(				self ,       snake_case_ ,       snake_case_ ,       snake_case_=None ,       snake_case_=1						)   ->			Tuple:







								'''simple docstring'''

								__lowercase 					=   tokenizer
								__lowercase 					=   dataset
								__lowercase 					=   len(__a						) if n_tasks is None else n_tasks
								__lowercase 					=   n_copies






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







								'''simple docstring'''

								__lowercase 					=   []
								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()						)
								__lowercase 					=   self.tokenizer(__a ,       padding=__a ,       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 	lowerCamelCase_       (				_UpperCamelCase				):



					'''simple docstring'''







					def __init__(				self ,       snake_case_ ,       snake_case_ ,       snake_case_						)   ->			int:







								'''simple docstring'''

								__lowercase 					=   start_length
								__lowercase 					=   eof_strings
								__lowercase 					=   tokenizer






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







								'''simple docstring'''

								__lowercase 					=   self.tokenizer.batch_decode(input_ids[:, self.start_length :]						)
								__lowercase 					=   []
								for decoded_generation in decoded_generations:
											done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings						)						)
								return all(__a						)




def 		lowercase_  (     _UpperCamelCase  ):


			'''simple docstring'''







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


def 		lowercase_  (     _UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase=20       ,			**_UpperCamelCase  ):


			'''simple docstring'''







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

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

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

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


def 		lowercase_  (     ):


			'''simple docstring'''







			__lowercase 					=   HfArgumentParser(lowercase_  )
			__lowercase 					=   parser.parse_args()

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

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

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

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

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

			# Load evaluation dataset and metric
			__lowercase 					=   load_dataset('''openai_humaneval'''  )
			__lowercase 					=   load_metric('''code_eval'''  )

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

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

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

			__lowercase 					=   accelerator.prepare(lowercase_       ,			lowercase_  )

			__lowercase 					=   complete_code(
			    lowercase_       ,			lowercase_       ,			lowercase_       ,			lowercase_       ,			n_tasks=lowercase_       ,			batch_size=args.batch_size       ,			**lowercase_       ,			)

			if accelerator.is_main_process:
						__lowercase 					=   []

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

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


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








import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
	import jax
@require_flax
class 	lowerCamelCase_       (				lowerCAmelCase__     ,     unittest.TestCase				):



					'''simple docstring'''







					__UpperCAmelCase            =				FlaxAutoencoderKL






					@property
					def A      (				self						)   ->			Optional[int]:







								'''simple docstring'''

								__lowercase 					=   4
								__lowercase 					=   3
								__lowercase 					=   (3_2, 3_2)

								__lowercase 					=   jax.random.PRNGKey(0						)
								__lowercase 					=   jax.random.uniform(snake_case_ ,       ((batch_size, num_channels) + sizes)						)

								return {"sample": image, "prng_key": prng_key}






					def A      (				self						)   ->			Any:







								'''simple docstring'''

								__lowercase 					=   {
								    '''block_out_channels''': [3_2, 6_4],
								    '''in_channels''': 3,
								    '''out_channels''': 3,
								    '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
								    '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
								    '''latent_channels''': 4,
								}
								__lowercase 					=   self.dummy_input
								return init_dict, inputs_dict


