





"""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      __SCREAMING_SNAKE_CASE    (						_UpperCamelCase   ):







  '''simple docstring'''


  def __init__(		self   :       Optional[Any]					,   __a   :       str					,   __a   :       Union[str, Any]					,   __a   :       Optional[Any]=None					,   __a   :       Any=1	)		->				Tuple:
        _UpperCamelCase      :	Union[str, Any]						=       tokenizer
        _UpperCamelCase      :	Tuple						=       dataset
        _UpperCamelCase      :	str						=       len(__a	) if n_tasks is None else n_tasks
        _UpperCamelCase      :	Union[str, Any]						=       n_copies


  def __iter__(		self   :       str	)		->				Optional[int]:
        _UpperCamelCase      :	Optional[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      :	Dict						=       self.tokenizer(__a					,   padding=__a					,   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      __SCREAMING_SNAKE_CASE    (						_UpperCamelCase   ):







  '''simple docstring'''


  def __init__(		self   :       Tuple					,   __a   :       Any					,   __a   :       int					,   __a   :       Optional[int]	)		->				int:
        _UpperCamelCase      :	Union[str, Any]						=       start_length
        _UpperCamelCase      :	Union[str, Any]						=       eof_strings
        _UpperCamelCase      :	Tuple						=       tokenizer


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





def 					lowercase__    (					lowercase_			)					->		int:




      """simple docstring"""

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




def 					lowercase__    (					lowercase_		,lowercase_		,lowercase_		,lowercase_		,lowercase_		,lowercase_=20		,**lowercase_			)					->		Dict:




      """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      :	Optional[int]						=       batch["ids"].shape[-1]
                  _UpperCamelCase      :	int						=       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      :	Optional[Any]						=       batch["task_id"].repeat(lowercase_			)
                  _UpperCamelCase      :	Any						=       accelerator.pad_across_processes(
                      lowercase_		,dim=1		,pad_index=tokenizer.pad_token_id			)

                  _UpperCamelCase,  _UpperCamelCase      :	Dict						=       accelerator.gather((generated_tokens, generated_tasks)			)
                  _UpperCamelCase      :	Any						=       generated_tokens.cpu().numpy()
                  _UpperCamelCase      :	Optional[int]						=       generated_tasks.cpu().numpy()

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

      _UpperCamelCase      :	Tuple						=       [[] for _ in range(lowercase_			)]
      for task, generated_tokens in gen_token_dict.items():
            for s in generated_tokens:
                  _UpperCamelCase      :	Optional[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 					lowercase__    (					)					->		List[str]:




      """simple docstring"""

      _UpperCamelCase      :	Union[str, Any]						=       HfArgumentParser(lowercase_			)
      _UpperCamelCase      :	str						=       parser.parse_args()

      transformers.logging.set_verbosity_error()
      # enables code execution in code_eval metric
      _UpperCamelCase      :	str						=       args.HF_ALLOW_CODE_EVAL
      # make sure tokenizer plays nice with multiprocessing
      _UpperCamelCase      :	int						=       "false"

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

      # Use dataset load to feed to accelerate
      _UpperCamelCase      :	Optional[Any]						=       Accelerator()
      set_seed(args.seed		,device_specific=lowercase_			)

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

      # Generation settings
      _UpperCamelCase      :	Optional[int]						=       {
          "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      :	List[Any]						=       load_dataset("openai_humaneval"			)
      _UpperCamelCase      :	int						=       load_metric("code_eval"			)

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

      _UpperCamelCase      :	Tuple						=       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      :	List[Any]						=       DataLoader(lowercase_		,batch_size=1			)

      # Run a quick test to see if code evaluation is enabled
      try:
            _UpperCamelCase      :	Union[str, 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      :	Tuple						=       accelerator.prepare(lowercase_		,lowercase_			)

      _UpperCamelCase      :	Union[str, Any]						=       complete_code(
          lowercase_		,lowercase_		,lowercase_		,lowercase_		,n_tasks=lowercase_		,batch_size=args.batch_size		,**lowercase_		,)

      if accelerator.is_main_process:
            _UpperCamelCase      :	Union[str, Any]						=       []

            for task in tqdm(range(lowercase_			)			):
                  _UpperCamelCase      :	List[Any]						=       human_eval["test"][task]["test"]
                  _UpperCamelCase      :	List[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      :	List[str]						=       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()







"""simple docstring"""


from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Value
from .base import TaskTemplate



@dataclass(frozen=_UpperCamelCase   )
class      __SCREAMING_SNAKE_CASE    (						_UpperCamelCase   ):







  '''simple docstring'''


  SCREAMING_SNAKE_CASE__							:str	     =					field(default="summarization"		,  metadata={"include_in_asdict_even_if_is_default": True}   )
  SCREAMING_SNAKE_CASE__							:ClassVar[Features]	     =					Features({"text": Value("string"   )}   )
  SCREAMING_SNAKE_CASE__							:ClassVar[Features]	     =					Features({"summary": Value("string"   )}   )
  SCREAMING_SNAKE_CASE__							:str	     =					"text"
  SCREAMING_SNAKE_CASE__							:str	     =					"summary"

  @property
  def    __SCREAMING_SNAKE_CASE				(		self   :       List[Any]	)		->				Dict[str, str]:
        return {self.text_column: "text", self.summary_column: "summary"}
