


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


lowercase_	   =     ['\nclass', '\ndef', '\n#', '\n@', '\nprint', '\nif']


class 	A_ (	__UpperCamelCase   ):




      '''simple docstring'''


      def __init__(    self:	Optional[int]						,				a:	List[Any]						,				a:	List[Any]						,				a:	Tuple=None						,				a:	List[str]=1    ):
       __lowerCamelCase		:							Optional[int]		 =						tokenizer
       __lowerCamelCase		:							Optional[Any]		 =						dataset
       __lowerCamelCase		:							Dict		 =						len(a    ) if n_tasks is None else n_tasks
       __lowerCamelCase		:							List[Any]		 =						n_copies




      def __iter__(    self:	str    ):
       __lowerCamelCase		:							List[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()    )
       __lowerCamelCase		:							str		 =						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 	A_ (	__UpperCamelCase   ):




      '''simple docstring'''


      def __init__(    self:	Union[str, Any]						,				a:	Optional[int]						,				a:	Optional[int]						,				a:	int    ):
       __lowerCamelCase		:							Union[str, Any]		 =						start_length
       __lowerCamelCase		:							int		 =						eof_strings
       __lowerCamelCase		:							Tuple		 =						tokenizer




      def __call__(    self:	Dict						,				a:	List[Any]						,				a:	Tuple						,				**a:	List[str]    ):
       __lowerCamelCase		:							int		 =						self.tokenizer.batch_decode(input_ids[:, self.start_length :]    )
       __lowerCamelCase		:							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 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__						):
 __lowerCamelCase		:							Optional[Any]		 =						re.split('(%s)' % '|'.join(SCREAMING_SNAKE_CASE__						)      ,	SCREAMING_SNAKE_CASE__						)
 # last string should be ""
 return "".join(string_list[:-2]						)
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__=20      ,	**SCREAMING_SNAKE_CASE__						):
 __lowerCamelCase		:							Any		 =						defaultdict(SCREAMING_SNAKE_CASE__						)  # dict of list of generated tokens
 for step, batch in tqdm(enumerate(SCREAMING_SNAKE_CASE__						)						):
  with torch.no_grad():
   __lowerCamelCase		:							str		 =						batch['ids'].shape[-1]
   __lowerCamelCase		:							Optional[Any]		 =						accelerator.unwrap_model(SCREAMING_SNAKE_CASE__						).generate(
       input_ids=batch['ids'][:, : batch['input_len']]      ,	num_return_sequences=SCREAMING_SNAKE_CASE__      ,	**SCREAMING_SNAKE_CASE__						)
   # each task is generated batch_size times
   __lowerCamelCase		:							int		 =						batch['task_id'].repeat(SCREAMING_SNAKE_CASE__						)
   __lowerCamelCase		:							Optional[Any]		 =						accelerator.pad_across_processes(
       SCREAMING_SNAKE_CASE__      ,	dim=1      ,	pad_index=tokenizer.pad_token_id						)

   __lowerCamelCase ,					__lowerCamelCase		:							Optional[Any]		 =						accelerator.gather((generated_tokens, generated_tasks)						)
   __lowerCamelCase		:							int		 =						generated_tokens.cpu().numpy()
   __lowerCamelCase		:							List[str]		 =						generated_tasks.cpu().numpy()

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

 __lowerCamelCase		:							str		 =						[[] for _ in range(SCREAMING_SNAKE_CASE__						)]
 for task, generated_tokens in gen_token_dict.items():
  for s in generated_tokens:
   __lowerCamelCase		:							List[str]		 =						tokenizer.decode(SCREAMING_SNAKE_CASE__      ,	skip_special_tokens=SCREAMING_SNAKE_CASE__      ,	clean_up_tokenization_spaces=SCREAMING_SNAKE_CASE__						)
   code_gens[task].append(remove_last_block(SCREAMING_SNAKE_CASE__						)						)
 return code_gens
