


"""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__				:Dict			   =	["""\nclass""", """\ndef""", """\n#""", """\n@""", """\nprint""", """\nif"""]






class      A(  lowerCamelCase__		):

		"""simple docstring"""





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

					"""simple docstring"""



					_UpperCamelCase					:Any					     =	tokenizer
					_UpperCamelCase					:Union[str, Any]					     =	dataset
					_UpperCamelCase					:List[Any]					     =	len(SCREAMING_SNAKE_CASE__   ) if n_tasks is None else n_tasks
					_UpperCamelCase					:Any					     =	n_copies





		def __iter__( self   )      ->	str:

					"""simple docstring"""



					_UpperCamelCase					: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()   )
					_UpperCamelCase					:Any					     =	self.tokenizer(SCREAMING_SNAKE_CASE__					,				padding=SCREAMING_SNAKE_CASE__					,				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__		):

		"""simple docstring"""





		def __init__( self					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__   )      ->	Dict:

					"""simple docstring"""



					_UpperCamelCase					:int					     =	start_length
					_UpperCamelCase					:List[str]					     =	eof_strings
					_UpperCamelCase					:Union[str, Any]					     =	tokenizer





		def __call__( self					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				**SCREAMING_SNAKE_CASE__   )      ->	List[Any]:

					"""simple docstring"""



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





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





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

									_UpperCamelCase     ,      _UpperCamelCase					:Any					     =	accelerator.gather((generated_tokens, generated_tasks)		)
									_UpperCamelCase					:Any					     =	generated_tokens.cpu().numpy()
									_UpperCamelCase					:Dict					     =	generated_tasks.cpu().numpy()

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

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





def 				A_       ( )					->  str:
			# Setup configuration
			_UpperCamelCase					:Dict					     =	HfArgumentParser(snake_case__		)
			_UpperCamelCase					:Optional[Any]					     =	parser.parse_args()

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

			if args.num_workers is None:
						_UpperCamelCase					:Optional[Any]					     =	multiprocessing.cpu_count()

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

			# Load model and tokenizer
			_UpperCamelCase					:Tuple					     =	AutoTokenizer.from_pretrained(args.model_ckpt		)
			_UpperCamelCase					:int					     =	tokenizer.eos_token
			_UpperCamelCase					:int					     =	AutoModelForCausalLM.from_pretrained(args.model_ckpt		)

			# Generation settings
			_UpperCamelCase					:Union[str, 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       ,					snake_case__       ,					snake_case__		)]		),
			}

			# Load evaluation dataset and metric
			_UpperCamelCase					:Optional[Any]					     =	load_dataset('''openai_humaneval'''		)
			_UpperCamelCase					:Tuple					     =	load_metric('''code_eval'''		)

			_UpperCamelCase					:Optional[int]					     =	args.num_tasks if args.num_tasks is not None else len(human_eval['''test''']		)
			_UpperCamelCase					:Union[str, Any]					     =	args.n_samples // args.batch_size

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

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

			_UpperCamelCase     ,      _UpperCamelCase					:Tuple					     =	accelerator.prepare(snake_case__       ,					snake_case__		)

			_UpperCamelCase					:Any					     =	complete_code(
			    snake_case__       ,					snake_case__       ,					snake_case__       ,					snake_case__       ,					n_tasks=snake_case__       ,					batch_size=args.batch_size       ,					**snake_case__       ,					)

			if accelerator.is_main_process:
						_UpperCamelCase					:Optional[int]					     =	[]

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

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


# 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 unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate





def 				A_       ( snake_case__		)					->  Union[str, Any]:
			return x + 2






