





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


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






class 						__snake_case   (						__lowerCAmelCase	):





 """simple docstring"""


 def __init__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase=None   ,		__lowerCamelCase=1    ):



     '''simple docstring'''




     __A       :		List[str]							=	tokenizer
     __A       :		List[Any]							=	dataset
     __A       :		Tuple							=	len(_lowerCamelCase    ) if n_tasks is None else n_tasks
     __A       :		Any							=	n_copies







 def __iter__(       self    ):



     '''simple docstring'''




     __A       :		Tuple							=	[]
     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()    )
     __A       :		Any							=	self.tokenizer(_lowerCamelCase   ,		padding=_lowerCamelCase   ,		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 						__snake_case   (						__lowerCAmelCase	):





 """simple docstring"""


 def __init__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		Union[str, Any]							=	start_length
     __A       :		Union[str, Any]							=	eof_strings
     __A       :		List[Any]							=	tokenizer







 def __call__(       self   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		**__lowerCamelCase    ):



     '''simple docstring'''




     __A       :		str							=	self.tokenizer.batch_decode(input_ids[:, self.start_length :]    )
     __A       :		List[Any]							=	[]
     for decoded_generation in decoded_generations:
         done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings    )    )
     return all(_lowerCamelCase    )







def 						__lowercase	(  snake_case_    :    List[Any]		)  ->List[str]:



    '''simple docstring'''







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







def 						__lowercase	(  snake_case_    :    Any							,snake_case_    :    List[str]							,snake_case_    :    Union[str, Any]							,snake_case_    :    Any							,snake_case_    :    List[str]							,snake_case_    :    Optional[int]=20							,**snake_case_    :    Dict		)  ->Optional[int]:



    '''simple docstring'''







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

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

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

    __A       :		Union[str, Any]							=	[[] for _ in range(__lowerCAmelCase		)]
    for task, generated_tokens in gen_token_dict.items():
        for s in generated_tokens:
            __A       :		int							=	tokenizer.decode(__lowerCAmelCase							,skip_special_tokens=__lowerCAmelCase							,clean_up_tokenization_spaces=__lowerCAmelCase		)
            code_gens[task].append(remove_last_block(__lowerCAmelCase		)		)
    return code_gens







