





from sklearn.metrics import mean_squared_error

import datasets


__lowercase       :Dict       =  "\\n@article{scikit-learn,\n title={Scikit-learn: Machine Learning in {P}ython},\n author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.\n         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.\n         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and\n         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},\n journal={Journal of Machine Learning Research},\n volume={12},\n pages={2825--2830},\n year={2011}\n}\n"

__lowercase       :str       =  "\\nMean Squared Error(MSE) is the average of the square of difference between the predicted\nand actual values.\n"


__lowercase       :Tuple       =  "\nArgs:\n    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Estimated target values.\n    references: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Ground truth (correct) target values.\n    sample_weight: array-like of shape (n_samples,), default=None\n        Sample weights.\n    multioutput: {\"raw_values\", \"uniform_average\"} or array-like of shape (n_outputs,), default=\"uniform_average\"\n        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.\n\n        \"raw_values\" : Returns a full set of errors in case of multioutput input.\n\n        \"uniform_average\" : Errors of all outputs are averaged with uniform weight.\n\n    squared : bool, default=True\n        If True returns MSE value, if False returns RMSE (Root Mean Squared Error) value.\n\nReturns:\n    mse : mean squared error.\nExamples:\n\n    >>> mse_metric = datasets.load_metric(\"mse\")\n    >>> predictions = [2.5, 0.0, 2, 8]\n    >>> references = [3, -0.5, 2, 7]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'mse': 0.375}\n    >>> rmse_result = mse_metric.compute(predictions=predictions, references=references, squared=False)\n    >>> print(rmse_result)\n    {'mse': 0.6123724356957945}\n\n    If you're using multi-dimensional lists, then set the config as follows :\n\n    >>> mse_metric = datasets.load_metric(\"mse\", \"multilist\")\n    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]\n    >>> references = [[0, 2], [-1, 2], [8, -5]]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'mse': 0.7083333333333334}\n    >>> results = mse_metric.compute(predictions=predictions, references=references, multioutput='raw_values')\n    >>> print(results) # doctest: +NORMALIZE_WHITESPACE\n    {'mse': array([0.41666667, 1. ])}\n"



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION					,							_KWARGS_DESCRIPTION							)
class    _a     (				datasets.Metric							):






							"""simple docstring"""

							def      A_  (      self				:     List[Any]    )	->Any:
													return datasets.MetricInfo(
													    description=_DESCRIPTION							,	citation=_CITATION							,	inputs_description=_KWARGS_DESCRIPTION							,	features=datasets.Features(self._get_feature_types()    )							,	reference_urls=[
													        "https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html"
													    ]							,	)

							def      A_  (      self				:     int    )	->str:
													if self.config_name == "multilist":
																			return {
																			    "predictions": datasets.Sequence(datasets.Value("float"    )    ),
																			    "references": datasets.Sequence(datasets.Value("float"    )    ),
																			}
													else:
																			return {
																			    "predictions": datasets.Value("float"    ),
																			    "references": datasets.Value("float"    ),
																			}






							def      A_  (      self				:     Union[str, Any]							,	a				:     List[str]							,	a				:     Dict							,	a				:     Union[str, Any]=None							,	a				:     List[Any]="uniform_average"							,	a				:     List[Any]=True    )	->str:
													SCREAMING_SNAKE_CASE__     :				List[str]   =					mean_squared_error(
													    a							,	a							,	sample_weight=a							,	multioutput=a							,	squared=a    )

													return {"mse": mse}






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