


"""simple docstring"""


import json
import os
import re
import unittest

from transformers import CodeGenTokenizer, CodeGenTokenizerFast
from transformers.models.codegen.tokenization_codegen import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin



@require_tokenizers
class   _a      (      SCREAMING_SNAKE_CASE_		,		unittest.TestCase ):
    a_    :				Tuple             =    CodeGenTokenizer
    a_    :				List[str]             =    CodeGenTokenizerFast
    a_    :				Optional[int]             =    True
    a_    :				Union[str, Any]             =    {'add_prefix_space': True}
    a_    :				List[str]             =    False
    def      _UpperCamelCase       ( self							:					int				):
          super().setUp()

          # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
          lowerCamelCase__           =		[
              'l',
              'o',
              'w',
              'e',
              'r',
              's',
              't',
              'i',
              'd',
              'n',
              '\u0120',
              '\u0120l',
              '\u0120n',
              '\u0120lo',
              '\u0120low',
              'er',
              '\u0120lowest',
              '\u0120newer',
              '\u0120wider',
              '<unk>',
              '<|endoftext|>',
          ]
          lowerCamelCase__           =		dict(zip(SCREAMING_SNAKE_CASE__							,					range(len(SCREAMING_SNAKE_CASE__				)				)				)				)
          lowerCamelCase__           =		['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
          lowerCamelCase__           =		{'unk_token': '<unk>'}

          lowerCamelCase__           =		os.path.join(self.tmpdirname							,					VOCAB_FILES_NAMES['vocab_file']				)
          lowerCamelCase__           =		os.path.join(self.tmpdirname							,					VOCAB_FILES_NAMES['merges_file']				)
          with open(self.vocab_file							,					'w'							,					encoding='utf-8'				) as fp:
                fp.write(json.dumps(SCREAMING_SNAKE_CASE__				) + '\n'				)
          with open(self.merges_file							,					'w'							,					encoding='utf-8'				) as fp:
                fp.write('\n'.join(SCREAMING_SNAKE_CASE__				)				)
    def      _UpperCamelCase       ( self							:					Union[str, Any]							,					**SCREAMING_SNAKE_CASE__							:					Union[str, Any]				):
          kwargs.update(self.special_tokens_map				)
          return CodeGenTokenizer.from_pretrained(self.tmpdirname							,					**SCREAMING_SNAKE_CASE__				)
    def      _UpperCamelCase       ( self							:					List[str]							,					**SCREAMING_SNAKE_CASE__							:					str				):
          kwargs.update(self.special_tokens_map				)
          return CodeGenTokenizerFast.from_pretrained(self.tmpdirname							,					**SCREAMING_SNAKE_CASE__				)
    def      _UpperCamelCase       ( self							:					int							,					SCREAMING_SNAKE_CASE__							:					List[str]				):
          lowerCamelCase__           =		'lower newer'
          lowerCamelCase__           =		'lower newer'
          return input_text, output_text
    def      _UpperCamelCase       ( self							:					Union[str, Any]				):
          lowerCamelCase__           =		CodeGenTokenizer(self.vocab_file							,					self.merges_file							,					**self.special_tokens_map				)
          lowerCamelCase__           =		'lower newer'
          lowerCamelCase__           =		['\u0120low', 'er', '\u0120', 'n', 'e', 'w', 'er']
          lowerCamelCase__           =		tokenizer.tokenize(SCREAMING_SNAKE_CASE__							,					add_prefix_space=SCREAMING_SNAKE_CASE__				)
          self.assertListEqual(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__				)

          lowerCamelCase__           =		tokens + [tokenizer.unk_token]
          lowerCamelCase__           =		[14, 15, 10, 9, 3, 2, 15, 19]
          self.assertListEqual(tokenizer.convert_tokens_to_ids(SCREAMING_SNAKE_CASE__				)							,					SCREAMING_SNAKE_CASE__				)
    def      _UpperCamelCase       ( self							:					str				):
          if not self.test_rust_tokenizer:
                return

          lowerCamelCase__           =		self.get_tokenizer()
          lowerCamelCase__           =		self.get_rust_tokenizer(add_prefix_space=SCREAMING_SNAKE_CASE__				)

          lowerCamelCase__           =		'lower newer'

          # Testing tokenization
          lowerCamelCase__           =		tokenizer.tokenize(SCREAMING_SNAKE_CASE__							,					add_prefix_space=SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		rust_tokenizer.tokenize(SCREAMING_SNAKE_CASE__				)
          self.assertListEqual(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__				)

          # Testing conversion to ids without special tokens
          lowerCamelCase__           =		tokenizer.encode(SCREAMING_SNAKE_CASE__							,					add_special_tokens=SCREAMING_SNAKE_CASE__							,					add_prefix_space=SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		rust_tokenizer.encode(SCREAMING_SNAKE_CASE__							,					add_special_tokens=SCREAMING_SNAKE_CASE__				)
          self.assertListEqual(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__				)

          # Testing conversion to ids with special tokens
          lowerCamelCase__           =		self.get_rust_tokenizer(add_prefix_space=SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		tokenizer.encode(SCREAMING_SNAKE_CASE__							,					add_prefix_space=SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		rust_tokenizer.encode(SCREAMING_SNAKE_CASE__				)
          self.assertListEqual(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__				)

          # Testing the unknown token
          lowerCamelCase__           =		tokens + [rust_tokenizer.unk_token]
          lowerCamelCase__           =		[14, 15, 10, 9, 3, 2, 15, 19]
          self.assertListEqual(rust_tokenizer.convert_tokens_to_ids(SCREAMING_SNAKE_CASE__				)							,					SCREAMING_SNAKE_CASE__				)
    def      _UpperCamelCase       ( self							:					List[Any]							,					*SCREAMING_SNAKE_CASE__							:					Dict							,					**SCREAMING_SNAKE_CASE__							:					str				):
          # It's very difficult to mix/test pretokenization with byte-level
          # And get both CodeGen and Roberta to work at the same time (mostly an issue of adding a space before the string)
          pass
    def      _UpperCamelCase       ( self							:					Dict							,					SCREAMING_SNAKE_CASE__							:					Union[str, Any]=15				):
          for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
                with self.subTest(F'{tokenizer.__class__.__name__} ({pretrained_name})'				):
                      lowerCamelCase__           =		self.rust_tokenizer_class.from_pretrained(SCREAMING_SNAKE_CASE__							,					**SCREAMING_SNAKE_CASE__				)

                      # Simple input
                      lowerCamelCase__           =		'This is a simple input'
                      lowerCamelCase__           =		['This is a simple input 1', 'This is a simple input 2']
                      lowerCamelCase__           =		('This is a simple input', 'This is a pair')
                      lowerCamelCase__           =		[
                          ('This is a simple input 1', 'This is a simple input 2'),
                          ('This is a simple pair 1', 'This is a simple pair 2'),
                      ]

                      # Simple input tests
                      self.assertRaises(SCREAMING_SNAKE_CASE__							,					tokenizer_r.encode							,					SCREAMING_SNAKE_CASE__							,					max_length=SCREAMING_SNAKE_CASE__							,					padding='max_length'				)

                      # Simple input
                      self.assertRaises(SCREAMING_SNAKE_CASE__							,					tokenizer_r.encode_plus							,					SCREAMING_SNAKE_CASE__							,					max_length=SCREAMING_SNAKE_CASE__							,					padding='max_length'				)

                      # Simple input
                      self.assertRaises(
                          SCREAMING_SNAKE_CASE__							,					tokenizer_r.batch_encode_plus							,					SCREAMING_SNAKE_CASE__							,					max_length=SCREAMING_SNAKE_CASE__							,					padding='max_length'							,					)

                      # Pair input
                      self.assertRaises(SCREAMING_SNAKE_CASE__							,					tokenizer_r.encode							,					SCREAMING_SNAKE_CASE__							,					max_length=SCREAMING_SNAKE_CASE__							,					padding='max_length'				)

                      # Pair input
                      self.assertRaises(SCREAMING_SNAKE_CASE__							,					tokenizer_r.encode_plus							,					SCREAMING_SNAKE_CASE__							,					max_length=SCREAMING_SNAKE_CASE__							,					padding='max_length'				)

                      # Pair input
                      self.assertRaises(
                          SCREAMING_SNAKE_CASE__							,					tokenizer_r.batch_encode_plus							,					SCREAMING_SNAKE_CASE__							,					max_length=SCREAMING_SNAKE_CASE__							,					padding='max_length'							,					)
    def      _UpperCamelCase       ( self							:					List[Any]				):
          lowerCamelCase__           =		CodeGenTokenizer.from_pretrained(self.tmpdirname							,					pad_token='<pad>'				)

          # Simple input
          lowerCamelCase__           =		'This is a simple input'
          lowerCamelCase__           =		['This is a simple input looooooooong', 'This is a simple input']
          lowerCamelCase__           =		('This is a simple input', 'This is a pair')
          lowerCamelCase__           =		[
              ('This is a simple input loooooong', 'This is a simple input'),
              ('This is a simple pair loooooong', 'This is a simple pair'),
          ]

          lowerCamelCase__           =		tokenizer.pad_token_id

          lowerCamelCase__           =		tokenizer(SCREAMING_SNAKE_CASE__							,					padding='max_length'							,					max_length=30							,					return_tensors='np'				)
          lowerCamelCase__           =		tokenizer(SCREAMING_SNAKE_CASE__							,					padding=SCREAMING_SNAKE_CASE__							,					truncate=SCREAMING_SNAKE_CASE__							,					return_tensors='np'				)
          lowerCamelCase__           =		tokenizer(*SCREAMING_SNAKE_CASE__							,					padding='max_length'							,					max_length=60							,					return_tensors='np'				)
          lowerCamelCase__           =		tokenizer(SCREAMING_SNAKE_CASE__							,					padding=SCREAMING_SNAKE_CASE__							,					truncate=SCREAMING_SNAKE_CASE__							,					return_tensors='np'				)

          # s
          # test single string max_length padding
          self.assertEqual(out_s['input_ids'].shape[-1]							,					30				)
          self.assertTrue(pad_token_id in out_s['input_ids']				)
          self.assertTrue(0 in out_s['attention_mask']				)

          # s2
          # test automatic padding
          self.assertEqual(out_sa['input_ids'].shape[-1]							,					33				)
          # long slice doesn't have padding
          self.assertFalse(pad_token_id in out_sa['input_ids'][0]				)
          self.assertFalse(0 in out_sa['attention_mask'][0]				)
          # short slice does have padding
          self.assertTrue(pad_token_id in out_sa['input_ids'][1]				)
          self.assertTrue(0 in out_sa['attention_mask'][1]				)

          # p
          # test single pair max_length padding
          self.assertEqual(out_p['input_ids'].shape[-1]							,					60				)
          self.assertTrue(pad_token_id in out_p['input_ids']				)
          self.assertTrue(0 in out_p['attention_mask']				)

          # p2
          # test automatic padding pair
          self.assertEqual(out_pa['input_ids'].shape[-1]							,					52				)
          # long slice pair doesn't have padding
          self.assertFalse(pad_token_id in out_pa['input_ids'][0]				)
          self.assertFalse(0 in out_pa['attention_mask'][0]				)
          # short slice pair does have padding
          self.assertTrue(pad_token_id in out_pa['input_ids'][1]				)
          self.assertTrue(0 in out_pa['attention_mask'][1]				)
    def      _UpperCamelCase       ( self							:					int				):
          lowerCamelCase__           =		'$$$'
          lowerCamelCase__           =		CodeGenTokenizer.from_pretrained(self.tmpdirname							,					bos_token=SCREAMING_SNAKE_CASE__							,					add_bos_token=SCREAMING_SNAKE_CASE__				)

          lowerCamelCase__           =		'This is a simple input'
          lowerCamelCase__           =		['This is a simple input 1', 'This is a simple input 2']

          lowerCamelCase__           =		tokenizer.bos_token_id

          lowerCamelCase__           =		tokenizer(SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		tokenizer(SCREAMING_SNAKE_CASE__				)

          self.assertEqual(out_s.input_ids[0]							,					SCREAMING_SNAKE_CASE__				)
          self.assertTrue(all(o[0] == bos_token_id for o in out_sa.input_ids				)				)

          lowerCamelCase__           =		tokenizer.decode(out_s.input_ids				)
          lowerCamelCase__           =		tokenizer.batch_decode(out_sa.input_ids				)

          self.assertEqual(decode_s.split()[0]							,					SCREAMING_SNAKE_CASE__				)
          self.assertTrue(all(d.split()[0] == bos_token for d in decode_sa				)				)
    @slow
    def      _UpperCamelCase       ( self							:					Union[str, Any]				):
          lowerCamelCase__           =		CodeGenTokenizer.from_pretrained('Salesforce/codegen-350M-mono'				)

          lowerCamelCase__           =		'\nif len_a > len_b:\n    result = a\nelse:\n    result = b\n\n\n\n#'
          lowerCamelCase__           =		'\nif len_a > len_b:      result = a\nelse:      result = b'

          lowerCamelCase__           =		tokenizer.encode(SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		['^#', re.escape('<|endoftext|>'				), '^\'\'\'', '^"""', '\n\n\n']
          lowerCamelCase__           =		tokenizer.decode(SCREAMING_SNAKE_CASE__							,					truncate_before_pattern=SCREAMING_SNAKE_CASE__				)
          self.assertEqual(SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__				)


    def      _UpperCamelCase       ( self							:					Optional[int]				):
          pass







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


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



class   _a      (      SCREAMING_SNAKE_CASE_ ):
    def __init__( self							:					List[str]							,					SCREAMING_SNAKE_CASE__							:					Any							,					SCREAMING_SNAKE_CASE__							:					List[Any]							,					SCREAMING_SNAKE_CASE__							:					int=None							,					SCREAMING_SNAKE_CASE__							:					Optional[int]=1				):
          lowerCamelCase__           =		tokenizer
          lowerCamelCase__           =		dataset
          lowerCamelCase__           =		len(SCREAMING_SNAKE_CASE__				) if n_tasks is None else n_tasks
          lowerCamelCase__           =		n_copies


    def __iter__( self							:					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      (      SCREAMING_SNAKE_CASE_ ):
    def __init__( self							:					Dict							,					SCREAMING_SNAKE_CASE__							:					List[str]							,					SCREAMING_SNAKE_CASE__							:					List[Any]							,					SCREAMING_SNAKE_CASE__							:					Dict				):
          lowerCamelCase__           =		start_length
          lowerCamelCase__           =		eof_strings
          lowerCamelCase__           =		tokenizer


    def __call__( self							:					Any							,					SCREAMING_SNAKE_CASE__							:					Any							,					SCREAMING_SNAKE_CASE__							:					List[str]							,					**SCREAMING_SNAKE_CASE__							:					Dict				):
          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   snake_case						(				_a:     List[Any]							)->		Dict:
      '''simple docstring'''


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



def   snake_case						(				_a:     List[Any]		,   _a:     Optional[int]		,   _a:     str		,   _a:     Union[str, Any]		,   _a:     Dict		,   _a:     Optional[int]=20		,   **_a:     Optional[int]							)->		List[str]:
      '''simple docstring'''


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

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



def   snake_case						(				)->		Union[str, Any]:
      '''simple docstring'''


      lowerCamelCase__           =		HfArgumentParser(_a							)
      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=_a							)

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

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

      lowerCamelCase__           =		complete_code(
          _a		,   _a		,   _a		,   _a		,   n_tasks=_a		,   batch_size=args.batch_size		,   **_a		,   )

      if accelerator.is_main_process:
            lowerCamelCase__           =		[]

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


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



