






"""simple docstring"""





import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import MaMaaaTokenizer, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
)
from transformers.utils import is_sentencepiece_available


if is_sentencepiece_available():
  from transformers.models.mam_aaa.tokenization_mam_aaa import VOCAB_FILES_NAMES, save_json

from ...test_tokenization_common import TokenizerTesterMixin


if is_sentencepiece_available():
  _SCREAMING_SNAKE_CASE					:				Any					    =						get_tests_dir('''fixtures/test_sentencepiece.model''')


if is_torch_available():
  from transformers.models.mam_aaa.modeling_mam_aaa import shift_tokens_right

_SCREAMING_SNAKE_CASE					:				int					    =						12_8022
_SCREAMING_SNAKE_CASE					:				Optional[int]					    =						12_8028






@require_sentencepiece
class      a						(					__snake_case  ,      unittest.TestCase		):
    SCREAMING_SNAKE_CASE						:     Optional[Any]						= MaMaaaTokenizer
    SCREAMING_SNAKE_CASE						:     Union[str, Any]						= False
    SCREAMING_SNAKE_CASE						:     int						= False
    SCREAMING_SNAKE_CASE						:     Tuple						= True






    def    UpperCamelCase       (       self     :							int					)     ->		List[str]:
      super().setUp()

      lowerCamelCase_         =						['</s>', '<unk>', '▁This', '▁is', '▁a', '▁t', 'est', '\u0120', '<pad>']
      lowerCamelCase_         =						dict(zip(__SCREAMING_SNAKE_CASE       ,						range(len(__SCREAMING_SNAKE_CASE					)					)					)					)
      lowerCamelCase_         =						Path(self.tmpdirname					)
      save_json(__SCREAMING_SNAKE_CASE       ,						save_dir / VOCAB_FILES_NAMES['vocab_file']					)
      if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
        copyfile(__SCREAMING_SNAKE_CASE       ,						save_dir / VOCAB_FILES_NAMES['spm_file']					)

      lowerCamelCase_         =						MaMaaaTokenizer.from_pretrained(self.tmpdirname					)
      tokenizer.save_pretrained(self.tmpdirname					)






    def    UpperCamelCase       (       self     :							str       ,						**__SCREAMING_SNAKE_CASE     :							List[Any]					)     ->		List[str]:
      return MaMaaaTokenizer.from_pretrained(self.tmpdirname       ,						**__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							List[str]       ,						__SCREAMING_SNAKE_CASE     :							str					)     ->		Union[str, Any]:
      return (
          "This is a test",
          "This is a test",
      )






    def    UpperCamelCase       (       self     :							Dict					)     ->		Dict:
      lowerCamelCase_         =						'</s>'
      lowerCamelCase_         =						0

      self.assertEqual(self.get_tokenizer()._convert_token_to_id(__SCREAMING_SNAKE_CASE					)       ,						__SCREAMING_SNAKE_CASE					)
      self.assertEqual(self.get_tokenizer()._convert_id_to_token(__SCREAMING_SNAKE_CASE					)       ,						__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							List[Any]					)     ->		Optional[Any]:
      lowerCamelCase_         =						self.get_tokenizer()
      lowerCamelCase_         =						list(tokenizer.get_vocab().keys()					)

      self.assertEqual(vocab_keys[0]       ,						'</s>'					)
      self.assertEqual(vocab_keys[1]       ,						'<unk>'					)
      self.assertEqual(vocab_keys[-1]       ,						'<s>'					)
      self.assertEqual(len(__SCREAMING_SNAKE_CASE					)       ,						tokenizer.vocab_size + len(tokenizer.get_added_vocab()					)					)






    @unittest.skip('Skip this test while all models are still to be uploaded.'					)
    def    UpperCamelCase       (       self     :							int					)     ->		Dict:
      pass






    def    UpperCamelCase       (       self     :							List[Any]					)     ->		Any:
      lowerCamelCase_         =						self.get_tokenizer()

      lowerCamelCase_         =						tokenizer.tokenize('This is a test'					)
      self.assertListEqual(__SCREAMING_SNAKE_CASE       ,						['▁This', '▁is', '▁a', '▁t', 'est']					)

      self.assertListEqual(
          tokenizer.convert_tokens_to_ids(__SCREAMING_SNAKE_CASE					)       ,						[2, 3, 4, 5, 6]       ,						)

      lowerCamelCase_         =						tokenizer.convert_ids_to_tokens([2, 3, 4, 5, 6]					)
      self.assertListEqual(__SCREAMING_SNAKE_CASE       ,						['▁This', '▁is', '▁a', '▁t', 'est']					)

      lowerCamelCase_         =						tokenizer.convert_tokens_to_string(__SCREAMING_SNAKE_CASE					)
      self.assertEqual(__SCREAMING_SNAKE_CASE       ,						'This is a test'					)






    @slow
    def    UpperCamelCase       (       self     :							Any					)     ->		Optional[Any]:
      # fmt: off
      lowerCamelCase_         =						{'input_ids': [[128022, 110108, 397, 11, 38272, 2247, 124811, 285, 18105, 1586, 207, 7, 39534, 4428, 397, 1019, 18105, 1586, 207, 7, 41337, 16786, 241, 7, 20214, 17, 125690, 10398, 7, 44378, 58069, 68342, 7798, 7343, 11, 299, 33310, 4, 158, 37350, 94077, 4569, 299, 33310, 90, 4, 52840, 290, 4, 31270, 112, 299, 682, 4, 52840, 39953, 14079, 193, 52519, 90894, 17894, 120697, 11, 40445, 551, 17, 1019, 52519, 90894, 17756, 963, 11, 40445, 480, 17, 9792, 1120, 5173, 1393, 6240, 16786, 241, 120996, 28, 1245, 1393, 118240, 11123, 1019, 93612, 2691, 10618, 98058, 120409, 1928, 279, 4, 40683, 367, 178, 207, 1019, 103, 103121, 506, 65296, 5, 2], [128022, 21217, 367, 117, 125450, 128, 719, 7, 7308, 40, 93612, 12669, 1116, 16704, 71, 17785, 3699, 15592, 35, 144, 9584, 241, 11943, 713, 950, 799, 2247, 88427, 150, 149, 118813, 120706, 1019, 106906, 81518, 28, 1224, 22799, 397, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [128022, 1658, 123311, 5155, 5578, 4722, 279, 14947, 2366, 1120, 1197, 14, 1348, 9232, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
      # fmt: on

      self.tokenizer_integration_test_util(
          expected_encoding=__SCREAMING_SNAKE_CASE       ,						model_name='facebook/m2m100_418M'       ,						revision='c168bae485c864188cf9aa0e4108b0b6934dc91e'       ,						)






@require_torch
@require_sentencepiece
@require_tokenizers
class      a						(					unittest.TestCase		):
    SCREAMING_SNAKE_CASE						:     Dict						= """facebook/m2m100_418M"""
    SCREAMING_SNAKE_CASE						:     Optional[int]						= [
        """In my opinion, there are two levels of response from the French government.""",
        """NSA Affair Emphasizes Complete Lack of Debate on Intelligence""",
    ]
    SCREAMING_SNAKE_CASE						:     int						= [
        """Selon moi, il y a deux niveaux de réponse de la part du gouvernement français.""",
        """L'affaire NSA souligne l'absence totale de débat sur le renseignement""",
    ]

    # fmt: off
    SCREAMING_SNAKE_CASE						:     Tuple						= [EN_CODE, 5_9_3, 1_9_4_9, 1_1_5_7_8_1, 4, 7_1_5_8_6, 4_2_3_4, 6_0_6_3_3, 1_2_6_2_3_3, 4_3_2, 1_2_3_8_0_8, 1_5_5_9_2, 1_1_9_7, 1_1_7_1_3_2, 1_2_0_6_1_8, 5, 2]






    @classmethod
    def    UpperCamelCase       (       cls     :							int					)     ->		Union[str, Any]:
      lowerCamelCase_         =						MaMaaaTokenizer.from_pretrained(
          cls.checkpoint_name       ,						src_lang='en'       ,						tgt_lang='fr'					)
      lowerCamelCase_         =						1
      return cls






    def    UpperCamelCase       (       self     :							str					)     ->		Any:
      self.assertEqual(self.tokenizer.get_lang_id('ar'					)       ,						128006					)
      self.assertEqual(self.tokenizer.get_lang_id('en'					)       ,						128022					)
      self.assertEqual(self.tokenizer.get_lang_id('ro'					)       ,						128076					)
      self.assertEqual(self.tokenizer.get_lang_id('mr'					)       ,						128063					)






    def    UpperCamelCase       (       self     :							Optional[int]					)     ->		Union[str, Any]:
      lowerCamelCase_         =						self.tokenizer.get_vocab()
      self.assertEqual(len(__SCREAMING_SNAKE_CASE					)       ,						self.tokenizer.vocab_size					)
      self.assertEqual(vocab['<unk>']       ,						3					)
      self.assertIn(self.tokenizer.get_lang_token('en'					)       ,						__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							Union[str, Any]					)     ->		Any:
      lowerCamelCase_         =						'en'
      lowerCamelCase_         =						self.tokenizer.batch_encode_plus(self.src_text					).input_ids[0]
      self.assertListEqual(self.expected_src_tokens       ,						__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							Optional[Any]					)     ->		Optional[int]:
      self.assertIn(__SCREAMING_SNAKE_CASE       ,						self.tokenizer.all_special_ids					)
      # fmt: off
      lowerCamelCase_         =						[FR_CODE, 5364, 82, 8642, 4, 294, 47, 8, 14028, 136, 3286, 9706, 6, 90797, 6, 144012, 162, 88128, 30061, 5, 2]
      # fmt: on
      lowerCamelCase_         =						self.tokenizer.decode(__SCREAMING_SNAKE_CASE       ,						skip_special_tokens=__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						self.tokenizer.decode(generated_ids[1:]       ,						skip_special_tokens=__SCREAMING_SNAKE_CASE					)
      self.assertEqual(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE					)
      self.assertNotIn(self.tokenizer.eos_token       ,						__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							Optional[Any]					)     ->		Union[str, Any]:
      lowerCamelCase_         =						tempfile.mkdtemp()
      lowerCamelCase_         =						self.tokenizer.lang_token_to_id
      self.tokenizer.save_pretrained(__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						MaMaaaTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE					)
      self.assertDictEqual(new_tok.lang_token_to_id       ,						__SCREAMING_SNAKE_CASE					)






    @require_torch
    def    UpperCamelCase       (       self     :							Dict					)     ->		List[Any]:
      lowerCamelCase_         =						'en'
      lowerCamelCase_         =						'fr'

      lowerCamelCase_         =						self.tokenizer(self.src_text       ,						text_target=self.tgt_text       ,						padding=__SCREAMING_SNAKE_CASE       ,						return_tensors='pt'					)

      lowerCamelCase_         =						shift_tokens_right(
          batch['labels']       ,						self.tokenizer.pad_token_id       ,						self.tokenizer.eos_token_id					)

      for k in batch:
        lowerCamelCase_         =						batch[k].tolist()
      # batch = {k: v.tolist() for k,v in batch.items()}
      # fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
      # batch.decoder_inputs_ids[0][0] ==
      assert batch.input_ids[1][0] == EN_CODE
      assert batch.input_ids[1][-1] == 2
      assert batch.labels[1][0] == FR_CODE
      assert batch.labels[1][-1] == 2
      assert batch.decoder_input_ids[1][:2] == [2, FR_CODE]






    @require_torch
    def    UpperCamelCase       (       self     :							str					)     ->		Tuple:
      lowerCamelCase_         =						'mr'
      self.assertListEqual(self.tokenizer.prefix_tokens       ,						[self.tokenizer.get_lang_id('mr'					)]					)
      self.assertListEqual(self.tokenizer.suffix_tokens       ,						[self.tokenizer.eos_token_id]					)

      lowerCamelCase_         =						'zh'
      self.assertListEqual(self.tokenizer.prefix_tokens       ,						[self.tokenizer.get_lang_id('zh'					)]					)
      self.assertListEqual(self.tokenizer.suffix_tokens       ,						[self.tokenizer.eos_token_id]					)






    @require_torch
    def    UpperCamelCase       (       self     :							int					)     ->		List[Any]:
      lowerCamelCase_         =						'mr'
      self.tokenizer._switch_to_target_mode()
      self.assertListEqual(self.tokenizer.prefix_tokens       ,						[self.tokenizer.get_lang_id('mr'					)]					)
      self.assertListEqual(self.tokenizer.suffix_tokens       ,						[self.tokenizer.eos_token_id]					)
      self.tokenizer._switch_to_input_mode()
      self.assertListEqual(self.tokenizer.prefix_tokens       ,						[self.tokenizer.get_lang_id(self.tokenizer.src_lang					)]					)

      lowerCamelCase_         =						'zh'
      self.tokenizer._switch_to_target_mode()
      self.assertListEqual(self.tokenizer.prefix_tokens       ,						[self.tokenizer.get_lang_id('zh'					)]					)
      self.assertListEqual(self.tokenizer.suffix_tokens       ,						[self.tokenizer.eos_token_id]					)
      self.tokenizer._switch_to_input_mode()
      self.assertListEqual(self.tokenizer.prefix_tokens       ,						[self.tokenizer.get_lang_id(self.tokenizer.src_lang					)]					)






    @require_torch
    def    UpperCamelCase       (       self     :							int					)     ->		Tuple:
      lowerCamelCase_         =						self.tokenizer._build_translation_inputs('A test'       ,						return_tensors='pt'       ,						src_lang='en'       ,						tgt_lang='ar'					)

      self.assertEqual(
          nested_simplify(__SCREAMING_SNAKE_CASE					)       ,						{
              # en_XX, A, test, EOS
              'input_ids': [[128022, 58, 4183, 2]],
              'attention_mask': [[1, 1, 1, 1]],
              # ar_AR
              'forced_bos_token_id': 128006,
          }       ,						)












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




