




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     __lowerCAmelCase (       __snake_case ):

			"""simple docstring"""







			def __init__(      self			,						lowerCamelCase__			,						lowerCamelCase__			,						lowerCamelCase__=None			,						lowerCamelCase__=1							)    ->      str:



							'''simple docstring'''
							__lowerCamelCase		       =    tokenizer
							__lowerCamelCase		       =    dataset
							__lowerCamelCase		       =    len(__SCREAMING_SNAKE_CASE							) if n_tasks is None else n_tasks
							__lowerCamelCase		       =    n_copies






			def __iter__(      self							)    ->      Any:



							'''simple docstring'''
							__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     __lowerCAmelCase (       __snake_case ):

			"""simple docstring"""







			def __init__(      self			,						lowerCamelCase__			,						lowerCamelCase__			,						lowerCamelCase__							)    ->      List[Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    start_length
							__lowerCamelCase		       =    eof_strings
							__lowerCamelCase		       =    tokenizer






			def __call__(      self			,						lowerCamelCase__			,						lowerCamelCase__			,						**lowerCamelCase__							)    ->      Optional[Any]:



							'''simple docstring'''
							__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 			lowerCamelCase_		(       UpperCamelCase__				:      List[Any]      )     ->       Tuple:





				"""simple docstring"""


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







def 			lowerCamelCase_		(       UpperCamelCase__				:      str ,     UpperCamelCase__				:      str ,     UpperCamelCase__				:      Optional[Any] ,     UpperCamelCase__				:      Optional[Any] ,     UpperCamelCase__				:      str ,     UpperCamelCase__				:      Dict=20 ,     **UpperCamelCase__				:      Dict      )     ->       List[str]:





				"""simple docstring"""


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

				__lowerCamelCase		       =    [[] for _ in range(_lowerCamelCase      )]
				for task, generated_tokens in gen_token_dict.items():
								for s in generated_tokens:
												__lowerCamelCase		       =    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_		(       )     ->       Tuple:





				"""simple docstring"""


				__lowerCamelCase		       =    HfArgumentParser(_lowerCamelCase      )
				__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=_lowerCamelCase      )

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

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

				__lowerCamelCase		       =    complete_code(
				    _lowerCamelCase ,     _lowerCamelCase ,     _lowerCamelCase ,     _lowerCamelCase ,     n_tasks=_lowerCamelCase ,     batch_size=args.batch_size ,     **_lowerCamelCase ,     )

				if accelerator.is_main_process:
								__lowerCamelCase		       =    []

								for task in tqdm(range(_lowerCamelCase      )      ):
												__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=_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 __future__ import annotations

from typing import Generic, TypeVar

_SCREAMING_SNAKE_CASE					:				Optional[Any]					    =						TypeVar('''T''')






class      a						(					Generic[T]		):






    def __init__(       self     :							List[str]       ,						__SCREAMING_SNAKE_CASE     :							T					)     ->		None:
      lowerCamelCase_         =						data
      lowerCamelCase_         =						self
      lowerCamelCase_         =						0






class      a						(					Generic[T]		):






    def __init__(       self     :							Any					)     ->		None:
      # map from node name to the node object
      lowerCamelCase_         =						{}






    def    UpperCamelCase       (       self     :							List[str]       ,						__SCREAMING_SNAKE_CASE     :							T					)     ->		None:
      # create a new set with x as its member
      lowerCamelCase_         =						DisjointSetTreeNode(__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							Optional[int]       ,						__SCREAMING_SNAKE_CASE     :							T					)     ->		DisjointSetTreeNode[T]:
      # find the set x belongs to (with path-compression)
      lowerCamelCase_         =						self.map[data]
      if elem_ref != elem_ref.parent:
        lowerCamelCase_         =						self.find_set(elem_ref.parent.data					)
      return elem_ref.parent






    def    UpperCamelCase       (       self     :							str       ,						__SCREAMING_SNAKE_CASE     :							DisjointSetTreeNode[T]       ,						__SCREAMING_SNAKE_CASE     :							DisjointSetTreeNode[T]					)     ->		None:
      # helper function for union operation
      if nodea.rank > nodea.rank:
        lowerCamelCase_         =						nodea
      else:
        lowerCamelCase_         =						nodea
        if nodea.rank == nodea.rank:
          nodea.rank += 1






    def    UpperCamelCase       (       self     :							List[str]       ,						__SCREAMING_SNAKE_CASE     :							T       ,						__SCREAMING_SNAKE_CASE     :							T					)     ->		None:
      # merge 2 disjoint sets
      self.link(self.find_set(__SCREAMING_SNAKE_CASE					)       ,						self.find_set(__SCREAMING_SNAKE_CASE					)					)






class      a						(					Generic[T]		):






    def __init__(       self     :							Optional[int]					)     ->		None:
      # connections: map from the node to the neighbouring nodes (with weights)
      lowerCamelCase_         =						{}






    def    UpperCamelCase       (       self     :							Optional[int]       ,						__SCREAMING_SNAKE_CASE     :							T					)     ->		None:
      # add a node ONLY if its not present in the graph
      if node not in self.connections:
        lowerCamelCase_         =						{}






    def    UpperCamelCase       (       self     :							Tuple       ,						__SCREAMING_SNAKE_CASE     :							T       ,						__SCREAMING_SNAKE_CASE     :							T       ,						__SCREAMING_SNAKE_CASE     :							int					)     ->		None:
      # add an edge with the given weight
      self.add_node(__SCREAMING_SNAKE_CASE					)
      self.add_node(__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						weight
      lowerCamelCase_         =						weight






    def    UpperCamelCase       (       self     :							List[Any]					)     ->		GraphUndirectedWeighted[T]:
      lowerCamelCase_         =						[]
      lowerCamelCase_         =						set()
      for start in self.connections:
        for end in self.connections[start]:
          if (start, end) not in seen:
            seen.add((end, start)					)
            edges.append((start, end, self.connections[start][end])					)
      edges.sort(key=lambda __SCREAMING_SNAKE_CASE					: x[2]					)

      # creating the disjoint set
      lowerCamelCase_         =						DisjointSetTree[T]()
      for node in self.connections:
        disjoint_set.make_set(__SCREAMING_SNAKE_CASE					)

      # MST generation
      lowerCamelCase_         =						0
      lowerCamelCase_         =						0
      lowerCamelCase_         =						GraphUndirectedWeighted[T]()
      while num_edges < len(self.connections					) - 1:
        lowerCamelCase_     ,	lowerCamelCase_     ,	lowerCamelCase_         =						edges[index]
        index += 1
        lowerCamelCase_         =						disjoint_set.find_set(__SCREAMING_SNAKE_CASE					)
        lowerCamelCase_         =						disjoint_set.find_set(__SCREAMING_SNAKE_CASE					)
        if parent_u != parent_v:
          num_edges += 1
          graph.add_edge(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE					)
          disjoint_set.union(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE					)
      return graph




