





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


__UpperCAmelCase       :   Optional[int]				       =							["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif"]







class     __snake_case       (							__lowerCamelCase			):





       '''simple docstring'''







       def __init__(  self :    int							,				A :    Optional[int]							,				A :    Any							,				A :    Tuple=None							,				A :    Tuple=1		):
         __snake_case:						Optional[int]   =						tokenizer
         __snake_case:						str   =						dataset
         __snake_case:						List[str]   =						len(__a		) if n_tasks is None else n_tasks
         __snake_case:						Union[str, Any]   =						n_copies






       def __iter__(  self :    List[str]		):
         __snake_case:						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()		)
         __snake_case:						Optional[int]   =						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     __snake_case       (							__lowerCamelCase			):





       '''simple docstring'''







       def __init__(  self :    Union[str, Any]							,				A :    int							,				A :    Optional[Any]							,				A :    Optional[Any]		):
         __snake_case:						Optional[int]   =						start_length
         __snake_case:						int   =						eof_strings
         __snake_case:						List[str]   =						tokenizer






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







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

      __snake_case  ,   __snake_case:						List[Any]   =						accelerator.gather((generated_tokens, generated_tasks))
      __snake_case:						List[Any]   =						generated_tokens.cpu().numpy()
      __snake_case:						Any   =						generated_tasks.cpu().numpy()

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

  __snake_case:						int   =						[[] for _ in range(__snake_case)]
  for task, generated_tokens in gen_token_dict.items():
    for s in generated_tokens:
      __snake_case:						Dict   =						tokenizer.decode(__snake_case					,					skip_special_tokens=__snake_case					,					clean_up_tokenization_spaces=__snake_case)
      code_gens[task].append(remove_last_block(__snake_case))
  return code_gens







def 			A__							(    )   ->    Any:
  __snake_case:						int   =						HfArgumentParser(__snake_case)
  __snake_case:						List[Any]   =						parser.parse_args()

  transformers.logging.set_verbosity_error()
  # enables code execution in code_eval metric
  __snake_case:						Dict   =						args.HF_ALLOW_CODE_EVAL
  # make sure tokenizer plays nice with multiprocessing
  __snake_case:						Tuple   =						"""false"""

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

  # Use dataset load to feed to accelerate
  __snake_case:						str   =						Accelerator()
  set_seed(args.seed					,					device_specific=__snake_case)

  # Load model and tokenizer
  __snake_case:						int   =						AutoTokenizer.from_pretrained(args.model_ckpt)
  __snake_case:						Union[str, Any]   =						tokenizer.eos_token
  __snake_case:						List[Any]   =						AutoModelForCausalLM.from_pretrained(args.model_ckpt)

  # Generation settings
  __snake_case:						str   =						{
      """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					,					__snake_case					,					__snake_case)]),
  }

  # Load evaluation dataset and metric
  __snake_case:						Any   =						load_dataset("""openai_humaneval""")
  __snake_case:						Tuple   =						load_metric("""code_eval""")

  __snake_case:						int   =						args.num_tasks if args.num_tasks is not None else len(human_eval["""test"""])
  __snake_case:						Optional[Any]   =						args.n_samples // args.batch_size

  __snake_case:						Dict   =						TokenizedDataset(__snake_case					,					human_eval["""test"""]					,					n_copies=__snake_case					,					n_tasks=__snake_case)
  # do not confuse args.batch_size, which is actually the num_return_sequences
  __snake_case:						Tuple   =						DataLoader(__snake_case					,					batch_size=1)

  # Run a quick test to see if code evaluation is enabled
  try:
    __snake_case:						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

  __snake_case  ,   __snake_case:						Optional[int]   =						accelerator.prepare(__snake_case					,					__snake_case)

  __snake_case:						List[Any]   =						complete_code(
      __snake_case					,					__snake_case					,					__snake_case					,					__snake_case					,					n_tasks=__snake_case					,					batch_size=args.batch_size					,					**__snake_case					,					)

  if accelerator.is_main_process:
    __snake_case:						List[Any]   =						[]

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

    # Evaluate completions with "code_eval" metric
    __snake_case  ,   __snake_case:						Optional[int]   =						code_eval_metric.compute(
        references=__snake_case					,					predictions=__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(__snake_case					,					__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()








"""simple docstring"""





from graphs.minimum_spanning_tree_kruskal import kruskal



def       lowerCamelCase__		(			)							->		List[Any]:


					"""simple docstring"""





					_UpperCamelCase							 =  9
					_UpperCamelCase							 =  [
					    [0, 1, 4],
					    [0, 7, 8],
					    [1, 2, 8],
					    [7, 8, 7],
					    [7, 6, 1],
					    [2, 8, 2],
					    [8, 6, 6],
					    [2, 3, 7],
					    [2, 5, 4],
					    [6, 5, 2],
					    [3, 5, 14],
					    [3, 4, 9],
					    [5, 4, 10],
					    [1, 7, 11],
					]

					_UpperCamelCase							 =  kruskal(__snake_case,     __snake_case	)

					_UpperCamelCase							 =  [
					    [7, 6, 1],
					    [2, 8, 2],
					    [6, 5, 2],
					    [0, 1, 4],
					    [2, 5, 4],
					    [2, 3, 7],
					    [0, 7, 8],
					    [3, 4, 9],
					]

					assert sorted(__snake_case	) == sorted(__snake_case	)






