






"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


_SCREAMING_SNAKE_CASE					:				int					    =						{'''configuration_reformer''': ['''REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ReformerConfig''']}

try:
  if not is_sentencepiece_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _SCREAMING_SNAKE_CASE					:				List[str]					    =						['''ReformerTokenizer''']

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _SCREAMING_SNAKE_CASE					:				Union[str, Any]					    =						['''ReformerTokenizerFast''']

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _SCREAMING_SNAKE_CASE					:				List[Any]					    =						[
      '''REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''ReformerAttention''',
      '''ReformerForMaskedLM''',
      '''ReformerForQuestionAnswering''',
      '''ReformerForSequenceClassification''',
      '''ReformerLayer''',
      '''ReformerModel''',
      '''ReformerModelWithLMHead''',
      '''ReformerPreTrainedModel''',
  ]


if TYPE_CHECKING:
  from .configuration_reformer import REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, ReformerConfig

  try:
    if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_reformer import ReformerTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_reformer_fast import ReformerTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_reformer import (
        REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        ReformerAttention,
        ReformerForMaskedLM,
        ReformerForQuestionAnswering,
        ReformerForSequenceClassification,
        ReformerLayer,
        ReformerModel,
        ReformerModelWithLMHead,
        ReformerPreTrainedModel,
    )

else:
  import sys

  _SCREAMING_SNAKE_CASE					:				Dict					    =						_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)












"""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


_SCREAMING_SNAKE_CASE					:				Union[str, Any]					    =						['''\nclass''', '''\ndef''', '''\n#''', '''\n@''', '''\nprint''', '''\nif''']






class      a						(					__snake_case		):






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






    def __iter__(       self     :							Dict					)     ->		Any:
      lowerCamelCase_         =						[]
      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()					)
      lowerCamelCase_         =						self.tokenizer(__SCREAMING_SNAKE_CASE       ,						padding=__SCREAMING_SNAKE_CASE       ,						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						(					__snake_case		):






    def __init__(       self     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							str       ,						__SCREAMING_SNAKE_CASE     :							int       ,						__SCREAMING_SNAKE_CASE     :							Any					)     ->		List[Any]:
      lowerCamelCase_         =						start_length
      lowerCamelCase_         =						eof_strings
      lowerCamelCase_         =						tokenizer






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







def      lowerCamelCase__       (							_lowerCamelCase       :		List[Any] )      ->				Tuple:
  lowerCamelCase_         =						re.split('(%s)' % '|'.join(_lowerCamelCase )		,				_lowerCamelCase )
  # last string should be ""
  return "".join(string_list[:-2] )







def      lowerCamelCase__       (							_lowerCamelCase       :		str		,				_lowerCamelCase       :		str		,				_lowerCamelCase       :		Optional[Any]		,				_lowerCamelCase       :		Optional[Any]		,				_lowerCamelCase       :		str		,				_lowerCamelCase       :		Dict=20		,				**_lowerCamelCase       :		Dict )      ->				List[str]:
  lowerCamelCase_         =						defaultdict(_lowerCamelCase )  # dict of list of generated tokens
  for step, batch in tqdm(enumerate(_lowerCamelCase ) ):
    with torch.no_grad():
      lowerCamelCase_         =						batch['ids'].shape[-1]
      lowerCamelCase_         =						accelerator.unwrap_model(_lowerCamelCase ).generate(
          input_ids=batch['ids'][:, : batch['input_len']]		,				num_return_sequences=_lowerCamelCase		,				**_lowerCamelCase )
      # each task is generated batch_size times
      lowerCamelCase_         =						batch['task_id'].repeat(_lowerCamelCase )
      lowerCamelCase_         =						accelerator.pad_across_processes(
          _lowerCamelCase		,				dim=1		,				pad_index=tokenizer.pad_token_id )

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

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

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







def      lowerCamelCase__       (							)      ->				Tuple:
  # Setup configuration
  lowerCamelCase_         =						HfArgumentParser(_lowerCamelCase )
  lowerCamelCase_         =						parser.parse_args()

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

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

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

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

  # Generation settings
  lowerCamelCase_         =						{
      '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
  lowerCamelCase_         =						load_dataset('openai_humaneval' )
  lowerCamelCase_         =						load_metric('code_eval' )

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

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

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

  lowerCamelCase_     ,	lowerCamelCase_         =						accelerator.prepare(_lowerCamelCase		,				_lowerCamelCase )

  lowerCamelCase_         =						complete_code(
      _lowerCamelCase		,				_lowerCamelCase		,				_lowerCamelCase		,				_lowerCamelCase		,				n_tasks=_lowerCamelCase		,				batch_size=args.batch_size		,				**_lowerCamelCase		,				)

  if accelerator.is_main_process:
    lowerCamelCase_         =						[]

    for task in tqdm(range(_lowerCamelCase ) ):
      lowerCamelCase_         =						human_eval['test'][task]['test']
      lowerCamelCase_         =						F'''check({human_eval["test"][task]["entry_point"]})'''
      references.append('\n' + test_func + '\n' + entry_point )

    # Evaluate completions with "code_eval" metric
    lowerCamelCase_     ,	lowerCamelCase_         =						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()




