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


class     lowercase     (  lowercase_  ):




							def __init__(    self :       Any						,				_UpperCamelCase :       List[Any]						,				_UpperCamelCase :       Tuple						,				_UpperCamelCase :       Tuple=None						,				_UpperCamelCase :       int=1   )       ->					Tuple:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =							tokenizer
												SCREAMING_SNAKE_CASE           =							dataset
												SCREAMING_SNAKE_CASE           =							len(a__   ) if n_tasks is None else n_tasks
												SCREAMING_SNAKE_CASE           =							n_copies





							def __iter__(    self :       Tuple   )       ->					Dict:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =							[]
												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           =							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     lowercase     (  lowercase_  ):




							def __init__(    self :       Union[str, Any]						,				_UpperCamelCase :       str						,				_UpperCamelCase :       List[str]						,				_UpperCamelCase :       int   )       ->					Union[str, Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =							start_length
												SCREAMING_SNAKE_CASE           =							eof_strings
												SCREAMING_SNAKE_CASE           =							tokenizer





							def __call__(    self :       List[Any]						,				_UpperCamelCase :       Optional[int]						,				_UpperCamelCase :       Dict						,				**_UpperCamelCase :       int   )       ->					Union[str, Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE           =							self.tokenizer.batch_decode(input_ids[:, self.start_length :]   )
												SCREAMING_SNAKE_CASE           =							[]
												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 						__lowerCamelCase				(UpperCAmelCase__  :				Any	):
					SCREAMING_SNAKE_CASE           =							re.split("(%s)" % "|".join(snake_case__	)   ,			snake_case__	)
					# last string should be ""
					return "".join(string_list[:-2]	)





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

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

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

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





def 						__lowerCamelCase				():
					SCREAMING_SNAKE_CASE           =							HfArgumentParser(snake_case__	)
					SCREAMING_SNAKE_CASE           =							parser.parse_args()

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

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

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

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

					# Generation settings
					SCREAMING_SNAKE_CASE           =							{
					    """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   ,			snake_case__   ,			snake_case__	)]	),
					}

					# Load evaluation dataset and metric
					SCREAMING_SNAKE_CASE           =							load_dataset("openai_humaneval"	)
					SCREAMING_SNAKE_CASE           =							load_metric("code_eval"	)

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

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

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

					SCREAMING_SNAKE_CASE           =							complete_code(
					    snake_case__   ,			snake_case__   ,			snake_case__   ,			snake_case__   ,			n_tasks=snake_case__   ,			batch_size=args.batch_size   ,			**snake_case__   ,			)

					if accelerator.is_main_process:
										SCREAMING_SNAKE_CASE           =							[]

										for task in tqdm(range(snake_case__	)	):
															SCREAMING_SNAKE_CASE           =							human_eval["""test"""][task]["""test"""]
															SCREAMING_SNAKE_CASE           =							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           =							code_eval_metric.compute(
										    references=snake_case__   ,			predictions=snake_case__   ,			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(snake_case__   ,			snake_case__	)


# 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 importlib
import os

import fsspec
import pytest
from fsspec import register_implementation
from fsspec.registry import _registry as _fsspec_registry

from datasets.filesystems import COMPRESSION_FILESYSTEMS, HfFileSystem, extract_path_from_uri, is_remote_filesystem

from .utils import require_lza, require_zstandard






def 	UpperCamelCase	(							snake_case__  : int	):




     '''simple docstring'''






     assert "mock" in _fsspec_registry
     assert "bz2" in _fsspec_registry






def 	UpperCamelCase	(							):




     '''simple docstring'''






     assert "mock" not in _fsspec_registry
     assert "bz2" in _fsspec_registry






def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :Tuple 							=    """mock-s3-bucket"""
     __snake_case      :Optional[Any] 							=    f'''s3://{mock_bucket}'''
     __snake_case      :Optional[int] 							=    extract_path_from_uri(snake_case__	)
     assert dataset_path.startswith("""s3://"""	) is False

     __snake_case      :Union[str, Any] 							=    """./local/path"""
     __snake_case      :List[str] 							=    extract_path_from_uri(snake_case__	)
     assert dataset_path == new_dataset_path






def 	UpperCamelCase	(							snake_case__  : str	):




     '''simple docstring'''






     __snake_case      :List[str] 							=    is_remote_filesystem(snake_case__	)
     assert is_remote is True

     __snake_case      :Any 							=    fsspec.filesystem("""file"""	)

     __snake_case      :Optional[Any] 							=    is_remote_filesystem(snake_case__	)
     assert is_remote is False






@pytest.mark.parametrize("""compression_fs_class"""						,snake_case__	)
def 	UpperCamelCase	(							snake_case__  : Optional[Any]						,snake_case__  : List[str]						,snake_case__  : int						,snake_case__  : List[Any]						,snake_case__  : Tuple						,snake_case__  : str						,snake_case__  : Optional[Any]	):




     '''simple docstring'''






     __snake_case      :Optional[int] 							=    {"""gzip""": gz_file, """xz""": xz_file, """zstd""": zstd_file, """bz2""": bza_file, """lz4""": lza_file}
     __snake_case      :Optional[Any] 							=    input_paths[compression_fs_class.protocol]
     if input_path is None:
          __snake_case      :Optional[int] 							=    f'''for \'{compression_fs_class.protocol}\' compression protocol, '''
          if compression_fs_class.protocol == "lz4":
               reason += require_lza.kwargs["reason"]
          elif compression_fs_class.protocol == "zstd":
               reason += require_zstandard.kwargs["reason"]
          pytest.skip(snake_case__	)
     __snake_case      :Optional[Any] 							=    fsspec.filesystem(compression_fs_class.protocol						,fo=snake_case__	)
     assert isinstance(snake_case__						,snake_case__	)
     __snake_case      :Dict 							=    os.path.basename(snake_case__	)
     __snake_case      :Dict 							=    expected_filename[: expected_filename.rindex("""."""	)]
     assert fs.glob("""*"""	) == [expected_filename]
     with fs.open(snake_case__						,"""r"""						,encoding="""utf-8"""	) as f, open(snake_case__						,encoding="""utf-8"""	) as expected_file:
          assert f.read() == expected_file.read()






@pytest.mark.parametrize("""protocol"""						,["""zip""", """gzip"""]	)
def 	UpperCamelCase	(							snake_case__  : List[Any]						,snake_case__  : List[Any]						,snake_case__  : Dict	):




     '''simple docstring'''






     __snake_case      :Optional[int] 							=    {"""zip""": zip_jsonl_path, """gzip""": jsonl_gz_path}
     __snake_case      :List[str] 							=    compressed_file_paths[protocol]
     __snake_case      :Optional[int] 							=    """dataset.jsonl"""
     __snake_case      :Optional[int] 							=    f'''{protocol}://{member_file_path}::{compressed_file_path}'''
     __snake_case							,			*__snake_case      :Optional[Any] 							=    fsspec.get_fs_token_paths(snake_case__	)
     assert fs.isfile(snake_case__	)
     assert not fs.isfile("""non_existing_""" + member_file_path	)






@pytest.mark.integration
def 	UpperCamelCase	(							snake_case__  : Optional[int]						,snake_case__  : Dict						,snake_case__  : List[str]						,snake_case__  : Dict	):




     '''simple docstring'''






     __snake_case      :List[Any] 							=    hf_api.dataset_info(snake_case__						,token=snake_case__	)
     __snake_case      :List[str] 							=    HfFileSystem(repo_info=snake_case__						,token=snake_case__	)
     assert sorted(hffs.glob("""*"""	)	) == [".gitattributes", "data"]
     assert hffs.isdir("""data"""	)
     assert hffs.isfile(""".gitattributes"""	) and hffs.isfile("""data/text_data.txt"""	)
     with open(snake_case__	) as f:
          assert hffs.open("""data/text_data.txt"""						,"""r"""	).read() == f.read()






def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :str 							=    """bz2"""

     # Import module
     import datasets.filesystems

     # Overwrite protocol and reload
     register_implementation(snake_case__						,snake_case__						,clobber=snake_case__	)
     with pytest.warns(snake_case__	) as warning_info:
          importlib.reload(datasets.filesystems	)

     assert len(snake_case__	) == 1
     assert (
         str(warning_info[0].message	)
         == f'''A filesystem protocol was already set for {protocol} and will be overwritten.'''
     )


