


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   __lowercase				(  __snake_case  ):





      def __init__(		self						:					Optional[int]					,       __lowerCamelCase						:					Optional[Any]					,       __lowerCamelCase						:					Union[str, Any]					,       __lowerCamelCase						:					Optional[Any]=None					,       __lowerCamelCase						:					str=1       )    ->   Union[str, Any]:






         """simple docstring"""







         UpperCAmelCase							=     tokenizer
         UpperCAmelCase							=     dataset
         UpperCAmelCase							=     len(__lowerCamelCase       ) if n_tasks is None else n_tasks
         UpperCAmelCase							=     n_copies





      def __iter__(		self						:					Dict       )    ->   Any:






         """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(__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   __lowercase				(  __snake_case  ):





      def __init__(		self						:					Tuple					,       __lowerCamelCase						:					Tuple					,       __lowerCamelCase						:					Tuple					,       __lowerCamelCase						:					str       )    ->   Optional[int]:






         """simple docstring"""







         UpperCAmelCase							=     start_length
         UpperCAmelCase							=     eof_strings
         UpperCAmelCase							=     tokenizer





      def __call__(		self						:					int					,       __lowerCamelCase						:					str					,       __lowerCamelCase						:					Dict					,       **__lowerCamelCase						:					List[str]       )    ->   Tuple:






         """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(__lowerCamelCase       )

def   _UpperCamelCase				(						lowerCAmelCase_   )						->Any:
   UpperCAmelCase							=     re.split("""(%s)""" % """|""".join(lowerCAmelCase_   )  ,						lowerCAmelCase_   )
   # last string should be ""
   return "".join(string_list[:-2]   )
def   _UpperCamelCase				(						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_=2_0  ,						**lowerCAmelCase_   )						->Tuple:
   UpperCAmelCase							=     defaultdict(lowerCAmelCase_   )  # dict of list of generated tokens
   for step, batch in tqdm(enumerate(lowerCAmelCase_   )   ):
      with torch.no_grad():
         UpperCAmelCase							=     batch["""ids"""].shape[-1]
         UpperCAmelCase							=     accelerator.unwrap_model(lowerCAmelCase_   ).generate(
             input_ids=batch["""ids"""][:, : batch["""input_len"""]]  ,						num_return_sequences=lowerCAmelCase_  ,						**lowerCAmelCase_   )
         # each task is generated batch_size times
         UpperCAmelCase							=     batch["""task_id"""].repeat(lowerCAmelCase_   )
         UpperCAmelCase							=     accelerator.pad_across_processes(
             lowerCAmelCase_  ,						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(lowerCAmelCase_  ,						lowerCAmelCase_   ):
            gen_token_dict[task].append(lowerCAmelCase_   )

   UpperCAmelCase							=     [[] for _ in range(lowerCAmelCase_   )]
   for task, generated_tokens in gen_token_dict.items():
      for s in generated_tokens:
         UpperCAmelCase							=     tokenizer.decode(lowerCAmelCase_  ,						skip_special_tokens=lowerCAmelCase_  ,						clean_up_tokenization_spaces=lowerCAmelCase_   )
         code_gens[task].append(remove_last_block(lowerCAmelCase_   )   )
   return code_gens
def   _UpperCamelCase				(						)						->List[str]:
   # Setup configuration
   UpperCAmelCase							=     HfArgumentParser(lowerCAmelCase_   )
   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=lowerCAmelCase_   )

   # 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  ,						lowerCAmelCase_  ,						lowerCAmelCase_   )]   ),
   }

   # 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(lowerCAmelCase_  ,						human_eval["""test"""]  ,						n_copies=lowerCAmelCase_  ,						n_tasks=lowerCAmelCase_   )
   # do not confuse args.batch_size, which is actually the num_return_sequences
   UpperCAmelCase							=     DataLoader(lowerCAmelCase_  ,						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(lowerCAmelCase_  ,						lowerCAmelCase_   )

   UpperCAmelCase							=     complete_code(
       lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_  ,						n_tasks=lowerCAmelCase_  ,						batch_size=args.batch_size  ,						**lowerCAmelCase_  ,						)

   if accelerator.is_main_process:
      UpperCAmelCase							=     []

      for task in tqdm(range(lowerCAmelCase_   )   ):
         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=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()







__a     =			[
    (1000, """M"""),
    (900, """CM"""),
    (500, """D"""),
    (400, """CD"""),
    (100, """C"""),
    (90, """XC"""),
    (50, """L"""),
    (40, """XL"""),
    (10, """X"""),
    (9, """IX"""),
    (5, """V"""),
    (4, """IV"""),
    (1, """I"""),
]
def   _UpperCamelCase				(						lowerCAmelCase_   )						->int:
   UpperCAmelCase							=     {"""I""": 1, """V""": 5, """X""": 1_0, """L""": 5_0, """C""": 1_0_0, """D""": 5_0_0, """M""": 1_0_0_0}
   UpperCAmelCase							=     0
   UpperCAmelCase							=     0
   while place < len(lowerCAmelCase_   ):
      if (place + 1 < len(lowerCAmelCase_   )) and (vals[roman[place]] < vals[roman[place + 1]]):
         total += vals[roman[place + 1]] - vals[roman[place]]
         place += 2
      else:
         total += vals[roman[place]]
         place += 1
   return total
def   _UpperCamelCase				(						lowerCAmelCase_   )						->str:
   UpperCAmelCase							=     []
   for arabic, roman in ROMAN:
      ((UpperCAmelCase)				,  (UpperCAmelCase))							=     divmod(lowerCAmelCase_  ,						lowerCAmelCase_   )
      result.append(roman * factor   )
      if number == 0:
         break
   return "".join(lowerCAmelCase_   )


if __name__ == "__main__":
 import doctest

 doctest.testmod()



