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   _lowerCAmelCase		(     __SCREAMING_SNAKE_CASE				):







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







							def __iter__(				self     :    Tuple		):
										lowerCamelCase						:Dict						  =     []
										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()		)
										lowerCamelCase						:List[str]						  =     self.tokenizer(__snake_case	,					padding=__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   _lowerCAmelCase		(     __SCREAMING_SNAKE_CASE				):







							def __init__(				self     :    Optional[int]	,					__snake_case     :    Union[str, Any]	,					__snake_case     :    str	,					__snake_case     :    int		):
										lowerCamelCase						:Dict						  =     start_length
										lowerCamelCase						:Any						  =     eof_strings
										lowerCamelCase						:int						  =     tokenizer







							def __call__(				self     :    Any	,					__snake_case     :    str	,					__snake_case     :    Dict	,					**__snake_case     :    List[str]		):
										lowerCamelCase						:Union[str, Any]						  =     self.tokenizer.batch_decode(input_ids[:, self.start_length :]		)
										lowerCamelCase						: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(__snake_case		)





def  _lowerCamelCase					(   a_				: Tuple):
			lowerCamelCase						:Tuple						  =     re.split('''(%s)''' % '''|'''.join(a_)						,			a_)
			# last string should be ""
			return "".join(string_list[:-2])





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

									lowerCamelCase						:List[str]						  =     accelerator.gather((generated_tokens, generated_tasks))
									lowerCamelCase						:Tuple						  =     generated_tokens.cpu().numpy()
									lowerCamelCase						:Tuple						  =     generated_tasks.cpu().numpy()

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

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





def  _lowerCamelCase					(   ):
			# Setup configuration
			lowerCamelCase						:int						  =     HfArgumentParser(a_)
			lowerCamelCase						:Union[str, Any]						  =     parser.parse_args()

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

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

			# Use dataset load to feed to accelerate
			lowerCamelCase						:Union[str, Any]						  =     Accelerator()
			set_seed(args.seed						,			device_specific=a_)

			# Load model and tokenizer
			lowerCamelCase						:Union[str, Any]						  =     AutoTokenizer.from_pretrained(args.model_ckpt)
			lowerCamelCase						:str						  =     tokenizer.eos_token
			lowerCamelCase						:str						  =     AutoModelForCausalLM.from_pretrained(args.model_ckpt)

			# Generation settings
			lowerCamelCase						: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						,			a_						,			a_)]),
			}

			# Load evaluation dataset and metric
			lowerCamelCase						:int						  =     load_dataset('''openai_humaneval''')
			lowerCamelCase						:Dict						  =     load_metric('''code_eval''')

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

			lowerCamelCase						:Union[str, Any]						  =     TokenizedDataset(a_						,			human_eval['''test''']						,			n_copies=a_						,			n_tasks=a_)
			# do not confuse args.batch_size, which is actually the num_return_sequences
			lowerCamelCase						:Dict						  =     DataLoader(a_						,			batch_size=1)

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

			lowerCamelCase						:Union[str, Any]						  =     accelerator.prepare(a_						,			a_)

			lowerCamelCase						:int						  =     complete_code(
			    a_						,			a_						,			a_						,			a_						,			n_tasks=a_						,			batch_size=args.batch_size						,			**a_						,			)

			if accelerator.is_main_process:
						lowerCamelCase						:Dict						  =     []

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

						# Evaluate completions with "code_eval" metric
						lowerCamelCase						:Any						  =     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()




import json
import os
import unittest

