






"""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        UpperCamelCase			(	__lowerCamelCase      ):


 def __init__(     self		:						str							,						UpperCAmelCase__		:						Dict							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						Any=None							,						UpperCAmelCase__		:						Union[str, Any]=1							)					->       Optional[Any]:
   _a  :       Optional[int]             =	tokenizer
   _a  :       List[Any]             =	dataset
   _a  :       Any             =	len(__lowercase							) if n_tasks is None else n_tasks
   _a  :       List[Any]             =	n_copies







 def __iter__(     self		:						Tuple							)					->       Optional[Any]:
   _a  :       Union[str, 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()							)
   _a  :       Optional[int]             =	self.tokenizer(__lowercase							,						padding=__lowercase							,						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			(	__lowerCamelCase      ):


 def __init__(     self		:						Union[str, Any]							,						UpperCAmelCase__		:						List[Any]							,						UpperCAmelCase__		:						int							,						UpperCAmelCase__		:						List[str]							)					->       Tuple:
   _a  :       List[str]             =	start_length
   _a  :       Any             =	eof_strings
   _a  :       Optional[Any]             =	tokenizer







 def __call__(     self		:						str							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						Any							,						**UpperCAmelCase__		:						Optional[int]							)					->       Dict:
   _a  :       Tuple             =	self.tokenizer.batch_decode(input_ids[:, self.start_length :]							)
   _a  :       Union[str, Any]             =	[]
   for decoded_generation in decoded_generations:
     done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings							)							)
   return all(__lowercase							)







def 			lowerCAmelCase__     (     UpperCamelCase__				):



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

def 			lowerCAmelCase__     (     UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__=2_0     ,   **UpperCamelCase__				):



  '''simple docstring'''
  _a  :       Optional[int]             =	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  :       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  :       str             =	batch['''task_id'''].repeat(__lowerCAmelCase				)
      _a  :       Optional[int]             =	accelerator.pad_across_processes(
          __lowerCAmelCase     ,   dim=1     ,   pad_index=tokenizer.pad_token_id				)

      _a  :       Any             =	accelerator.gather((generated_tokens, generated_tasks)				)
      _a  :       Union[str, Any]             =	generated_tokens.cpu().numpy()
      _a  :       Optional[Any]             =	generated_tasks.cpu().numpy()

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

  _a  :       List[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 			lowerCAmelCase__     (     ):



  '''simple docstring'''
  _a  :       List[str]             =	HfArgumentParser(__lowerCAmelCase				)
  _a  :       Optional[Any]             =	parser.parse_args()

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

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

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

  # Load model and tokenizer
  _a  :       str             =	AutoTokenizer.from_pretrained(args.model_ckpt				)
  _a  :       List[str]             =	tokenizer.eos_token
  _a  :       Optional[Any]             =	AutoModelForCausalLM.from_pretrained(args.model_ckpt				)

  # Generation settings
  _a  :       List[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     ,   __lowerCAmelCase     ,   __lowerCAmelCase				)]				),
  }

  # Load evaluation dataset and metric
  _a  :       List[str]             =	load_dataset("""openai_humaneval"""				)
  _a  :       str             =	load_metric("""code_eval"""				)

  _a  :       List[Any]             =	args.num_tasks if args.num_tasks is not None else len(human_eval["""test"""]				)
  _a  :       str             =	args.n_samples // args.batch_size

  _a  :       Tuple             =	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  :       Tuple             =	DataLoader(__lowerCAmelCase     ,   batch_size=1				)

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

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

    # Evaluate completions with "code_eval" metric
    _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()






import math
import sys



def 			_lowerCAmelCase       (					__lowerCAmelCase      )   ->					str:


				"""simple docstring"""

				snake_case__	:      Optional[Any]     		=	''''''
				try:
								with open(__lowerCAmelCase  ,  '''rb'''      ) as binary_file:
												snake_case__	:      int     		=	binary_file.read()
								for dat in data:
												snake_case__	:      Any     		=	f"""{dat:08b}"""
												result += curr_byte
								return result
				except OSError:
								print('''File not accessible'''      )
								sys.exit()



def 			_lowerCAmelCase       (					__lowerCAmelCase      )   ->					str:


				"""simple docstring"""

				snake_case__	:      List[str]     		=	{'''0''': '''0''', '''1''': '''1'''}
				snake_case__		, snake_case__	:      List[Any]     		=	'''''', ''''''
				snake_case__	:      Tuple     		=	len(__lowerCAmelCase      )

				for i in range(len(__lowerCAmelCase      )      ):
								curr_string += data_bits[i]
								if curr_string not in lexicon:
												continue

								snake_case__	:      Tuple     		=	lexicon[curr_string]
								result += last_match_id
								snake_case__	:      Any     		=	last_match_id + '''0'''

								if math.loga(__lowerCAmelCase      ).is_integer():
												snake_case__	:      Tuple     		=	{}
												for curr_key in list(__lowerCAmelCase      ):
																snake_case__	:      Union[str, Any]     		=	lexicon.pop(__lowerCAmelCase      )
												snake_case__	:      Optional[Any]     		=	new_lex

								snake_case__	:      Tuple     		=	last_match_id + '''1'''
								index += 1
								snake_case__	:      Dict     		=	''''''
				return result



def 			_lowerCAmelCase       (					__lowerCAmelCase  ,  __lowerCAmelCase      )   ->					None:


				"""simple docstring"""

				snake_case__	:      Dict     		=	8
				try:
								with open(__lowerCAmelCase  ,  '''wb'''      ) as opened_file:
												snake_case__	:      Union[str, Any]     		=	[
												    to_write[i : i + byte_length]
												    for i in range(0  ,  len(__lowerCAmelCase      )  ,  __lowerCAmelCase      )
												]

												if len(result_byte_array[-1]      ) % byte_length == 0:
																result_byte_array.append('''10000000'''      )
												else:
																result_byte_array[-1] += "1" + "0" * (
																    byte_length - len(result_byte_array[-1]      ) - 1
																)

												for elem in result_byte_array[:-1]:
																opened_file.write(int(__lowerCAmelCase  ,  2      ).to_bytes(1  ,  byteorder='''big'''      )      )
				except OSError:
								print('''File not accessible'''      )
								sys.exit()



def 			_lowerCAmelCase       (					__lowerCAmelCase      )   ->					str:


				"""simple docstring"""

				snake_case__	:      Any     		=	0
				for letter in data_bits:
								if letter == "1":
												break
								counter += 1

				snake_case__	:      Optional[int]     		=	data_bits[counter:]
				snake_case__	:      int     		=	data_bits[counter + 1 :]
				return data_bits



def 			_lowerCAmelCase       (					__lowerCAmelCase  ,  __lowerCAmelCase      )   ->					None:


				"""simple docstring"""

				snake_case__	:      Union[str, Any]     		=	read_file_binary(__lowerCAmelCase      )
				snake_case__	:      List[str]     		=	remove_prefix(__lowerCAmelCase      )
				snake_case__	:      Any     		=	decompress_data(__lowerCAmelCase      )
				write_file_binary(__lowerCAmelCase  ,  __lowerCAmelCase      )


if __name__ == "__main__":
				compress(sys.argv[1], sys.argv[2])

