





from typing import Dict, List, Optional, Union

import numpy as np

from .feature_extraction_utils import BatchFeature, FeatureExtractionMixin
from .utils import PaddingStrategy, TensorType, is_tf_tensor, is_torch_tensor, logging, to_numpy


_SCREAMING_SNAKE_CASE 		=						logging.get_logger(__name__)
class SCREAMING_SNAKE_CASE_  (   a_   ):






						def __init__( self : str						,				_A : Optional[Any]						,				_A : Tuple						,				_A : Optional[int]						,				**_A : str		)							->      Dict:







							"""simple docstring"""


							snake_case_						:     Dict     =		feature_size
							snake_case_						:     Dict     =		sampling_rate
							snake_case_						:     Any     =		padding_value

							snake_case_						:     Any     =		kwargs.pop('padding_side'						,				'right'		)
							snake_case_						:     List[Any]     =		kwargs.pop('return_attention_mask'						,				_A		)

							super().__init__(**_A		)






						def UpperCAmelCase_			( self : str						,				_A : Any						,				_A : str = True						,				_A : int = None						,				_A : Optional[int] = False						,				_A : Optional[Any] = None						,				_A : Union[str, Any] = None						,				_A : Optional[Any] = None						,				)							->      BatchFeature:







							"""simple docstring"""


							if isinstance(_A						,				(list, tuple)		) and isinstance(processed_features[0]						,				(dict, BatchFeature)		):
								snake_case_						:     Dict     =		{
								    key: [example[key] for example in processed_features] for key in processed_features[0].keys()
								}

							# The model's main input name, usually `input_values`, has be passed for padding
							if self.model_input_names[0] not in processed_features:
								raise ValueError(
								    'You should supply an instance of `transformers.BatchFeature` or list of `transformers.BatchFeature`'
								    F""" to this method that includes {self.model_input_names[0]}, but you provided"""
								    F""" {list(processed_features.keys()		)}"""		)

							snake_case_						:     Optional[int]     =		processed_features[self.model_input_names[0]]
							snake_case_						:     List[str]     =		(
							    return_attention_mask if return_attention_mask is not None else self.return_attention_mask
							)

							if len(_A		) == 0:
								if return_attention_mask:
									snake_case_						:     List[str]     =		[]
								return processed_features

							# If we have PyTorch/TF tensors or lists as inputs, we cast them as Numpy arrays
							# and rebuild them afterwards if no return_tensors is specified
							# Note that we lose the specific device the tensor may be on for PyTorch

							snake_case_						:     Union[str, Any]     =		required_input[0]
							if isinstance(_A						,				(list, tuple)		):
								# first_element might be an empty list/tuple in some edge cases so we grab the first non empty element.
								snake_case_						:     List[Any]     =		0
								while len(required_input[index]		) == 0:
									index += 1
								if index < len(_A		):
									snake_case_						:     Any     =		required_input[index][0]

							if return_tensors is None:
								if is_tf_tensor(_A		):
									snake_case_						:     str     =		'tf'
								elif is_torch_tensor(_A		):
									snake_case_						:     Tuple     =		'pt'
								elif isinstance(_A						,				(int, float, list, tuple, np.ndarray)		):
									snake_case_						:     List[Any]     =		'np'
								else:
									raise ValueError(
									    F"""type of {first_element} unknown: {type(_A		)}. """
									    'Should be one of a python, numpy, pytorch or tensorflow object.'		)

							for key, value in processed_features.items():
								if isinstance(value[0]						,				(int, float)		):
									snake_case_						:     Any     =		to_numpy(_A		)
								else:
									snake_case_						:     Optional[int]     =		[to_numpy(_A		) for v in value]

        # Convert padding_strategy in PaddingStrategy
							snake_case_						:     List[Any]     =		self._get_padding_strategies(padding=_A						,				max_length=_A		)

							snake_case_						:     Optional[int]     =		processed_features[self.model_input_names[0]]

							snake_case_						:     Any     =		len(_A		)
							if not all(len(_A		) == batch_size for v in processed_features.values()		):
								raise ValueError('Some items in the output dictionary have a different batch size than others.'		)

							snake_case_						:     Union[str, Any]     =		[]
							for i in range(_A		):
								snake_case_						:     int     =		{k: v[i] for k, v in processed_features.items()}
								# truncation
								snake_case_						:     Tuple     =		self._truncate(
								    _A						,				max_length=_A						,				pad_to_multiple_of=_A						,				truncation=_A						,				)
								truncated_inputs.append(_A		)

							if padding_strategy == PaddingStrategy.LONGEST:
								# make sure that `max_length` cannot be longer than the longest truncated length
								snake_case_						:     Dict     =		max(len(input_slice[self.model_input_names[0]]		) for input_slice in truncated_inputs		)
								snake_case_						:     Union[str, Any]     =		PaddingStrategy.MAX_LENGTH

							snake_case_						:     Any     =		{}
							for i in range(_A		):
								# padding
								snake_case_						:     Union[str, Any]     =		self._pad(
								    truncated_inputs[i]						,				max_length=_A						,				padding_strategy=_A						,				pad_to_multiple_of=_A						,				return_attention_mask=_A						,				)

								for key, value in outputs.items():
									if key not in batch_outputs:
										snake_case_						:     str     =		[]
									if value.dtype is np.dtype(np.floataa		):
										snake_case_						:     int     =		value.astype(np.floataa		)
									batch_outputs[key].append(_A		)

							return BatchFeature(_A						,				tensor_type=_A		)






						def UpperCAmelCase_			( self : Optional[int]						,				_A : str						,				_A : Optional[Any] = None						,				_A : Dict = PaddingStrategy.DO_NOT_PAD						,				_A : Dict = None						,				_A : List[str] = None						,				)							->      dict:







							"""simple docstring"""


							snake_case_						:     str     =		processed_features[self.model_input_names[0]]

							if padding_strategy == PaddingStrategy.LONGEST:
								snake_case_						:     List[Any]     =		len(_A		)

							if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
								snake_case_						:     Optional[int]     =		((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of

							snake_case_						:     str     =		padding_strategy != PaddingStrategy.DO_NOT_PAD and len(_A		) < max_length

							if return_attention_mask and "attention_mask" not in processed_features:
								snake_case_						:     Optional[int]     =		np.ones(len(_A		)						,				dtype=np.intaa		)

							if needs_to_be_padded:
								snake_case_						:     str     =		max_length - len(_A		)
								if self.padding_side == "right":
									if return_attention_mask:
										snake_case_						:     str     =		np.pad(
										    processed_features['attention_mask']						,				(0, difference)		)
									snake_case_						:     Dict     =		((0, difference), (0, 0)) if self.feature_size > 1 else (0, difference)
									snake_case_						:     Any     =		np.pad(
									    _A						,				_A						,				'constant'						,				constant_values=self.padding_value		)
								elif self.padding_side == "left":
									if return_attention_mask:
										snake_case_						:     Tuple     =		np.pad(
										    processed_features['attention_mask']						,				(difference, 0)		)
									snake_case_						:     Optional[int]     =		((difference, 0), (0, 0)) if self.feature_size > 1 else (difference, 0)
									snake_case_						:     List[str]     =		np.pad(
									    _A						,				_A						,				'constant'						,				constant_values=self.padding_value		)
								else:
									raise ValueError('Invalid padding strategy:' + str(self.padding_side		)		)

							return processed_features






						def UpperCAmelCase_			( self : str						,				_A : List[str]						,				_A : Dict = None						,				_A : int = None						,				_A : Optional[int] = None						,				)							->      List[str]:







							"""simple docstring"""


							if not truncation:
								return processed_features
							elif truncation and max_length is None:
								raise ValueError('When setting ``truncation=True``, make sure that ``max_length`` is defined.'		)

							snake_case_						:     Union[str, Any]     =		processed_features[self.model_input_names[0]]

							# find `max_length` that fits `pad_to_multiple_of`
							if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
								snake_case_						:     List[Any]     =		((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of

							snake_case_						:     str     =		len(_A		) > max_length

							if needs_to_be_truncated:
								snake_case_						:     Optional[int]     =		processed_features[self.model_input_names[0]][:max_length]
								if "attention_mask" in processed_features:
									snake_case_						:     int     =		processed_features['attention_mask'][:max_length]

							return processed_features






						def UpperCAmelCase_			( self : List[Any]						,				_A : Union[str, Any]=False						,				_A : Dict=None		)							->      Any:







							"""simple docstring"""


							if padding is not False:
								if padding is True:
									snake_case_						:     str     =		PaddingStrategy.LONGEST  # Default to pad to the longest sequence in the batch
								elif not isinstance(_A						,				_A		):
									snake_case_						:     List[str]     =		PaddingStrategy(_A		)
								elif isinstance(_A						,				_A		):
									snake_case_						:     Union[str, Any]     =		padding
							else:
								snake_case_						:     List[str]     =		PaddingStrategy.DO_NOT_PAD

							# Set max length if needed
							if max_length is None:
								if padding_strategy == PaddingStrategy.MAX_LENGTH:
									raise ValueError(
									    F"""When setting ``padding={PaddingStrategy.MAX_LENGTH}``, make sure that max_length is defined"""		)

        # Test if we have a padding value
							if padding_strategy != PaddingStrategy.DO_NOT_PAD and (self.padding_value is None):
								raise ValueError(
								    'Asking to pad but the feature_extractor does not have a padding value. Please select a value to use'
								    ' as `padding_value`. For example: `feature_extractor.padding_value = 0.0`.'		)

							return padding_strategy








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


_SCREAMING_SNAKE_CASE 		=						["""\nclass""", """\ndef""", """\n#""", """\n@""", """\nprint""", """\nif"""]
class SCREAMING_SNAKE_CASE_  (   snake_case_   ):






						def __init__( self : Tuple						,				_A : Any						,				_A : str						,				_A : int=None						,				_A : str=1		)							->      List[str]:







							"""simple docstring"""


							snake_case_						:     Union[str, Any]     =		tokenizer
							snake_case_						:     Optional[int]     =		dataset
							snake_case_						:     List[str]     =		len(_A		) if n_tasks is None else n_tasks
							snake_case_						:     List[Any]     =		n_copies






						def __iter__( self : Any		)							->      List[str]:







							"""simple docstring"""


							snake_case_						:     List[str]     =		[]
							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()		)
							snake_case_						:     Optional[int]     =		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 SCREAMING_SNAKE_CASE_  (   snake_case_   ):






						def __init__( self : List[Any]						,				_A : Optional[int]						,				_A : str						,				_A : Dict		)							->      Any:







							"""simple docstring"""


							snake_case_						:     List[str]     =		start_length
							snake_case_						:     int     =		eof_strings
							snake_case_						:     Dict     =		tokenizer






						def __call__( self : Any						,				_A : Union[str, Any]						,				_A : Dict						,				**_A : List[Any]		)							->      List[str]:







							"""simple docstring"""


							snake_case_						:     Optional[int]     =		self.tokenizer.batch_decode(input_ids[:, self.start_length :]		)
							snake_case_						:     Optional[int]     =		[]
							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 							SCREAMING_SNAKE_CASE__  (      __a				):
	snake_case_						:     List[str]     =		re.split('(%s)' % '|'.join(__a				)				,	__a				)
	# last string should be ""
	return "".join(string_list[:-2]				)


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

			snake_case_  ,snake_case_						:     Union[str, Any]     =		accelerator.gather((generated_tokens, generated_tasks)				)
			snake_case_						:     Optional[Any]     =		generated_tokens.cpu().numpy()
			snake_case_						:     Dict     =		generated_tasks.cpu().numpy()

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

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


def 							SCREAMING_SNAKE_CASE__  (      ):
	# Setup configuration
	snake_case_						:     Optional[int]     =		HfArgumentParser(__a				)
	snake_case_						:     Tuple     =		parser.parse_args()

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

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

	# Use dataset load to feed to accelerate
	snake_case_						:     List[Any]     =		Accelerator()
	set_seed(args.seed				,	device_specific=__a				)

	# Load model and tokenizer
	snake_case_						:     Any     =		AutoTokenizer.from_pretrained(args.model_ckpt				)
	snake_case_						:     int     =		tokenizer.eos_token
	snake_case_						:     List[str]     =		AutoModelForCausalLM.from_pretrained(args.model_ckpt				)

	# Generation settings
	snake_case_						:     List[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
	snake_case_						:     Dict     =		load_dataset('openai_humaneval'				)
	snake_case_						:     Optional[Any]     =		load_metric('code_eval'				)

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

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

	# Run a quick test to see if code evaluation is enabled
	try:
		snake_case_						:     Union[str, Any]     =		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

	snake_case_  ,snake_case_						:     Union[str, Any]     =		accelerator.prepare(__a				,	__a				)

	snake_case_						:     str     =		complete_code(
	    __a				,	__a				,	__a				,	__a				,	n_tasks=__a				,	batch_size=args.batch_size				,	**__a				,	)

	if accelerator.is_main_process:
		snake_case_						:     Tuple     =		[]

		for task in tqdm(range(__a				)				):
			snake_case_						:     Union[str, Any]     =		human_eval['test'][task]['test']
			snake_case_						:     Union[str, Any]     =		f"""check({human_eval['test'][task]['entry_point']})"""
			references.append('\n' + test_func + '\n' + entry_point				)

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

