


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


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






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_ ,       A_ ,       A_=None ,       A_=1   )  ->	Optional[Any]:
         __UpperCamelCase   							=tokenizer
         __UpperCamelCase   							=dataset
         __UpperCamelCase   							=len(A_   ) if n_tasks is None else n_tasks
         __UpperCamelCase   							=n_copies






    def __iter__(      self   )  ->	int:
         __UpperCamelCase   							=[]
         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   							=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    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_ ,       A_ ,       A_   )  ->	List[Any]:
         __UpperCamelCase   							=start_length
         __UpperCamelCase   							=eof_strings
         __UpperCamelCase   							=tokenizer






    def __call__(      self ,       A_ ,       A_ ,       **A_   )  ->	str:
         __UpperCamelCase   							=self.tokenizer.batch_decode(input_ids[:, self.start_length :]   )
         __UpperCamelCase   							=[]
         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__      : Tuple    ):
     __UpperCamelCase   							=re.split('(%s)' % '|'.join(SCREAMING_SNAKE_CASE__    )      ,					SCREAMING_SNAKE_CASE__    )
     # last string should be ""
     return "".join(string_list[:-2]    )

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Optional[int]      ,					SCREAMING_SNAKE_CASE__      : Tuple      ,					SCREAMING_SNAKE_CASE__      : List[str]      ,					SCREAMING_SNAKE_CASE__      : str      ,					SCREAMING_SNAKE_CASE__      : Optional[int]      ,					SCREAMING_SNAKE_CASE__      : List[Any]=20      ,					**SCREAMING_SNAKE_CASE__      : Any    ):
     __UpperCamelCase   							=defaultdict(SCREAMING_SNAKE_CASE__    )  # dict of list of generated tokens
     for step, batch in tqdm(enumerate(SCREAMING_SNAKE_CASE__    )    ):
          with torch.no_grad():
               __UpperCamelCase   							=batch['ids'].shape[-1]
               __UpperCamelCase   							=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
               __UpperCamelCase   							=batch['task_id'].repeat(SCREAMING_SNAKE_CASE__    )
               __UpperCamelCase   							=accelerator.pad_across_processes(
                   SCREAMING_SNAKE_CASE__      ,					dim=1      ,					pad_index=tokenizer.pad_token_id    )

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

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

     __UpperCamelCase   							=[[] for _ in range(SCREAMING_SNAKE_CASE__    )]
     for task, generated_tokens in gen_token_dict.items():
          for s in generated_tokens:
               __UpperCamelCase   							=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
     __UpperCamelCase   							=HfArgumentParser(SCREAMING_SNAKE_CASE__    )
     __UpperCamelCase   							=parser.parse_args()

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

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

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

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

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

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

     __UpperCamelCase   							=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
     __UpperCamelCase   							=DataLoader(SCREAMING_SNAKE_CASE__      ,					batch_size=1    )

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

     __UpperCamelCase   							=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:
          __UpperCamelCase   							=[]

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

          # Evaluate completions with "code_eval" metric
          __UpperCamelCase							,				__UpperCamelCase   							=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 warnings

from ...utils import logging
from .image_processing_imagegpt import ImageGPTImageProcessor


_A			=						logging.get_logger(__name__)






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       *A_ ,       **A_   )  ->	None:
         warnings.warn(
             'The class ImageGPTFeatureExtractor is deprecated and will be removed in version 5 of Transformers.'
             ' Please use ImageGPTImageProcessor instead.' ,       A_ ,       )
         super().__init__(*A_ ,       **A_   )





