'''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


__lowerCAmelCase	:       Optional[Any]						       =	["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif"]






class 						A   (   UpperCAmelCase   ):

							'''simple docstring'''





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






							def __iter__(      self				:							Union[str, Any]	)							->	Dict:
								__UpperCAmelCase							=			[]
								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							=			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 						A   (   UpperCAmelCase   ):

							'''simple docstring'''





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






							def __call__(      self				:							Optional[Any]     ,		__a				:							List[Any]     ,		__a				:							Union[str, Any]     ,		**__a				:							List[Any]	)							->	Tuple:
								__UpperCAmelCase							=			self.tokenizer.batch_decode(input_ids[:, self.start_length :]	)
								__UpperCAmelCase							=			[]
								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					(		UpperCamelCase__				:   List[str]       ):


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





def       lowerCAmelCase					(		UpperCamelCase__				:   Any    ,					UpperCamelCase__				:   Tuple    ,					UpperCamelCase__				:   int    ,					UpperCamelCase__				:   List[Any]    ,					UpperCamelCase__				:   int    ,					UpperCamelCase__				:   int=2_0    ,					**UpperCamelCase__				:   Optional[int]       ):


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

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

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

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





def       lowerCAmelCase					(		):


	"""simple docstring"""
	__UpperCAmelCase							=			HfArgumentParser(UpperCamelCase__       )
	__UpperCAmelCase							=			parser.parse_args()

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

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

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

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

	# Generation settings
	__UpperCAmelCase							=			{
	    '''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    ,					UpperCamelCase__    ,					UpperCamelCase__       )]       ),
	}

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

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

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

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

	__UpperCAmelCase							=			complete_code(
	    UpperCamelCase__    ,					UpperCamelCase__    ,					UpperCamelCase__    ,					UpperCamelCase__    ,					n_tasks=UpperCamelCase__    ,					batch_size=args.batch_size    ,					**UpperCamelCase__    ,					)

	if accelerator.is_main_process:
		__UpperCAmelCase							=			[]

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

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


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


from typing import TYPE_CHECKING

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


__lowerCAmelCase	:       List[Any]						       =	{
    "configuration_audio_spectrogram_transformer": [
        "AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "ASTConfig",
    ]
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__lowerCAmelCase	:       str						       =	[
					    "AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "ASTForAudioClassification",
					    "ASTModel",
					    "ASTPreTrainedModel",
					]

try:
					if not is_speech_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__lowerCAmelCase	:       Union[str, Any]						       =	["ASTFeatureExtractor"]

if TYPE_CHECKING:
					from .configuration_audio_spectrogram_transformer import (
					    AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
					    ASTConfig,
					)

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_audio_spectrogram_transformer import (
										    AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    ASTForAudioClassification,
										    ASTModel,
										    ASTPreTrainedModel,
										)

					try:
										if not is_speech_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .feature_extraction_audio_spectrogram_transformer import ASTFeatureExtractor


else:
					import sys

					__lowerCAmelCase	:       Optional[Any]						       =	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
