


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


_snake_case  :	Union[str, Any]			      =       ["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif"]

class        a	(_lowerCAmelCase   ):




							"""simple docstring"""


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






							def __iter__(							self :						List[str]       )			->					Union[str, Any]:
														__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()       )
														__snake_case	:    Optional[Any]											=     self.tokenizer(lowerCamelCase					,			padding=lowerCamelCase					,			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	(_lowerCAmelCase   ):




							"""simple docstring"""


							def __init__(							self :						str					,			lowerCamelCase :						List[str]					,			lowerCamelCase :						Union[str, Any]					,			lowerCamelCase :						Dict       )			->					Optional[Any]:
														__snake_case	:    int											=     start_length
														__snake_case	:    Optional[Any]											=     eof_strings
														__snake_case	:    str											=     tokenizer






							def __call__(							self :						Optional[Any]					,			lowerCamelCase :						str					,			lowerCamelCase :						List[str]					,			**lowerCamelCase :						Tuple       )			->					int:
														__snake_case	:    List[str]											=     self.tokenizer.batch_decode(input_ids[:, self.start_length :]       )
														__snake_case	:    Tuple											=     []
														for decoded_generation in decoded_generations:
																					done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings       )       )
														return all(lowerCamelCase       )







def        lowerCAmelCase_  (      __lowerCamelCase			):
							__snake_case	:    Union[str, Any]											=     re.split("(%s)" % "|".join(__lowerCamelCase			)							,			__lowerCamelCase			)
							# last string should be ""
							return "".join(string_list[:-2]			)







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase=2_0							,			**__lowerCamelCase			):
							__snake_case	:    int											=     defaultdict(__lowerCamelCase			)  # dict of list of generated tokens
							for step, batch in tqdm(enumerate(__lowerCamelCase			)			):
														with torch.no_grad():
																					__snake_case	:    Optional[int]											=     batch["ids"].shape[-1]
																					__snake_case	:    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
																					__snake_case	:    Optional[Any]											=     batch["task_id"].repeat(__lowerCamelCase			)
																					__snake_case	:    Any											=     accelerator.pad_across_processes(
																					    __lowerCamelCase							,			dim=1							,			pad_index=tokenizer.pad_token_id			)

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

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

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







def        lowerCAmelCase_  (      ):
							# Setup configuration
							__snake_case	:    Union[str, Any]											=     HfArgumentParser(__lowerCamelCase			)
							__snake_case	:    str											=     parser.parse_args()

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

							if args.num_workers is None:
														__snake_case	:    Union[str, Any]											=     multiprocessing.cpu_count()

							# Use dataset load to feed to accelerate
							__snake_case	:    Union[str, Any]											=     Accelerator()
							set_seed(args.seed							,			device_specific=__lowerCamelCase			)

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

							# Generation settings
							__snake_case	:    List[str]											=     {
							    "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
							__snake_case	:    Optional[Any]											=     load_dataset("openai_humaneval"			)
							__snake_case	:    Any											=     load_metric("code_eval"			)

							__snake_case	:    Dict											=     args.num_tasks if args.num_tasks is not None else len(human_eval["test"]			)
							__snake_case	:    List[str]											=     args.n_samples // args.batch_size

							__snake_case	:    Optional[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
							__snake_case	:    List[str]											=     DataLoader(__lowerCamelCase							,			batch_size=1			)

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

							__snake_case	:    Any											=     complete_code(
							    __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			n_tasks=__lowerCamelCase							,			batch_size=args.batch_size							,			**__lowerCamelCase							,			)

							if accelerator.is_main_process:
														__snake_case	:    Tuple											=     []

														for task in tqdm(range(__lowerCamelCase			)			):
																					__snake_case	:    Any											=     human_eval["test"][task]["test"]
																					__snake_case	:    str											=     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	:    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()





import unittest

from transformers import DonutProcessor


_snake_case  :	Dict			      =       "naver-clova-ix/donut-base"

class        a	(unittest.TestCase   ):




							"""simple docstring"""


							def 					__snake_case      (							self :						Dict       )			->					int:
														__snake_case	:    Any											=     DonutProcessor.from_pretrained(lowerCamelCase       )






							def 					__snake_case      (							self :						str       )			->					List[Any]:
														__snake_case	:    Union[str, Any]											=     {
														    "name": "John Doe",
														    "age": "99",
														    "city": "Atlanta",
														    "state": "GA",
														    "zip": "30301",
														    "phone": "123-4567",
														    "nicknames": [{"nickname": "Johnny"}, {"nickname": "JD"}],
														}

														__snake_case	:    Union[str, Any]											=     (
														    "<s_name>John Doe</s_name><s_age>99</s_age><s_city>Atlanta</s_city>"
														    "<s_state>GA</s_state><s_zip>30301</s_zip><s_phone>123-4567</s_phone>"
														    "<s_nicknames><s_nickname>Johnny</s_nickname>"
														    "<sep/><s_nickname>JD</s_nickname></s_nicknames>"
														)
														__snake_case	:    Tuple											=     self.processor.tokenajson(lowerCamelCase       )

														self.assertDictEqual(lowerCamelCase					,			lowerCamelCase       )

