


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


__UpperCamelCase    		=       ['\nclass', '\ndef', '\n#', '\n@', '\nprint', '\nif']







class        lowerCamelCase__	(							UpperCAmelCase      ):






			"""simple docstring"""






			def __init__(							self      ,						snake_case      ,						snake_case      ,						snake_case=None      ,						snake_case=1     ):

						'''simple docstring'''


						UpperCamelCase__					=       tokenizer
						UpperCamelCase__					=       dataset
						UpperCamelCase__					=       len(snake_case     ) if n_tasks is None else n_tasks
						UpperCamelCase__					=       n_copies







			def __iter__(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       []
						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()     )
						UpperCamelCase__					=       self.tokenizer(snake_case      ,						padding=snake_case      ,						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      ):






			"""simple docstring"""






			def __init__(							self      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       start_length
						UpperCamelCase__					=       eof_strings
						UpperCamelCase__					=       tokenizer







			def __call__(							self      ,						snake_case      ,						snake_case      ,						**snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.tokenizer.batch_decode(input_ids[:, self.start_length :]     )
						UpperCamelCase__					=       []
						for decoded_generation in decoded_generations:
									done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings     )     )
						return all(snake_case     )



def        UpperCamelCase_(				_A   :Optional[Any]  )->		Union[str, Any]:
			UpperCamelCase__					=       re.split("(%s)" % "|".join(_A  )				,     _A  )
			# last string should be ""
			return "".join(string_list[:-2]  )
def        UpperCamelCase_(				_A   :Union[str, Any]				,     _A   :Optional[Any]				,     _A   :List[Any]				,     _A   :Optional[int]				,     _A   :str				,     _A   :List[str]=20				,     **_A   :Tuple  )->		int:
			UpperCamelCase__					=       defaultdict(_A  )  # dict of list of generated tokens
			for step, batch in tqdm(enumerate(_A  )  ):
						with torch.no_grad():
									UpperCamelCase__					=       batch["ids"].shape[-1]
									UpperCamelCase__					=       accelerator.unwrap_model(_A  ).generate(
									    input_ids=batch["ids"][:, : batch["input_len"]]				,     num_return_sequences=_A				,     **_A  )
									# each task is generated batch_size times
									UpperCamelCase__					=       batch["task_id"].repeat(_A  )
									UpperCamelCase__					=       accelerator.pad_across_processes(
									    _A				,     dim=1				,     pad_index=tokenizer.pad_token_id  )

									UpperCamelCase__,      UpperCamelCase__					=       accelerator.gather((generated_tokens, generated_tasks)  )
									UpperCamelCase__					=       generated_tokens.cpu().numpy()
									UpperCamelCase__					=       generated_tasks.cpu().numpy()

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

			UpperCamelCase__					=       [[] for _ in range(_A  )]
			for task, generated_tokens in gen_token_dict.items():
						for s in generated_tokens:
									UpperCamelCase__					=       tokenizer.decode(_A				,     skip_special_tokens=_A				,     clean_up_tokenization_spaces=_A  )
									code_gens[task].append(remove_last_block(_A  )  )
			return code_gens
