




import os

import pytest
import yaml

from datasets.features.features import Features, Value
from datasets.info import DatasetInfo, DatasetInfosDict







@pytest.mark.parametrize(
    '''files'''					,		[
        ['''full:README.md''', '''dataset_infos.json'''],
        ['''empty:README.md''', '''dataset_infos.json'''],
        ['''dataset_infos.json'''],
        ['''full:README.md'''],
    ]					,		)
def lowerCamelCase__  (     _lowercase					,		_lowercase ):

     '''simple docstring'''


     UpperCAmelCase_     :				int  =     tmp_path_factory.mktemp('''dset_infos_dir''' )
     if "full:README.md" in files:
          with open(dataset_infos_dir / '''README.md'''					,		'''w''' ) as f:
               f.write('''---\ndataset_info:\n  dataset_size: 42\n---''' )
     if "empty:README.md" in files:
          with open(dataset_infos_dir / '''README.md'''					,		'''w''' ) as f:
               f.write('''''' )
    # we want to support dataset_infos.json for backward compatibility
     if "dataset_infos.json" in files:
          with open(dataset_infos_dir / '''dataset_infos.json'''					,		'''w''' ) as f:
               f.write('''{"default": {"dataset_size": 42}}''' )
     UpperCAmelCase_     :				List[Any]  =     DatasetInfosDict.from_directory(_lowercase )
     assert dataset_infos
     assert dataset_infos["default"].dataset_size == 42







@pytest.mark.parametrize(
    '''dataset_info'''					,		[
        DatasetInfo(),
        DatasetInfo(
            description='''foo'''					,		features=Features({'''a''': Value('''int32''' )} )					,		builder_name='''builder'''					,		config_name='''config'''					,		version='''1.0.0'''					,		splits=[{'''name''': '''train'''}]					,		download_size=42					,		),
    ]					,		)
def lowerCamelCase__  (     _lowercase					,		_lowercase ):

     '''simple docstring'''


     UpperCAmelCase_     :				List[str]  =     str(_lowercase )
     dataset_info.write_to_directory(_lowercase )
     UpperCAmelCase_     :				int  =     DatasetInfo.from_directory(_lowercase )
     assert dataset_info == reloaded
     assert os.path.exists(os.path.join(_lowercase					,		'''dataset_info.json''' ) )







def lowerCamelCase__  (     ):

     '''simple docstring'''


     UpperCAmelCase_     :				Tuple  =     DatasetInfo(
         description='''foo'''					,		citation='''bar'''					,		homepage='''https://foo.bar'''					,		license='''CC0'''					,		features=Features({'''a''': Value('''int32''' )} )					,		post_processed={}					,		supervised_keys=()					,		task_templates=[]					,		builder_name='''builder'''					,		config_name='''config'''					,		version='''1.0.0'''					,		splits=[{'''name''': '''train''', '''num_examples''': 42}]					,		download_checksums={}					,		download_size=1337					,		post_processing_size=442					,		dataset_size=1234					,		size_in_bytes=1337 + 442 + 1234					,		)
     UpperCAmelCase_     :				List[str]  =     dataset_info._to_yaml_dict()
     assert sorted(_lowercase ) == sorted(DatasetInfo._INCLUDED_INFO_IN_YAML )
     for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
          assert key in dataset_info_yaml_dict
          assert isinstance(dataset_info_yaml_dict[key]					,		(list, dict, int, str) )
     UpperCAmelCase_     :				str  =     yaml.safe_dump(_lowercase )
     UpperCAmelCase_     :				List[Any]  =     yaml.safe_load(_lowercase )
     assert dataset_info_yaml_dict == reloaded







def lowerCamelCase__  (     ):

     '''simple docstring'''


     UpperCAmelCase_     :				Optional[int]  =     DatasetInfo()
     UpperCAmelCase_     :				Union[str, Any]  =     dataset_info._to_yaml_dict()
     assert dataset_info_yaml_dict == {}







@pytest.mark.parametrize(
    '''dataset_infos_dict'''					,		[
        DatasetInfosDict(),
        DatasetInfosDict({'''default''': DatasetInfo()} ),
        DatasetInfosDict({'''my_config_name''': DatasetInfo()} ),
        DatasetInfosDict(
            {
                '''default''': DatasetInfo(
                    description='''foo'''					,		features=Features({'''a''': Value('''int32''' )} )					,		builder_name='''builder'''					,		config_name='''config'''					,		version='''1.0.0'''					,		splits=[{'''name''': '''train'''}]					,		download_size=42					,		)
            } ),
        DatasetInfosDict(
            {
                '''v1''': DatasetInfo(dataset_size=42 ),
                '''v2''': DatasetInfo(dataset_size=1337 ),
            } ),
    ]					,		)
def lowerCamelCase__  (     _lowercase					,		_lowercase ):

     '''simple docstring'''


     UpperCAmelCase_     :				Optional[int]  =     str(_lowercase )
     dataset_infos_dict.write_to_directory(_lowercase )
     UpperCAmelCase_     :				Dict  =     DatasetInfosDict.from_directory(_lowercase )

     # the config_name of the dataset_infos_dict take over the attribute
     for config_name, dataset_info in dataset_infos_dict.items():
          UpperCAmelCase_     :				List[str]  =     config_name
          # the yaml representation doesn't include fields like description or citation
          # so we just test that we can recover what we can from the yaml
          UpperCAmelCase_     :				Optional[Any]  =     DatasetInfo._from_yaml_dict(dataset_info._to_yaml_dict() )
     assert dataset_infos_dict == reloaded

     if dataset_infos_dict:
          assert os.path.exists(os.path.join(_lowercase					,		'''README.md''' ) )





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


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

