






'''simple docstring'''

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






class     a_						(      __lowercase    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE=None      ,		_SCREAMING_SNAKE_CASE=1				)  ->				Optional[Any]:



        """simple docstring"""






        UpperCamelCase      = tokenizer
        UpperCamelCase      = dataset
        UpperCamelCase      = len(UpperCAmelCase__				) if n_tasks is None else n_tasks
        UpperCamelCase      = n_copies


  def __iter__(				self				)  ->				Optional[int]:



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

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)  ->				str:



        """simple docstring"""






        UpperCamelCase      = start_length
        UpperCamelCase      = eof_strings
        UpperCamelCase      = tokenizer


  def __call__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		**_SCREAMING_SNAKE_CASE				)  ->				Union[str, Any]:



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







def    lowercase__			(				__UpperCamelCase   )->			List[Any]:
      UpperCamelCase      = re.split("""(%s)""" % """|""".join(__UpperCamelCase   )		,						__UpperCamelCase   )
      # last string should be ""
      return "".join(string_list[:-2]   )







def    lowercase__			(				__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase=20		,						**__UpperCamelCase   )->			Optional[Any]:
      UpperCamelCase      = defaultdict(__UpperCamelCase   )  # dict of list of generated tokens
      for step, batch in tqdm(enumerate(__UpperCamelCase   )   ):
            with torch.no_grad():
                  UpperCamelCase      = batch["""ids"""].shape[-1]
                  UpperCamelCase      = accelerator.unwrap_model(__UpperCamelCase   ).generate(
                      input_ids=batch["""ids"""][:, : batch["""input_len"""]]		,						num_return_sequences=__UpperCamelCase		,						**__UpperCamelCase   )
                  # each task is generated batch_size times
                  UpperCamelCase      = batch["""task_id"""].repeat(__UpperCamelCase   )
                  UpperCamelCase      = accelerator.pad_across_processes(
                      __UpperCamelCase		,						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(__UpperCamelCase		,						__UpperCamelCase   ):
                        gen_token_dict[task].append(__UpperCamelCase   )

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







def    lowercase__			(				)->			List[str]:
      # Setup configuration
      UpperCamelCase      = HfArgumentParser(__UpperCamelCase   )
      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=__UpperCamelCase   )

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

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

      UpperCamelCase      = complete_code(
          __UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase		,						n_tasks=__UpperCamelCase		,						batch_size=args.batch_size		,						**__UpperCamelCase		,						)

      if accelerator.is_main_process:
            UpperCamelCase      = []

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


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








'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


SCREAMING_SNAKE_CASE__     =       {
    'configuration_nllb_moe': [
        'NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'NllbMoeConfig',
    ]
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     SCREAMING_SNAKE_CASE__     =       [
         'NLLB_MOE_PRETRAINED_MODEL_ARCHIVE_LIST',
         'NllbMoeForConditionalGeneration',
         'NllbMoeModel',
         'NllbMoePreTrainedModel',
         'NllbMoeTop2Router',
         'NllbMoeSparseMLP',
     ]


if TYPE_CHECKING:
     from .configuration_nllb_moe import (
         NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP,
         NllbMoeConfig,
     )

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_nllb_moe import (
              NLLB_MOE_PRETRAINED_MODEL_ARCHIVE_LIST,
              NllbMoeForConditionalGeneration,
              NllbMoeModel,
              NllbMoePreTrainedModel,
              NllbMoeSparseMLP,
              NllbMoeTopaRouter,
          )


else:
     import sys

     SCREAMING_SNAKE_CASE__     =       _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