from transformers import BatchEncoding, LEDTokenizer, LEDTokenizerFast
from transformers.models.led.tokenization_led import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers, require_torch
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin
@require_tokenizers
class   _lowerCAmelCase		(     __SCREAMING_SNAKE_CASE      ,   unittest.TestCase				):
							_UpperCAmelCase						    =     LEDTokenizer
							_UpperCAmelCase						    =     LEDTokenizerFast
							_UpperCAmelCase						    =     True







							def 				snake_case  (				self     :    Any		):
										super().setUp()
										lowerCamelCase						:Optional[int]						  =     [
										    '''l''',
										    '''o''',
										    '''w''',
										    '''e''',
										    '''r''',
										    '''s''',
										    '''t''',
										    '''i''',
										    '''d''',
										    '''n''',
										    '''\u0120''',
										    '''\u0120l''',
										    '''\u0120n''',
										    '''\u0120lo''',
										    '''\u0120low''',
										    '''er''',
										    '''\u0120lowest''',
										    '''\u0120newer''',
										    '''\u0120wider''',
										    '''<unk>''',
										]
										lowerCamelCase						:Any						  =     dict(zip(__snake_case	,					range(len(__snake_case		)		)		)		)
										lowerCamelCase						:List[str]						  =     ['''#version: 0.2''', '''\u0120 l''', '''\u0120l o''', '''\u0120lo w''', '''e r''', '''''']
										lowerCamelCase						:int						  =     {'''unk_token''': '''<unk>'''}

										lowerCamelCase						:Optional[int]						  =     os.path.join(self.tmpdirname	,					VOCAB_FILES_NAMES['''vocab_file''']		)
										lowerCamelCase						:Any						  =     os.path.join(self.tmpdirname	,					VOCAB_FILES_NAMES['''merges_file''']		)
										with open(self.vocab_file	,					'''w'''	,					encoding='''utf-8'''		) as fp:
													fp.write(json.dumps(__snake_case		) + '''\n'''		)
										with open(self.merges_file	,					'''w'''	,					encoding='''utf-8'''		) as fp:
													fp.write('''\n'''.join(__snake_case		)		)







							def 				snake_case  (				self     :    int	,					**__snake_case     :    int		):
										kwargs.update(self.special_tokens_map		)
										return self.tokenizer_class.from_pretrained(self.tmpdirname	,					**__snake_case		)







							def 				snake_case  (				self     :    Dict	,					**__snake_case     :    Any		):
										kwargs.update(self.special_tokens_map		)
										return self.rust_tokenizer_class.from_pretrained(self.tmpdirname	,					**__snake_case		)







							def 				snake_case  (				self     :    Optional[Any]	,					__snake_case     :    Union[str, Any]		):
										return "lower newer", "lower newer"







							@cached_property
							def 				snake_case  (				self     :    Any		):
										return LEDTokenizer.from_pretrained('''allenai/led-base-16384'''		)







							@cached_property
							def 				snake_case  (				self     :    int		):
										return LEDTokenizerFast.from_pretrained('''allenai/led-base-16384'''		)







							@require_torch
							def 				snake_case  (				self     :    str		):
										lowerCamelCase						:Any						  =     ['''A long paragraph for summarization.''', '''Another paragraph for summarization.''']
										lowerCamelCase						:List[Any]						  =     [0, 250, 251, 17818, 13, 39186, 1938, 4, 2]

										for tokenizer in [self.default_tokenizer, self.default_tokenizer_fast]:
													lowerCamelCase						:Optional[Any]						  =     tokenizer(__snake_case	,					max_length=len(__snake_case		)	,					padding=__snake_case	,					return_tensors='''pt'''		)
													self.assertIsInstance(__snake_case	,					__snake_case		)

													self.assertEqual((2, 9)	,					batch.input_ids.shape		)
													self.assertEqual((2, 9)	,					batch.attention_mask.shape		)
													lowerCamelCase						:List[Any]						  =     batch.input_ids.tolist()[0]
													self.assertListEqual(__snake_case	,					__snake_case		)







							@require_torch
							def 				snake_case  (				self     :    Tuple		):
										lowerCamelCase						:Dict						  =     ['''A long paragraph for summarization.''', '''Another paragraph for summarization.''']
										for tokenizer in [self.default_tokenizer, self.default_tokenizer_fast]:
													lowerCamelCase						:Union[str, Any]						  =     tokenizer(__snake_case	,					padding=__snake_case	,					return_tensors='''pt'''		)
													self.assertIn('''input_ids'''	,					__snake_case		)
													self.assertIn('''attention_mask'''	,					__snake_case		)
													self.assertNotIn('''labels'''	,					__snake_case		)
													self.assertNotIn('''decoder_attention_mask'''	,					__snake_case		)







							@require_torch
							def 				snake_case  (				self     :    Union[str, Any]		):
										lowerCamelCase						:Union[str, Any]						  =     [
										    '''Summary of the text.''',
										    '''Another summary.''',
										]
										for tokenizer in [self.default_tokenizer, self.default_tokenizer_fast]:
													lowerCamelCase						:List[Any]						  =     tokenizer(text_target=__snake_case	,					max_length=32	,					padding='''max_length'''	,					return_tensors='''pt'''		)
													self.assertEqual(32	,					targets['''input_ids'''].shape[1]		)







							@require_torch
							def 				snake_case  (				self     :    List[Any]		):
										for tokenizer in [self.default_tokenizer, self.default_tokenizer_fast]:
													lowerCamelCase						:Optional[Any]						  =     tokenizer(
													    ['''I am a small frog''' * 1024, '''I am a small frog''']	,					padding=__snake_case	,					truncation=__snake_case	,					return_tensors='''pt'''		)
													self.assertIsInstance(__snake_case	,					__snake_case		)
													self.assertEqual(batch.input_ids.shape	,					(2, 5122)		)







							@require_torch
							def 				snake_case  (				self     :    Optional[int]		):
										lowerCamelCase						:Union[str, Any]						  =     ['''A long paragraph for summarization.''']
										lowerCamelCase						:Any						  =     [
										    '''Summary of the text.''',
										]
										for tokenizer in [self.default_tokenizer, self.default_tokenizer_fast]:
													lowerCamelCase						:Union[str, Any]						  =     tokenizer(__snake_case	,					return_tensors='''pt'''		)
													lowerCamelCase						:Any						  =     tokenizer(text_target=__snake_case	,					return_tensors='''pt'''		)
													lowerCamelCase						:Optional[int]						  =     inputs['''input_ids''']
													lowerCamelCase						:Any						  =     targets['''input_ids''']
													self.assertTrue((input_ids[:, 0] == tokenizer.bos_token_id).all().item()		)
													self.assertTrue((labels[:, 0] == tokenizer.bos_token_id).all().item()		)
													self.assertTrue((input_ids[:, -1] == tokenizer.eos_token_id).all().item()		)
													self.assertTrue((labels[:, -1] == tokenizer.eos_token_id).all().item()		)







							@require_torch
							def 				snake_case  (				self     :    Dict		):
										for tokenizer in [self.default_tokenizer, self.default_tokenizer_fast]:
													lowerCamelCase						:Optional[int]						  =     ['''Summary of the text.''', '''Another summary.''']
													lowerCamelCase						:List[Any]						  =     [[0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, -1, -1]]

													lowerCamelCase						:Optional[int]						  =     tokenizer(__snake_case	,					padding=__snake_case		)
													lowerCamelCase						:Union[str, Any]						  =     [[0] * len(__snake_case		) for x in encoded_output['''input_ids''']]
													lowerCamelCase						:str						  =     tokenizer.pad(__snake_case		)
													self.assertSequenceEqual(outputs['''global_attention_mask''']	,					__snake_case		)







							def 				snake_case  (				self     :    Tuple		):
										pass







							def 				snake_case  (				self     :    Optional[int]		):
										for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
													with self.subTest(F"{tokenizer.__class__.__name__} ({pretrained_name})"		):
																lowerCamelCase						:Optional[Any]						  =     self.rust_tokenizer_class.from_pretrained(__snake_case	,					**__snake_case		)
																lowerCamelCase						:Tuple						  =     self.tokenizer_class.from_pretrained(__snake_case	,					**__snake_case		)
																lowerCamelCase						:int						  =     '''A, <mask> AllenNLP sentence.'''
																lowerCamelCase						:str						  =     tokenizer_r.encode_plus(__snake_case	,					add_special_tokens=__snake_case	,					return_token_type_ids=__snake_case		)
																lowerCamelCase						:str						  =     tokenizer_p.encode_plus(__snake_case	,					add_special_tokens=__snake_case	,					return_token_type_ids=__snake_case		)
																self.assertEqual(sum(tokens_r['''token_type_ids''']		)	,					sum(tokens_p['''token_type_ids''']		)		)
																self.assertEqual(
																    sum(tokens_r['''attention_mask''']		) / len(tokens_r['''attention_mask''']		)	,					sum(tokens_p['''attention_mask''']		) / len(tokens_p['''attention_mask''']		)	,					)

																lowerCamelCase						:Optional[int]						  =     tokenizer_r.convert_ids_to_tokens(tokens_r['''input_ids''']		)
																lowerCamelCase						:Tuple						  =     tokenizer_p.convert_ids_to_tokens(tokens_p['''input_ids''']		)
																self.assertSequenceEqual(tokens_p['''input_ids''']	,					[0, 250, 6, 50264, 3823, 487, 21992, 3645, 4, 2]		)
																self.assertSequenceEqual(tokens_r['''input_ids''']	,					[0, 250, 6, 50264, 3823, 487, 21992, 3645, 4, 2]		)

																self.assertSequenceEqual(
																    __snake_case	,					['''<s>''', '''A''', ''',''', '''<mask>''', '''ĠAllen''', '''N''', '''LP''', '''Ġsentence''', '''.''', '''</s>''']		)
																self.assertSequenceEqual(
																    __snake_case	,					['''<s>''', '''A''', ''',''', '''<mask>''', '''ĠAllen''', '''N''', '''LP''', '''Ġsentence''', '''.''', '''</s>''']		)