class     __a(	_a  ):






    """simple docstring"""





    def __init__( self  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE=None  ,_SCREAMING_SNAKE_CASE=1    )					->						Dict:
         UpperCAmelCase_     :				List[Any]  =     tokenizer
         UpperCAmelCase_     :				int  =     dataset
         UpperCAmelCase_     :				Dict  =     len(_SCREAMING_SNAKE_CASE    ) if n_tasks is None else n_tasks
         UpperCAmelCase_     :				Optional[int]  =     n_copies





    def __iter__( self    )					->						Any:
         UpperCAmelCase_     :				List[Any]  =     []
         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_     :				Union[str, Any]  =     self.tokenizer(_SCREAMING_SNAKE_CASE  ,padding=_SCREAMING_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     __a(	_a  ):






    """simple docstring"""





    def __init__( self  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )					->						List[Any]:
         UpperCAmelCase_     :				str  =     start_length
         UpperCAmelCase_     :				Optional[int]  =     eof_strings
         UpperCAmelCase_     :				str  =     tokenizer





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







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


     UpperCAmelCase_     :				Tuple  =     re.split('''(%s)''' % '''|'''.join(_lowercase )					,		_lowercase )
     # last string should be ""
     return "".join(string_list[:-2] )







def lowerCamelCase__  (     _lowercase					,		_lowercase					,		_lowercase					,		_lowercase					,		_lowercase					,		_lowercase=20					,		**_lowercase ):

     '''simple docstring'''


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

               UpperCAmelCase_,       UpperCAmelCase_     :				List[str]  =     accelerator.gather((generated_tokens, generated_tasks) )
               UpperCAmelCase_     :				Union[str, Any]  =     generated_tokens.cpu().numpy()
               UpperCAmelCase_     :				Union[str, Any]  =     generated_tasks.cpu().numpy()

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

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







def lowerCamelCase__  (     ):

     '''simple docstring'''


     UpperCAmelCase_     :				Optional[Any]  =     HfArgumentParser(_lowercase )
     UpperCAmelCase_     :				int  =     parser.parse_args()

     transformers.logging.set_verbosity_error()
     # enables code execution in code_eval metric
     UpperCAmelCase_     :				Optional[Any]  =     args.HF_ALLOW_CODE_EVAL
     # make sure tokenizer plays nice with multiprocessing
     UpperCAmelCase_     :				List[Any]  =     '''false'''

     if args.num_workers is None:
          UpperCAmelCase_     :				Optional[Any]  =     multiprocessing.cpu_count()

     # Use dataset load to feed to accelerate
     UpperCAmelCase_     :				int  =     Accelerator()
     set_seed(args.seed					,		device_specific=_lowercase )

     # Load model and tokenizer
     UpperCAmelCase_     :				Tuple  =     AutoTokenizer.from_pretrained(args.model_ckpt )
     UpperCAmelCase_     :				Any  =     tokenizer.eos_token
     UpperCAmelCase_     :				List[str]  =     AutoModelForCausalLM.from_pretrained(args.model_ckpt )

     # Generation settings
     UpperCAmelCase_     :				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					,		_lowercase					,		_lowercase )] ),
     }

     # Load evaluation dataset and metric
     UpperCAmelCase_     :				Tuple  =     load_dataset('''openai_humaneval''' )
     UpperCAmelCase_     :				Dict  =     load_metric('''code_eval''' )

     UpperCAmelCase_     :				Optional[int]  =     args.num_tasks if args.num_tasks is not None else len(human_eval['''test'''] )
     UpperCAmelCase_     :				str  =     args.n_samples // args.batch_size

     UpperCAmelCase_     :				str  =     TokenizedDataset(_lowercase					,		human_eval['''test''']					,		n_copies=_lowercase					,		n_tasks=_lowercase )
     # do not confuse args.batch_size, which is actually the num_return_sequences
     UpperCAmelCase_     :				Optional[Any]  =     DataLoader(_lowercase					,		batch_size=1 )

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

     UpperCAmelCase_,       UpperCAmelCase_     :				int  =     accelerator.prepare(_lowercase					,		_lowercase )

     UpperCAmelCase_     :				int  =     complete_code(
         _lowercase					,		_lowercase					,		_lowercase					,		_lowercase					,		n_tasks=_lowercase					,		batch_size=args.batch_size					,		**_lowercase					,		)

     if accelerator.is_main_process:
          UpperCAmelCase_     :				Any  =     []

          for task in tqdm(range(_lowercase ) ):
               UpperCAmelCase_     :				int  =     human_eval['''test'''][task]['''test''']
               UpperCAmelCase_     :				str  =     f'''check({human_eval["test"][task]["entry_point"]})'''
               references.append('''\n''' + test_func + '''\n''' + entry_point )

          # Evaluate completions with "code_eval" metric
          UpperCAmelCase_,       UpperCAmelCase_     :				Any  =     code_eval_metric.compute(
              references=_lowercase					,		predictions=_lowercase					,		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(_lowercase					,		_lowercase )


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