






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			:		List[str]   				=    ["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif"]
class  lowerCAmelCase_       (				snake_case__					):




							def __init__(   self     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   Optional[Any] ,							SCREAMING_SNAKE_CASE_     :   Optional[Any]=None ,							SCREAMING_SNAKE_CASE_     :   Any=1	):
											lowerCAmelCase__      			=						tokenizer
											lowerCAmelCase__      			=						dataset
											lowerCAmelCase__      			=						len(__a	) if n_tasks is None else n_tasks
											lowerCAmelCase__      			=						n_copies




							def __iter__(   self     :   int	):
											lowerCAmelCase__      			=						[]
											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__      			=						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_       (				snake_case__					):




							def __init__(   self     :   List[Any] ,							SCREAMING_SNAKE_CASE_     :   str ,							SCREAMING_SNAKE_CASE_     :   Any ,							SCREAMING_SNAKE_CASE_     :   List[Any]	):
											lowerCAmelCase__      			=						start_length
											lowerCAmelCase__      			=						eof_strings
											lowerCAmelCase__      			=						tokenizer




							def __call__(   self     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   Any ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							**SCREAMING_SNAKE_CASE_     :   str	):
											lowerCAmelCase__      			=						self.tokenizer.batch_decode(input_ids[:, self.start_length :]	)
											lowerCAmelCase__      			=						[]
											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 	lowerCAmelCase_		(lowercase__						:       Any       )		->    Dict:


				'''simple docstring'''




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





def 	lowerCAmelCase_		(lowercase__						:       Any     ,						lowercase__						:       Union[str, Any]     ,						lowercase__						:       int     ,						lowercase__						:       Union[str, Any]     ,						lowercase__						:       Tuple     ,						lowercase__						:       List[str]=20     ,						**lowercase__						:       str       )		->    Any:


				'''simple docstring'''




				lowerCAmelCase__      			=						defaultdict(__snake_case       )  # dict of list of generated tokens
				for step, batch in tqdm(enumerate(__snake_case       )       ):
								with torch.no_grad():
												lowerCAmelCase__      			=						batch['''ids'''].shape[-1]
												lowerCAmelCase__      			=						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
												lowerCAmelCase__      			=						batch['''task_id'''].repeat(__snake_case       )
												lowerCAmelCase__      			=						accelerator.pad_across_processes(
												    __snake_case     ,						dim=1     ,						pad_index=tokenizer.pad_token_id       )

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

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

				lowerCAmelCase__      			=						[[] for _ in range(__snake_case       )]
				for task, generated_tokens in gen_token_dict.items():
								for s in generated_tokens:
												lowerCAmelCase__      			=						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 	lowerCAmelCase_		()		->    Optional[Any]:


				'''simple docstring'''




				lowerCAmelCase__      			=						HfArgumentParser(__snake_case       )
				lowerCAmelCase__      			=						parser.parse_args()

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

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

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

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

				# Generation settings
				lowerCAmelCase__      			=						{
				    '''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
				lowerCAmelCase__      			=						load_dataset('''openai_humaneval'''       )
				lowerCAmelCase__      			=						load_metric('''code_eval'''       )

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

				lowerCAmelCase__      			=						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
				lowerCAmelCase__      			=						DataLoader(__snake_case     ,						batch_size=1       )

				# Run a quick test to see if code evaluation is enabled
				try:
								lowerCAmelCase__      			=						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__			,		lowerCAmelCase__      			=						accelerator.prepare(__snake_case     ,						__snake_case       )

				lowerCAmelCase__      			=						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:
								lowerCAmelCase__      			=						[]

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

								# Evaluate completions with "code_eval" metric
								lowerCAmelCase__			,		lowerCAmelCase__      			=						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()













from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_UpperCAmelCase			:		Optional[Any]   				=    {
    "configuration_mctct": ["MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP", "MCTCTConfig"],
    "feature_extraction_mctct": ["MCTCTFeatureExtractor"],
    "processing_mctct": ["MCTCTProcessor"],
}


try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_UpperCAmelCase			:		str   				=    [
			    "MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "MCTCTForCTC",
			    "MCTCTModel",
			    "MCTCTPreTrainedModel",
			]


if TYPE_CHECKING:
			from .configuration_mctct import MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP, MCTCTConfig
			from .feature_extraction_mctct import MCTCTFeatureExtractor
			from .processing_mctct import MCTCTProcessor

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_mctct import MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST, MCTCTForCTC, MCTCTModel, MCTCTPreTrainedModel

else:
			import sys

			_UpperCAmelCase			:		List[str]   				=    _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