def        UpperCamelCase_(				)->		Any:
			# Setup configuration
			UpperCamelCase__					=       HfArgumentParser(_A  )
			UpperCamelCase__					=       parser.parse_args()

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

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

			# Use dataset load to feed to accelerate
			UpperCamelCase__					=       Accelerator()
			set_seed(args.seed				,     device_specific=_A  )

			# Load model and tokenizer
			UpperCamelCase__					=       AutoTokenizer.from_pretrained(args.model_ckpt  )
			UpperCamelCase__					=       tokenizer.eos_token
			UpperCamelCase__					=       AutoModelForCausalLM.from_pretrained(args.model_ckpt  )

			# Generation settings
			UpperCamelCase__					=       {
			    "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
			UpperCamelCase__					=       load_dataset("openai_humaneval"  )
			UpperCamelCase__					=       load_metric("code_eval"  )

			UpperCamelCase__					=       args.num_tasks if args.num_tasks is not None else len(human_eval["test"]  )
			UpperCamelCase__					=       args.n_samples // args.batch_size

			UpperCamelCase__					=       TokenizedDataset(_A				,     human_eval["test"]				,     n_copies=_A				,     n_tasks=_A  )
			# do not confuse args.batch_size, which is actually the num_return_sequences
			UpperCamelCase__					=       DataLoader(_A				,     batch_size=1  )

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

			UpperCamelCase__,      UpperCamelCase__					=       accelerator.prepare(_A				,     _A  )

			UpperCamelCase__					=       complete_code(
			    _A				,     _A				,     _A				,     _A				,     n_tasks=_A				,     batch_size=args.batch_size				,     **_A				,     )

			if accelerator.is_main_process:
						UpperCamelCase__					=       []

						for task in tqdm(range(_A  )  ):
									UpperCamelCase__					=       human_eval["test"][task]["test"]
									UpperCamelCase__					=       F'''check({human_eval['test'][task]['entry_point']})'''
									references.append("\n" + test_func + "\n" + entry_point  )

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






import os
import textwrap

import pyarrow as pa
import pytest

from datasets import ClassLabel, Features, Image
from datasets.packaged_modules.csv.csv import Csv

from ..utils import require_pil
@pytest.fixture
def        UpperCamelCase_(				_A   :Union[str, Any]  )->		int:
			UpperCamelCase__					=       tmp_path / "file.csv"
			UpperCamelCase__					=       textwrap.dedent(
			    "\\n        header1,header2\n        1,2\n        10,20\n        "  )
			with open(_A				,     "w"  ) as f:
						f.write(_A  )
			return str(_A  )
@pytest.fixture
def        UpperCamelCase_(				_A   :int  )->		Union[str, Any]:
			UpperCamelCase__					=       tmp_path / "malformed_file.csv"
			UpperCamelCase__					=       textwrap.dedent(
			    "\\n        header1,header2\n        1,2\n        10,20,\n        "  )
			with open(_A				,     "w"  ) as f:
						f.write(_A  )
			return str(_A  )
@pytest.fixture
def        UpperCamelCase_(				_A   :List[Any]				,     _A   :Dict  )->		Optional[Any]:
			UpperCamelCase__					=       tmp_path / "csv_with_image.csv"
			UpperCamelCase__					=       textwrap.dedent(
			    F'''\
        image
        {image_file}
        '''  )
			with open(_A				,     "w"  ) as f:
						f.write(_A  )
			return str(_A  )
@pytest.fixture
def        UpperCamelCase_(				_A   :List[Any]  )->		Any:
			UpperCamelCase__					=       tmp_path / "csv_with_label.csv"
			UpperCamelCase__					=       textwrap.dedent(
			    "\\n        label\n        good\n        bad\n        good\n        "  )
			with open(_A				,     "w"  ) as f:
						f.write(_A  )
			return str(_A  )
@pytest.fixture
def        UpperCamelCase_(				_A   :Any  )->		int:
			UpperCamelCase__					=       tmp_path / "csv_with_int_list.csv"
			UpperCamelCase__					=       textwrap.dedent(
			    "\\n        int_list\n        1 2 3\n        4 5 6\n        7 8 9\n        "  )
			with open(_A				,     "w"  ) as f:
						f.write(_A  )
			return str(_A  )
def        UpperCamelCase_(				_A   :Optional[Any]				,     _A   :Optional[Any]				,     _A   :List[str]  )->		Union[str, Any]:
			UpperCamelCase__					=       Csv()
			UpperCamelCase__					=       csv._generate_tables([[csv_file, malformed_csv_file]]  )
			with pytest.raises(_A				,     match="Error tokenizing data"  ):
						for _ in generator:
									pass
			assert any(
			    record.levelname == "ERROR"
			    and "Failed to read file" in record.message
			    and os.path.basename(_A  ) in record.message
			    for record in caplog.records  )
@require_pil
def        UpperCamelCase_(				_A   :Any  )->		int:
			with open(_A				,     encoding="utf-8"  ) as f:
						UpperCamelCase__					=       f.read().splitlines()[1]
			UpperCamelCase__					=       Csv(encoding="utf-8"				,     features=Features({"image": Image()}  )  )
			UpperCamelCase__					=       csv._generate_tables([[csv_file_with_image]]  )
			UpperCamelCase__					=       pa.concat_tables([table for _, table in generator]  )
			assert pa_table.schema.field("image"  ).type == Image()()
			UpperCamelCase__					=       pa_table.to_pydict()["image"]
			assert generated_content == [{"path": image_file, "bytes": None}]
def        UpperCamelCase_(				_A   :Any  )->		Dict:
			with open(_A				,     encoding="utf-8"  ) as f:
						UpperCamelCase__					=       f.read().splitlines()[1:]
			UpperCamelCase__					=       Csv(encoding="utf-8"				,     features=Features({"label": ClassLabel(names=["good", "bad"]  )}  )  )
			UpperCamelCase__					=       csv._generate_tables([[csv_file_with_label]]  )
			UpperCamelCase__					=       pa.concat_tables([table for _, table in generator]  )
			assert pa_table.schema.field("label"  ).type == ClassLabel(names=["good", "bad"]  )()
			UpperCamelCase__					=       pa_table.to_pydict()["label"]
			assert generated_content == [ClassLabel(names=["good", "bad"]  ).straint(_A  ) for label in labels]







def        UpperCamelCase_(				_A   :Dict  )->		Optional[Any]:
			UpperCamelCase__					=       Csv(encoding="utf-8"				,     sep=","				,     converters={"int_list": lambda _A  : [int(_A  ) for i in x.split()]}  )
			UpperCamelCase__					=       csv._generate_tables([[csv_file_with_int_list]]  )
			UpperCamelCase__					=       pa.concat_tables([table for _, table in generator]  )
			assert pa.types.is_list(pa_table.schema.field("int_list"  ).type  )
			UpperCamelCase__					=       pa_table.to_pydict()["int_list"]
			assert generated_content == [[1, 2, 3], [4, 5, 6], [7, 8, 9]]


