





from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


__lowercase       :int       =  logging.get_logger(__name__)



class    _a     (				lowercase__							):






							"""simple docstring"""

							snake_case_							 =						["pixel_values"]

							def __init__(      self				:     int							,	a				:     bool = True							,	a				:     Optional[Dict[str, int]] = None							,	a				:     PILImageResampling = PILImageResampling.BILINEAR							,	a				:     bool = True							,	a				:     Dict[str, int] = None							,	a				:     bool = True							,	a				:     Union[int, float] = 1 / 2_55							,	a				:     bool = True							,	a				:     Optional[Union[float, List[float]]] = None							,	a				:     Optional[Union[float, List[float]]] = None							,	**a				:     List[str]							,	)	->None:
													super().__init__(**a    )
													SCREAMING_SNAKE_CASE__     :				List[str]   =					size if size is not None else {"shortest_edge": 2_56}
													SCREAMING_SNAKE_CASE__     :				Any   =					get_size_dict(a							,	default_to_square=a    )
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					crop_size if crop_size is not None else {"height": 2_24, "width": 2_24}
													SCREAMING_SNAKE_CASE__     :				Dict   =					get_size_dict(a    )
													SCREAMING_SNAKE_CASE__     :				List[str]   =					do_resize
													SCREAMING_SNAKE_CASE__     :				List[str]   =					size
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					resample
													SCREAMING_SNAKE_CASE__     :				int   =					do_center_crop
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					crop_size
													SCREAMING_SNAKE_CASE__     :				Any   =					do_rescale
													SCREAMING_SNAKE_CASE__     :				Any   =					rescale_factor
													SCREAMING_SNAKE_CASE__     :				int   =					do_normalize
													SCREAMING_SNAKE_CASE__     :				str   =					image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
													SCREAMING_SNAKE_CASE__     :				int   =					image_std if image_std is not None else IMAGENET_STANDARD_STD

							def      A_  (      self				:     Tuple							,	a				:     np.ndarray							,	a				:     Dict[str, int]							,	a				:     PILImageResampling = PILImageResampling.BICUBIC							,	a				:     Optional[Union[str, ChannelDimension]] = None							,	**a				:     Optional[int]							,	)	->np.ndarray:
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					get_size_dict(a							,	default_to_square=a    )
													if "shortest_edge" not in size:
																			raise ValueError(f"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""    )
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					get_resize_output_image_size(a							,	size=size["shortest_edge"]							,	default_to_square=a    )
													return resize(a							,	size=a							,	resample=a							,	data_format=a							,	**a    )

							def      A_  (      self				:     List[Any]							,	a				:     np.ndarray							,	a				:     Dict[str, int]							,	a				:     Optional[Union[str, ChannelDimension]] = None							,	**a				:     List[Any]							,	)	->np.ndarray:
													SCREAMING_SNAKE_CASE__     :				Tuple   =					get_size_dict(a    )
													return center_crop(a							,	size=(size["height"], size["width"])							,	data_format=a							,	**a    )

							def      A_  (      self				:     Optional[int]							,	a				:     np.ndarray							,	a				:     float							,	a				:     Optional[Union[str, ChannelDimension]] = None							,	**a				:     Dict    )	->np.ndarray:
													return rescale(a							,	scale=a							,	data_format=a							,	**a    )

							def      A_  (      self				:     Union[str, Any]							,	a				:     np.ndarray							,	a				:     Union[float, List[float]]							,	a				:     Union[float, List[float]]							,	a				:     Optional[Union[str, ChannelDimension]] = None							,	**a				:     Union[str, Any]							,	)	->np.ndarray:
													return normalize(a							,	mean=a							,	std=a							,	data_format=a							,	**a    )






							def      A_  (      self				:     Tuple							,	a				:     ImageInput							,	a				:     Optional[bool] = None							,	a				:     Dict[str, int] = None							,	a				:     PILImageResampling = None							,	a				:     bool = None							,	a				:     Dict[str, int] = None							,	a				:     Optional[bool] = None							,	a				:     Optional[float] = None							,	a				:     Optional[bool] = None							,	a				:     Optional[Union[float, List[float]]] = None							,	a				:     Optional[Union[float, List[float]]] = None							,	a				:     Optional[Union[str, TensorType]] = None							,	a				:     Union[str, ChannelDimension] = ChannelDimension.FIRST							,	**a				:     Any							,	)	->Optional[int]:
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					do_resize if do_resize is not None else self.do_resize
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					size if size is not None else self.size
													SCREAMING_SNAKE_CASE__     :				Dict   =					get_size_dict(a							,	default_to_square=a    )
													SCREAMING_SNAKE_CASE__     :				str   =					resample if resample is not None else self.resample
													SCREAMING_SNAKE_CASE__     :				List[str]   =					do_center_crop if do_center_crop is not None else self.do_center_crop
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					crop_size if crop_size is not None else self.crop_size
													SCREAMING_SNAKE_CASE__     :				Dict   =					get_size_dict(a    )
													SCREAMING_SNAKE_CASE__     :				List[str]   =					do_rescale if do_rescale is not None else self.do_rescale
													SCREAMING_SNAKE_CASE__     :				int   =					rescale_factor if rescale_factor is not None else self.rescale_factor
													SCREAMING_SNAKE_CASE__     :				Dict   =					do_normalize if do_normalize is not None else self.do_normalize
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					image_mean if image_mean is not None else self.image_mean
													SCREAMING_SNAKE_CASE__     :				Tuple   =					image_std if image_std is not None else self.image_std

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					make_list_of_images(a    )

													if not valid_images(a    ):
																			raise ValueError(
																			    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
																			    "torch.Tensor, tf.Tensor or jax.ndarray."    )

													if do_resize and size is None:
																			raise ValueError("Size must be specified if do_resize is True."    )

													if do_center_crop and crop_size is None:
																			raise ValueError("Crop size must be specified if do_center_crop is True."    )

													if do_rescale and rescale_factor is None:
																			raise ValueError("Rescale factor must be specified if do_rescale is True."    )

													if do_normalize and (image_mean is None or image_std is None):
																			raise ValueError("Image mean and std must be specified if do_normalize is True."    )

													# All transformations expect numpy arrays.
													SCREAMING_SNAKE_CASE__     :				List[str]   =					[to_numpy_array(a    ) for image in images]

													if do_resize:
																			SCREAMING_SNAKE_CASE__     :				Tuple   =					[self.resize(image=a							,	size=a							,	resample=a    ) for image in images]

													if do_center_crop:
																			SCREAMING_SNAKE_CASE__     :				List[Any]   =					[self.center_crop(image=a							,	size=a    ) for image in images]

													if do_rescale:
																			SCREAMING_SNAKE_CASE__     :				List[str]   =					[self.rescale(image=a							,	scale=a    ) for image in images]

													if do_normalize:
																			SCREAMING_SNAKE_CASE__     :				Dict   =					[self.normalize(image=a							,	mean=a							,	std=a    ) for image in images]

													SCREAMING_SNAKE_CASE__     :				Dict   =					[to_channel_dimension_format(a							,	a    ) for image in images]

													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					{"pixel_values": images}
													return BatchFeature(data=a							,	tensor_type=a    )






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


__lowercase       :str       =  ["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif"]



class    _a     (				lowercase__							):






							"""simple docstring"""

							def __init__(      self				:     List[str]							,	a				:     Optional[int]							,	a				:     str							,	a				:     int=None							,	a				:     Optional[Any]=1    )	->Optional[Any]:
													SCREAMING_SNAKE_CASE__     :				Dict   =					tokenizer
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					dataset
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					len(a    ) if n_tasks is None else n_tasks
													SCREAMING_SNAKE_CASE__     :				Dict   =					n_copies






							def __iter__(      self				:     str    )	->Tuple:
													SCREAMING_SNAKE_CASE__     :				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()    )
													SCREAMING_SNAKE_CASE__     :				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    _a     (				lowercase__							):






							"""simple docstring"""

							def __init__(      self				:     Dict							,	a				:     int							,	a				:     int							,	a				:     Tuple    )	->Dict:
													SCREAMING_SNAKE_CASE__     :				Dict   =					start_length
													SCREAMING_SNAKE_CASE__     :				Any   =					eof_strings
													SCREAMING_SNAKE_CASE__     :				Any   =					tokenizer






							def __call__(      self				:     Any							,	a				:     Optional[int]							,	a				:     int							,	**a				:     Union[str, Any]    )	->List[str]:
													SCREAMING_SNAKE_CASE__     :				Dict   =					self.tokenizer.batch_decode(input_ids[:, self.start_length :]    )
													SCREAMING_SNAKE_CASE__     :				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       UpperCAmelCase							(     _lowerCamelCase      :  Dict     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				List[Any]   =					re.split("(%s)" % "|".join(_lowerCamelCase     )			,   _lowerCamelCase     )
						# last string should be ""
						return "".join(string_list[:-2]     )







def       UpperCAmelCase							(     _lowerCamelCase      :  Dict			,   _lowerCamelCase      :  Optional[Any]			,   _lowerCamelCase      :  Tuple			,   _lowerCamelCase      :  str			,   _lowerCamelCase      :  int			,   _lowerCamelCase      :  str=20			,   **_lowerCamelCase      :  Dict     ):





						'''simple docstring'''


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

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

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

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







def       UpperCAmelCase							(     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					HfArgumentParser(_lowerCamelCase     )
						SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					parser.parse_args()

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

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

						# Use dataset load to feed to accelerate
						SCREAMING_SNAKE_CASE__     :				Dict   =					Accelerator()
						set_seed(args.seed			,   device_specific=_lowerCamelCase     )

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

						# Generation settings
						SCREAMING_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			,   _lowerCamelCase			,   _lowerCamelCase     )]     ),
						}

						# Load evaluation dataset and metric
						SCREAMING_SNAKE_CASE__     :				str   =					load_dataset("openai_humaneval"     )
						SCREAMING_SNAKE_CASE__     :				Any   =					load_metric("code_eval"     )

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

						SCREAMING_SNAKE_CASE__     :				Dict   =					TokenizedDataset(_lowerCamelCase			,   human_eval["test"]			,   n_copies=_lowerCamelCase			,   n_tasks=_lowerCamelCase     )
						# do not confuse args.batch_size, which is actually the num_return_sequences
						SCREAMING_SNAKE_CASE__     :				Optional[int]   =					DataLoader(_lowerCamelCase			,   batch_size=1     )

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

						SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Optional[int]   =					accelerator.prepare(_lowerCamelCase			,   _lowerCamelCase     )

						SCREAMING_SNAKE_CASE__     :				Tuple   =					complete_code(
						    _lowerCamelCase			,   _lowerCamelCase			,   _lowerCamelCase			,   _lowerCamelCase			,   n_tasks=_lowerCamelCase			,   batch_size=args.batch_size			,   **_lowerCamelCase			,   )

						if accelerator.is_main_process:
												SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					[]

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

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


# 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()