



'''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							:int				    =					['''\nclass''', '''\ndef''', '''\n#''', '''\n@''', '''\nprint''', '''\nif''']







class 	_lowerCAmelCase							(			__UpperCAmelCase ):


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


	def __iter__(self							):
				A_	:						int	=     []
				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()							)
				A_	:						Tuple	=     self.tokenizer(lowercase       ,      padding=lowercase       ,      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							(			__UpperCAmelCase ):


	def __init__(self       ,      lowercase       ,      lowercase       ,      lowercase							):
				A_	:						Any	=     start_length
				A_	:						Optional[Any]	=     eof_strings
				A_	:						str	=     tokenizer


	def __call__(self       ,      lowercase       ,      lowercase       ,      **lowercase							):
				A_	:						Tuple	=     self.tokenizer.batch_decode(input_ids[:, self.start_length :]							)
				A_	:						List[str]	=     []
				for decoded_generation in decoded_generations:
							done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings							)							)
				return all(lowercase							)






def 	a			( lowerCamelCase__    ):
			'''simple docstring'''





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



def 	a			( lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__=20						,      **lowerCamelCase__    ):
			'''simple docstring'''





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

									A_,    A_	:						List[Any]	=     accelerator.gather((generated_tokens, generated_tasks)    )
									A_	:						Tuple	=     generated_tokens.cpu().numpy()
									A_	:						Dict	=     generated_tasks.cpu().numpy()

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

			A_	:						Union[str, Any]	=     [[] for _ in range(lowerCamelCase__    )]
			for task, generated_tokens in gen_token_dict.items():
						for s in generated_tokens:
									A_	:						Optional[int]	=     tokenizer.decode(lowerCamelCase__						,      skip_special_tokens=lowerCamelCase__						,      clean_up_tokenization_spaces=lowerCamelCase__    )
									code_gens[task].append(remove_last_block(lowerCamelCase__    )    )
			return code_gens



def 	a			( ):
			'''simple docstring'''





			A_	:						Tuple	=     HfArgumentParser(lowerCamelCase__    )
			A_	:						Union[str, Any]	=     parser.parse_args()

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

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

			# Use dataset load to feed to accelerate
			A_	:						Optional[Any]	=     Accelerator()
			set_seed(args.seed						,      device_specific=lowerCamelCase__    )

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

			# Generation settings
			A_	:						Union[str, 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
			A_	:						List[Any]	=     load_dataset("""openai_humaneval"""    )
			A_	:						int	=     load_metric("""code_eval"""    )

			A_	:						Optional[Any]	=     args.num_tasks if args.num_tasks is not None else len(human_eval["""test"""]    )
			A_	:						Optional[int]	=     args.n_samples // args.batch_size

			A_	:						List[Any]	=     TokenizedDataset(lowerCamelCase__						,      human_eval["""test"""]						,      n_copies=lowerCamelCase__						,      n_tasks=lowerCamelCase__    )
			# do not confuse args.batch_size, which is actually the num_return_sequences
			A_	:						Union[str, Any]	=     DataLoader(lowerCamelCase__						,      batch_size=1    )

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

			A_,    A_	:						Optional[int]	=     accelerator.prepare(lowerCamelCase__						,      lowerCamelCase__    )

			A_	:						Any	=     complete_code(
			    lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      n_tasks=lowerCamelCase__						,      batch_size=args.batch_size						,      **lowerCamelCase__						,      )

			if accelerator.is_main_process:
						A_	:						int	=     []

						for task in tqdm(range(lowerCamelCase__    )    ):
									A_	:						Optional[int]	=     human_eval["""test"""][task]["""test"""]
									A_	:						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
						A_,    A_	:						List[Any]	=     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()




'''simple docstring'''




# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from .config import config_command_parser
from .config_args import default_config_file, load_config_from_file  # noqa: F401
from .default import default_command_parser
from .update import update_command_parser



def 	a			( lowerCamelCase__=None    ):
			'''simple docstring'''





			A_	:						int	=     argparse.ArgumentParser(add_help=lowerCamelCase__						,      allow_abbrev=lowerCamelCase__    )
			# The main config parser
			A_	:						int	=     config_command_parser(lowerCamelCase__    )
			# The subparser to add commands to
			A_	:						int	=     config_parser.add_subparsers(title="""subcommands"""						,      dest="""subcommand"""    )

			# Then add other parsers with the parent parser
			default_command_parser(lowerCamelCase__						,      parents=[parent_parser]    )
			update_command_parser(lowerCamelCase__						,      parents=[parent_parser]    )

			return config_parser



def 	a			( ):
			'''simple docstring'''





			A_	:						Optional[int]	=     get_config_parser()
			A_	:						List[str]	=     config_parser.parse_args()

			if not hasattr(lowerCamelCase__						,      """func"""    ):
						config_parser.print_help()
						exit(1    )

			# Run
			args.func(lowerCamelCase__    )


if __name__ == "__main__":
						main()