


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_		(				lowerCamelCase__ ):

 '''simple docstring'''

 def __init__(			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase=None	,   _UpperCAmelCase=1      ):
       snake_case_ =   tokenizer
       snake_case_ =   dataset
       snake_case_ =   len(_UpperCAmelCase      ) if n_tasks is None else n_tasks
       snake_case_ =   n_copies





 def __iter__(			self      ):
       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()      )
       snake_case_ =   self.tokenizer(_UpperCAmelCase	,   padding=_UpperCAmelCase	,   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_		(				lowerCamelCase__ ):

 '''simple docstring'''

 def __init__(			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       snake_case_ =   start_length
       snake_case_ =   eof_strings
       snake_case_ =   tokenizer





 def __call__(			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   **_UpperCAmelCase      ):
       snake_case_ =   self.tokenizer.batch_decode(input_ids[:, self.start_length :]      )
       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(_UpperCAmelCase      )






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      )->						List[Any]:




      """simple docstring"""

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






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE=20      ,      **SCREAMING_SNAKE_CASE      )->						List[Any]:




      """simple docstring"""

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

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

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

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






def  __lowerCAmelCase					()->						str:




      """simple docstring"""

      snake_case_ =   HfArgumentParser(SCREAMING_SNAKE_CASE      )
      snake_case_ =   parser.parse_args()

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

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

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

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

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

      # Load evaluation dataset and metric
      snake_case_ =   load_dataset('''openai_humaneval'''      )
      snake_case_ =   load_metric('''code_eval'''      )

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

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

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

      snake_case_ ,					snake_case_ =   accelerator.prepare(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )

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

      if accelerator.is_main_process:
            snake_case_ =   []

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






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      )->						None:




      """simple docstring"""

      snake_case_ =   {'''Content-Type''': '''application/json'''}
      snake_case_ =   requests.post(SCREAMING_SNAKE_CASE      ,      json={'''text''': message_body}      ,      headers=SCREAMING_SNAKE_CASE      )
      if response.status_code != 200:
            snake_case_ =   (
                '''Request to slack returned an error '''
                f'''{response.status_code}, the response is:\n{response.text}'''
            )
            raise ValueError(SCREAMING_SNAKE_CASE      )


if __name__ == "__main__":
   # Set the slack url to the one provided by Slack when you create the webhook at
   # https://my.slack.com/services/new/incoming-webhook/
   send_slack_message("""<YOUR MESSAGE BODY>""", """<SLACK CHANNEL URL>""")