





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


SCREAMING_SNAKE_CASE_						       =	["""\nclass""", """\ndef""", """\n#""", """\n@""", """\nprint""", """\nif"""]

class 		UpperCamelCase__   (				lowerCAmelCase_       ):



     '''simple docstring'''

     def __init__(							self						:   Tuple						,lowerCamelCase__						:   Any						,lowerCamelCase__						:   Tuple						,lowerCamelCase__						:   Any=None						,lowerCamelCase__						:   str=1	)					->		Tuple:




           '''simple docstring'''


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






     def __iter__(							self						:   Any	)					->		List[Any]:




           '''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(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 		UpperCamelCase__   (				lowerCAmelCase_       ):



     '''simple docstring'''

     def __init__(							self						:   Dict						,lowerCamelCase__						:   Tuple						,lowerCamelCase__						:   List[str]						,lowerCamelCase__						:   Optional[int]	)					->		str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						start_length
           SCREAMING_SNAKE_CASE		=						eof_strings
           SCREAMING_SNAKE_CASE		=						tokenizer






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




           '''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(lowerCamelCase__	)



def        __lowercase       (				_SCREAMING_SNAKE_CASE     )			->       str:


      '''simple docstring'''


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



def        __lowercase       (				_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE=20			,						**_SCREAMING_SNAKE_CASE     )			->       int:


      '''simple docstring'''


      SCREAMING_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():
                  SCREAMING_SNAKE_CASE		=						batch["""ids"""].shape[-1]
                  SCREAMING_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
                  SCREAMING_SNAKE_CASE		=						batch["""task_id"""].repeat(_SCREAMING_SNAKE_CASE     )
                  SCREAMING_SNAKE_CASE		=						accelerator.pad_across_processes(
                      _SCREAMING_SNAKE_CASE			,						dim=1			,						pad_index=tokenizer.pad_token_id     )

                  SCREAMING_SNAKE_CASE, 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(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE     ):
                        gen_token_dict[task].append(_SCREAMING_SNAKE_CASE     )

      SCREAMING_SNAKE_CASE		=						[[] for _ in range(_SCREAMING_SNAKE_CASE     )]
      for task, generated_tokens in gen_token_dict.items():
            for s in generated_tokens:
                  SCREAMING_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        __lowercase       (				)			->       Union[str, Any]:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE		=						HfArgumentParser(_SCREAMING_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=_SCREAMING_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			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_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(_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
      SCREAMING_SNAKE_CASE		=						DataLoader(_SCREAMING_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, SCREAMING_SNAKE_CASE		=						accelerator.prepare(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE     )

      SCREAMING_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:
            SCREAMING_SNAKE_CASE		=						[]

            for task in tqdm(range(_SCREAMING_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, SCREAMING_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()









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


SCREAMING_SNAKE_CASE_						       =	{
    """configuration_mvp""": ["""MVP_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MvpConfig""", """MvpOnnxConfig"""],
    """tokenization_mvp""": ["""MvpTokenizer"""],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_						       =	["""MvpTokenizerFast"""]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_						       =	[
          """MVP_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """MvpForCausalLM""",
          """MvpForConditionalGeneration""",
          """MvpForQuestionAnswering""",
          """MvpForSequenceClassification""",
          """MvpModel""",
          """MvpPreTrainedModel""",
      ]

if TYPE_CHECKING:
      from .configuration_mvp import MVP_PRETRAINED_CONFIG_ARCHIVE_MAP, MvpConfig, MvpOnnxConfig
      from .tokenization_mvp import MvpTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_mvp_fast import MvpTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_mvp import (
                MVP_PRETRAINED_MODEL_ARCHIVE_LIST,
                MvpForCausalLM,
                MvpForConditionalGeneration,
                MvpForQuestionAnswering,
                MvpForSequenceClassification,
                MvpModel,
                MvpPreTrainedModel,
            )

else:
      import sys

      SCREAMING_SNAKE_CASE_						       =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