class      A(  unittest.TestCase		):

		"""simple docstring"""





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

					"""simple docstring"""



					_UpperCamelCase					:Union[str, Any]					     =	'''x = 3'''
					_UpperCamelCase					:List[str]					     =	{}
					_UpperCamelCase					:int					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{}					,				state=SCREAMING_SNAKE_CASE__   )
					assert result == 3
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3}   )

					_UpperCamelCase					:Tuple					     =	'''x = y'''
					_UpperCamelCase					:Optional[int]					     =	{'''y''': 5}
					_UpperCamelCase					:Optional[Any]					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{}					,				state=SCREAMING_SNAKE_CASE__   )
					# evaluate returns the value of the last assignment.
					assert result == 5
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 5, '''y''': 5}   )
		def   _UpperCamelCase( self   )      ->	int:

					"""simple docstring"""



					_UpperCamelCase					:str					     =	'''y = add_two(x)'''
					_UpperCamelCase					:Optional[int]					     =	{'''x''': 3}
					_UpperCamelCase					:Union[str, Any]					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{'''add_two''': add_two}					,				state=SCREAMING_SNAKE_CASE__   )
					assert result == 5
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3, '''y''': 5}   )

					# Won't work without the tool
					with CaptureStdout() as out:
								_UpperCamelCase					:Union[str, Any]					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{}					,				state=SCREAMING_SNAKE_CASE__   )
					assert result is None
					assert "tried to execute add_two" in out.out
		def   _UpperCamelCase( self   )      ->	List[Any]:

					"""simple docstring"""



					_UpperCamelCase					:Any					     =	'''x = 3'''
					_UpperCamelCase					:int					     =	{}
					_UpperCamelCase					:List[Any]					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{}					,				state=SCREAMING_SNAKE_CASE__   )
					assert result == 3
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3}   )
		def   _UpperCamelCase( self   )      ->	List[Any]:

					"""simple docstring"""



					_UpperCamelCase					:str					     =	'''test_dict = {\'x\': x, \'y\': add_two(x)}'''
					_UpperCamelCase					:str					     =	{'''x''': 3}
					_UpperCamelCase					:Dict					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{'''add_two''': add_two}					,				state=SCREAMING_SNAKE_CASE__   )
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3, '''y''': 5}   )
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3, '''test_dict''': {'''x''': 3, '''y''': 5}}   )
		def   _UpperCamelCase( self   )      ->	Tuple:

					"""simple docstring"""



					_UpperCamelCase					:int					     =	'''x = 3\ny = 5'''
					_UpperCamelCase					:int					     =	{}
					_UpperCamelCase					:Optional[int]					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{}					,				state=SCREAMING_SNAKE_CASE__   )
					# evaluate returns the value of the last assignment.
					assert result == 5
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3, '''y''': 5}   )
		def   _UpperCamelCase( self   )      ->	Optional[int]:

					"""simple docstring"""



					_UpperCamelCase					:Union[str, Any]					     =	'''text = f\'This is x: {x}.\''''
					_UpperCamelCase					:List[str]					     =	{'''x''': 3}
					_UpperCamelCase					:Any					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{}					,				state=SCREAMING_SNAKE_CASE__   )
					# evaluate returns the value of the last assignment.
					assert result == "This is x: 3."
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3, '''text''': '''This is x: 3.'''}   )
		def   _UpperCamelCase( self   )      ->	Union[str, Any]:

					"""simple docstring"""



					_UpperCamelCase					:Optional[Any]					     =	'''if x <= 3:\n    y = 2\nelse:\n    y = 5'''
					_UpperCamelCase					:Dict					     =	{'''x''': 3}
					_UpperCamelCase					:Any					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{}					,				state=SCREAMING_SNAKE_CASE__   )
					# evaluate returns the value of the last assignment.
					assert result == 2
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3, '''y''': 2}   )

					_UpperCamelCase					:Union[str, Any]					     =	{'''x''': 8}
					_UpperCamelCase					:Optional[int]					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{}					,				state=SCREAMING_SNAKE_CASE__   )
					# evaluate returns the value of the last assignment.
					assert result == 5
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 8, '''y''': 5}   )
		def   _UpperCamelCase( self   )      ->	Tuple:

					"""simple docstring"""



					_UpperCamelCase					:List[str]					     =	'''test_list = [x, add_two(x)]'''
					_UpperCamelCase					:Optional[int]					     =	{'''x''': 3}
					_UpperCamelCase					:str					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{'''add_two''': add_two}					,				state=SCREAMING_SNAKE_CASE__   )
					self.assertListEqual(SCREAMING_SNAKE_CASE__					,				[3, 5]   )
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3, '''test_list''': [3, 5]}   )
		def   _UpperCamelCase( self   )      ->	Dict:

					"""simple docstring"""



					_UpperCamelCase					:List[Any]					     =	'''y = x'''
					_UpperCamelCase					:List[str]					     =	{'''x''': 3}
					_UpperCamelCase					:Union[str, Any]					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{}					,				state=SCREAMING_SNAKE_CASE__   )
					assert result == 3
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3, '''y''': 3}   )
		def   _UpperCamelCase( self   )      ->	Tuple:

					"""simple docstring"""



					_UpperCamelCase					:Tuple					     =	'''test_list = [x, add_two(x)]\ntest_list[1]'''
					_UpperCamelCase					:Optional[Any]					     =	{'''x''': 3}
					_UpperCamelCase					:Any					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{'''add_two''': add_two}					,				state=SCREAMING_SNAKE_CASE__   )
					assert result == 5
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3, '''test_list''': [3, 5]}   )

					_UpperCamelCase					:str					     =	'''test_dict = {\'x\': x, \'y\': add_two(x)}\ntest_dict[\'y\']'''
					_UpperCamelCase					:Optional[int]					     =	{'''x''': 3}
					_UpperCamelCase					:Union[str, Any]					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{'''add_two''': add_two}					,				state=SCREAMING_SNAKE_CASE__   )
					assert result == 5
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 3, '''test_dict''': {'''x''': 3, '''y''': 5}}   )





		def   _UpperCamelCase( self   )      ->	int:

					"""simple docstring"""



					_UpperCamelCase					:str					     =	'''x = 0\nfor i in range(3):\n    x = i'''
					_UpperCamelCase					:Tuple					     =	{}
					_UpperCamelCase					:Union[str, Any]					     =	evaluate(SCREAMING_SNAKE_CASE__					,				{'''range''': range}					,				state=SCREAMING_SNAKE_CASE__   )
					assert result == 2
					self.assertDictEqual(SCREAMING_SNAKE_CASE__					,				{'''x''': 2, '''i''': 2}   )