def 				UpperCamelCase__					(  ):
 # Setup configuration
 __lowerCamelCase		:							List[str]		 =						HfArgumentParser(SCREAMING_SNAKE_CASE__						)
 __lowerCamelCase		:							str		 =						parser.parse_args()

 transformers.logging.set_verbosity_error()
 # enables code execution in code_eval metric
 __lowerCamelCase		:							List[Any]		 =						args.HF_ALLOW_CODE_EVAL
 # make sure tokenizer plays nice with multiprocessing
 __lowerCamelCase		:							Dict		 =						'false'

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

 # Use dataset load to feed to accelerate
 __lowerCamelCase		:							Union[str, Any]		 =						Accelerator()
 set_seed(args.seed      ,	device_specific=SCREAMING_SNAKE_CASE__						)

 # Load model and tokenizer
 __lowerCamelCase		:							int		 =						AutoTokenizer.from_pretrained(args.model_ckpt						)
 __lowerCamelCase		:							int		 =						tokenizer.eos_token
 __lowerCamelCase		:							Optional[int]		 =						AutoModelForCausalLM.from_pretrained(args.model_ckpt						)

 # Generation settings
 __lowerCamelCase		:							Union[str, Any]		 =						{
     '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      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						)]						),
 }

 # Load evaluation dataset and metric
 __lowerCamelCase		:							int		 =						load_dataset('openai_humaneval'						)
 __lowerCamelCase		:							List[Any]		 =						load_metric('code_eval'						)

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

 __lowerCamelCase		:							Union[str, Any]		 =						TokenizedDataset(SCREAMING_SNAKE_CASE__      ,	human_eval['test']      ,	n_copies=SCREAMING_SNAKE_CASE__      ,	n_tasks=SCREAMING_SNAKE_CASE__						)
 # do not confuse args.batch_size, which is actually the num_return_sequences
 __lowerCamelCase		:							Tuple		 =						DataLoader(SCREAMING_SNAKE_CASE__      ,	batch_size=1						)

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

 __lowerCamelCase		:							Tuple		 =						complete_code(
     SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	n_tasks=SCREAMING_SNAKE_CASE__      ,	batch_size=args.batch_size      ,	**SCREAMING_SNAKE_CASE__      ,	)

 if accelerator.is_main_process:
  __lowerCamelCase		:							Dict		 =						[]

  for task in tqdm(range(SCREAMING_SNAKE_CASE__						)						):
   __lowerCamelCase		:							List[Any]		 =						human_eval['test'][task]['test']
   __lowerCamelCase		:							Any		 =						f'check({human_eval["test"][task]["entry_point"]})'
   references.append('\n' + test_func + '\n' + entry_point						)

  # Evaluate completions with "code_eval" metric
  __lowerCamelCase ,					__lowerCamelCase		:							int		 =						code_eval_metric.compute(
      references=SCREAMING_SNAKE_CASE__      ,	predictions=SCREAMING_SNAKE_CASE__      ,	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(SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						)


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





import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


lowercase_	   =     5_0_0_0_0
lowercase_	   =     5_0_0_0

lowercase_   ,lowercase_	   =     os.path.split(__file__)
lowercase_	   =     os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))
@get_duration
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 for i in range(SCREAMING_SNAKE_CASE__						):
  __lowerCamelCase		:							Tuple		 =						dataset[i]
@get_duration
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 for i in range(0      ,	len(SCREAMING_SNAKE_CASE__						)      ,	SCREAMING_SNAKE_CASE__						):
  __lowerCamelCase		:							Optional[Any]		 =						dataset[i : i + batch_size]
@get_duration
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 with dataset.formatted_as(type=SCREAMING_SNAKE_CASE__						):
  for i in range(SCREAMING_SNAKE_CASE__						):
   __lowerCamelCase		:							str		 =						dataset[i]
@get_duration
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 with dataset.formatted_as(type=SCREAMING_SNAKE_CASE__						):
  for i in range(0      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
   __lowerCamelCase		:							int		 =						dataset[i : i + batch_size]
def 				UpperCamelCase__					(  ):
 __lowerCamelCase		:							Union[str, Any]		 =						{'num examples': SPEED_TEST_N_EXAMPLES}
 __lowerCamelCase		:							Optional[Any]		 =						[
     (read, {'length': SMALL_TEST}),
     (read, {'length': SPEED_TEST_N_EXAMPLES}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 10}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 100}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 1_000}),
     (read_formatted, {'type': 'numpy', 'length': SMALL_TEST}),
     (read_formatted, {'type': 'pandas', 'length': SMALL_TEST}),
     (read_formatted, {'type': 'torch', 'length': SMALL_TEST}),
     (read_formatted, {'type': 'tensorflow', 'length': SMALL_TEST}),
     (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 10}),
     (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 1_000}),
 ]

 __lowerCamelCase		:							Any		 =						[
     (read, {'length': SMALL_TEST}),
     (read, {'length': SPEED_TEST_N_EXAMPLES}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 10}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 100}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 1_000}),
     (read_formatted, {'type': 'numpy', 'length': SMALL_TEST}),
     (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 10}),
     (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 1_000}),
 ]
 with tempfile.TemporaryDirectory() as tmp_dir:
  print('generating dataset'						)
  __lowerCamelCase		:							Optional[int]		 =						datasets.Features(
      {'list': datasets.Sequence(datasets.Value('float32'						)						), 'numbers': datasets.Value('float32'						)}						)
  __lowerCamelCase		:							str		 =						generate_example_dataset(
      os.path.join(SCREAMING_SNAKE_CASE__      ,	'dataset.arrow'						)      ,	SCREAMING_SNAKE_CASE__      ,	num_examples=SCREAMING_SNAKE_CASE__      ,	seq_shapes={'list': (100,)}      ,	)
  print('first set of iterations'						)
  for func, kwargs in functions:
   print(func.__name__      ,	str(SCREAMING_SNAKE_CASE__						)						)
   __lowerCamelCase		:							Optional[int]		 =						func(SCREAMING_SNAKE_CASE__      ,	**SCREAMING_SNAKE_CASE__						)

  print('shuffling dataset'						)
  __lowerCamelCase		:							str		 =						dataset.shuffle()
  print('Second set of iterations (after shuffling'						)
  for func, kwargs in functions_shuffled:
   print('shuffled '      ,	func.__name__      ,	str(SCREAMING_SNAKE_CASE__						)						)
   __lowerCamelCase		:							int		 =						func(
       SCREAMING_SNAKE_CASE__      ,	**SCREAMING_SNAKE_CASE__						)

 with open(SCREAMING_SNAKE_CASE__      ,	'wb'						) as f:
  f.write(json.dumps(SCREAMING_SNAKE_CASE__						).encode('utf-8'						)						)


if __name__ == "__main__":  # useful to run the profiler
       benchmark_iterating()