def 						__lowercase	(  )  ->List[str]:



    '''simple docstring'''







    __A       :		int							=	HfArgumentParser(__lowerCAmelCase		)
    __A       :		List[str]							=	parser.parse_args()

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

    if args.num_workers is None:
        __A       :		Union[str, Any]							=	multiprocessing.cpu_count()

    # Use dataset load to feed to accelerate
    __A       :		int							=	Accelerator()
    set_seed(args.seed							,device_specific=__lowerCAmelCase		)

    # Load model and tokenizer
    __A       :		List[str]							=	AutoTokenizer.from_pretrained(args.model_ckpt		)
    __A       :		int							=	tokenizer.eos_token
    __A       :		Dict							=	AutoModelForCausalLM.from_pretrained(args.model_ckpt		)

    # Generation settings
    __A       :		int							=	{
        '''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
    __A       :		int							=	load_dataset('''openai_humaneval'''		)
    __A       :		Optional[Any]							=	load_metric('''code_eval'''		)

    __A       :		List[Any]							=	args.num_tasks if args.num_tasks is not None else len(human_eval['''test''']		)
    __A       :		Dict							=	args.n_samples // args.batch_size

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

    # Run a quick test to see if code evaluation is enabled
    try:
        __A       :		List[Any]							=	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

    __A					,			__A       :		Optional[Any]							=	accelerator.prepare(__lowerCAmelCase							,__lowerCAmelCase		)

    __A       :		Dict							=	complete_code(
        __lowerCAmelCase							,__lowerCAmelCase							,__lowerCAmelCase							,__lowerCAmelCase							,n_tasks=__lowerCAmelCase							,batch_size=args.batch_size							,**__lowerCAmelCase							,)

    if accelerator.is_main_process:
        __A       :		Union[str, Any]							=	[]

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

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





"""simple docstring"""


from math import sqrt





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   bool:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (
						    number >= 0
						), "'number' must been an int and positive"

						lowerCamelCase__   =True

						# 0 and 1 are none primes.
						if number <= 1:
												lowerCamelCase__   =False

						for divisor in range(2  ,       int(round(sqrt(__lowerCAmelCase							)							)							) + 1							):
												# if 'number' divisible by 'divisor' then sets 'status'
												# of false and break up the loop.
												if number % divisor == 0:
																		lowerCamelCase__   =False
																		break

    # precondition
						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							), "'status' must been from type bool"

						return status





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   List[str]:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (n > 2), "'N' must been an int and > 2"

						# beginList: contains all natural numbers from 2 up to N
						lowerCamelCase__   =list(range(2  ,       n + 1							)							)

						lowerCamelCase__   =[]  # this list will be returns.

						# actual sieve of erathostenes
						for i in range(len(__lowerCAmelCase							)							):
												for j in range(i + 1  ,       len(__lowerCAmelCase							)							):
																		if (begin_list[i] != 0) and (begin_list[j] % begin_list[i] == 0):
																								lowerCamelCase__   =0

    # filters actual prime numbers.
						lowerCamelCase__   =[x for x in begin_list if x != 0]

						# precondition
						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							), "'ans' must been from type list"

						return ans





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   Optional[int]:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (n > 2), "'N' must been an int and > 2"

						lowerCamelCase__   =[]

						# iterates over all numbers between 2 up to N+1
						# if a number is prime then appends to list 'ans'
						for number in range(2  ,       n + 1							):
												if is_prime(__lowerCAmelCase							):
																		ans.append(__lowerCAmelCase							)

    # precondition
						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							), "'ans' must been from type list"

						return ans





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   Union[str, Any]:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and number >= 0, "'number' must been an int and >= 0"

						lowerCamelCase__   =[]  # this list will be returns of the function.

						# potential prime number factors.

						lowerCamelCase__   =2

						lowerCamelCase__   =number

						if number == 0 or number == 1:
												ans.append(__lowerCAmelCase							)

						# if 'number' not prime then builds the prime factorization of 'number'
						elif not is_prime(__lowerCAmelCase							):
												while quotient != 1:
																		if is_prime(__lowerCAmelCase							) and (quotient % factor == 0):
																								ans.append(__lowerCAmelCase							)
																								quotient /= factor
																		else:
																								factor += 1

						else:
												ans.append(__lowerCAmelCase							)

						# precondition
						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							), "'ans' must been from type list"

						return ans





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   Optional[int]:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (
						    number >= 0
						), "'number' bust been an int and >= 0"

						lowerCamelCase__   =0

						# prime factorization of 'number'
						lowerCamelCase__   =prime_factorization(__lowerCAmelCase							)

						lowerCamelCase__   =max(__lowerCAmelCase							)

						# precondition
						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							), "'ans' must been from type int"

						return ans





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   List[Any]:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (
						    number >= 0
						), "'number' bust been an int and >= 0"

						lowerCamelCase__   =0

						# prime factorization of 'number'
						lowerCamelCase__   =prime_factorization(__lowerCAmelCase							)

						lowerCamelCase__   =min(__lowerCAmelCase							)

						# precondition
						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							), "'ans' must been from type int"

						return ans





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   Union[str, Any]:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							), "'number' must been an int"
						assert isinstance(number % 2 == 0  ,       __lowerCAmelCase							), "compare bust been from type bool"

						return number % 2 == 0





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   List[Any]:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							), "'number' must been an int"
						assert isinstance(number % 2 != 0  ,       __lowerCAmelCase							), "compare bust been from type bool"

						return number % 2 != 0





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   int:




						'''simple docstring'''



						assert (
						    isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (number > 2) and is_even(__lowerCAmelCase							)
						), "'number' must been an int, even and > 2"

						lowerCamelCase__   =[]  # this list will returned

						# creates a list of prime numbers between 2 up to 'number'
						lowerCamelCase__   =get_prime_numbers(__lowerCAmelCase							)
						lowerCamelCase__   =len(__lowerCAmelCase							)

						# run variable for while-loops.
						lowerCamelCase__   =0
						lowerCamelCase__   =None

						# exit variable. for break up the loops
						lowerCamelCase__   =True

						while i < len_pn and loop:
												lowerCamelCase__   =i + 1

												while j < len_pn and loop:
																		if prime_numbers[i] + prime_numbers[j] == number:
																								lowerCamelCase__   =False
																								ans.append(prime_numbers[i]							)
																								ans.append(prime_numbers[j]							)

																		j += 1

												i += 1

						# precondition
						assert (
						    isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							)
						    and (len(__lowerCAmelCase							) == 2)
						    and (ans[0] + ans[1] == number)
						    and is_prime(ans[0]							)
						    and is_prime(ans[1]							)
						), "'ans' must contains two primes. And sum of elements must been eq 'number'"

						return ans





def      lowerCamelCase_					(   __lowerCAmelCase  ,       __lowerCAmelCase							)				->   Tuple:




						'''simple docstring'''



						assert (
						    isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							)
						    and isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							)
						    and (numbera >= 0)
						    and (numbera >= 0)
						), "'number1' and 'number2' must been positive integer."

						lowerCamelCase__   =0

						while numbera != 0:
												lowerCamelCase__   =numbera % numbera
												lowerCamelCase__   =numbera
												lowerCamelCase__   =rest

						# precondition
						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (
						    numbera >= 0
						), "'number' must been from type int and positive"

						return numbera





def      lowerCamelCase_					(   __lowerCAmelCase  ,       __lowerCAmelCase							)				->   Union[str, Any]:




						'''simple docstring'''



						assert (
						    isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							)
						    and isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							)
						    and (numbera >= 1)
						    and (numbera >= 1)
						), "'number1' and 'number2' must been positive integer."

						lowerCamelCase__   =1  # actual answer that will be return.

						# for kgV (x,1)
						if numbera > 1 and numbera > 1:
												# builds the prime factorization of 'number1' and 'number2'
												lowerCamelCase__   =prime_factorization(__lowerCAmelCase							)
												lowerCamelCase__   =prime_factorization(__lowerCAmelCase							)

						elif numbera == 1 or numbera == 1:
												lowerCamelCase__   =[]
												lowerCamelCase__   =[]
												lowerCamelCase__   =max(__lowerCAmelCase  ,       __lowerCAmelCase							)

						lowerCamelCase__   =0
						lowerCamelCase__   =0

						lowerCamelCase__   =[]  # captured numbers int both 'primeFac1' and 'primeFac2'

						# iterates through primeFac1
						for n in prime_fac_a:
												if n not in done:
																		if n in prime_fac_a:
																								lowerCamelCase__   =prime_fac_a.count(__lowerCAmelCase							)
																								lowerCamelCase__   =prime_fac_a.count(__lowerCAmelCase							)

																								for _ in range(max(__lowerCAmelCase  ,       __lowerCAmelCase							)							):
																														ans *= n

																		else:
																								lowerCamelCase__   =prime_fac_a.count(__lowerCAmelCase							)

																								for _ in range(__lowerCAmelCase							):
																														ans *= n

																		done.append(__lowerCAmelCase							)

    # iterates through primeFac2
						for n in prime_fac_a:
												if n not in done:
																		lowerCamelCase__   =prime_fac_a.count(__lowerCAmelCase							)

																		for _ in range(__lowerCAmelCase							):
																								ans *= n

																		done.append(__lowerCAmelCase							)

    # precondition
						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (
						    ans >= 0
						), "'ans' must been from type int and positive"

						return ans





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   Optional[Any]:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (n >= 0), "'number' must been a positive int"

						lowerCamelCase__   =0
						lowerCamelCase__   =2  # this variable holds the answer

						while index < n:
												index += 1

												ans += 1  # counts to the next number

												# if ans not prime then
												# runs to the next prime number.
												while not is_prime(__lowerCAmelCase							):
																		ans += 1

    # precondition
						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and is_prime(
						    __lowerCAmelCase							), "'ans' must been a prime number and from type int"

						return ans





def      lowerCamelCase_					(   __lowerCAmelCase  ,       __lowerCAmelCase							)				->   int:




						'''simple docstring'''



						assert (
						    is_prime(__lowerCAmelCase							) and is_prime(__lowerCAmelCase							) and (p_number_a < p_number_a)
						), "The arguments must been prime numbers and 'pNumber1' < 'pNumber2'"

						lowerCamelCase__   =p_number_a + 1  # jump to the next number

						lowerCamelCase__   =[]  # this list will be returns.

						# if number is not prime then
						# fetch the next prime number.
						while not is_prime(__lowerCAmelCase							):
												number += 1

						while number < p_number_a:
												ans.append(__lowerCAmelCase							)

												number += 1

												# fetch the next prime number.
												while not is_prime(__lowerCAmelCase							):
																		number += 1

    # precondition
						assert (
						    isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							)
						    and ans[0] != p_number_a
						    and ans[len(__lowerCAmelCase							) - 1] != p_number_a
						), "'ans' must been a list without the arguments"

						# 'ans' contains not 'pNumber1' and 'pNumber2' !
						return ans





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   int:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (n >= 1), "'n' must been int and >= 1"

						lowerCamelCase__   =[]  # will be returned.

						for divisor in range(1  ,       n + 1							):
												if n % divisor == 0:
																		ans.append(__lowerCAmelCase							)

    # precondition
						assert ans[0] == 1 and ans[len(__lowerCAmelCase							) - 1] == n, "Error in function getDivisiors(...)"

						return ans





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   Union[str, Any]:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (
						    number > 1
						), "'number' must been an int and >= 1"

						lowerCamelCase__   =get_divisors(__lowerCAmelCase							)

						# precondition
						assert (
						    isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							)
						    and (divisors[0] == 1)
						    and (divisors[len(__lowerCAmelCase							) - 1] == number)
						), "Error in help-function getDivisiors(...)"

						# summed all divisors up to 'number' (exclusive), hence [:-1]
						return sum(divisors[:-1]							) == number





def      lowerCamelCase_					(   __lowerCAmelCase  ,       __lowerCAmelCase							)				->   Dict:




						'''simple docstring'''



						assert (
						    isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							)
						    and isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							)
						    and (denominator != 0)
						), "The arguments must been from type int and 'denominator' != 0"

						# build the greatest common divisor of numerator and denominator.
						lowerCamelCase__   =gcd(abs(__lowerCAmelCase							)  ,       abs(__lowerCAmelCase							)							)

						# precondition
						assert (
						    isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							)
						    and (numerator % gcd_of_fraction == 0)
						    and (denominator % gcd_of_fraction == 0)
						), "Error in function gcd(...,...)"

						return (numerator // gcd_of_fraction, denominator // gcd_of_fraction)





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   List[Any]:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (n >= 0), "'n' must been a int and >= 0"

						lowerCamelCase__   =1  # this will be return.

						for factor in range(1  ,       n + 1							):
												ans *= factor

						return ans





def      lowerCamelCase_					(   __lowerCAmelCase							)				->   Tuple:




						'''simple docstring'''



						assert isinstance(__lowerCAmelCase  ,       __lowerCAmelCase							) and (n >= 0), "'n' must been an int and >= 0"

						lowerCamelCase__   =0
						lowerCamelCase__   =1
						lowerCamelCase__   =1  # this will be return

						for _ in range(n - 1							):
												lowerCamelCase__   =ans
												ans += fiba
												lowerCamelCase__   =tmp

						return ans


