



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_						:     Tuple	     =		['\nclass', '\ndef', '\n#', '\n@', '\nprint', '\nif']





class    lowerCamelCase__  (						UpperCAmelCase_):
     """simple docstring"""



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






           '''simple docstring'''



           lowerCamelCase								=		tokenizer
           lowerCamelCase								=		dataset
           lowerCamelCase								=		len(__a    ) if n_tasks is None else n_tasks
           lowerCamelCase								=		n_copies







     def __iter__(self    ):






           '''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(__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    lowerCamelCase__  (						UpperCAmelCase_):
     """simple docstring"""



     def __init__(self				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		start_length
           lowerCamelCase								=		eof_strings
           lowerCamelCase								=		tokenizer







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






           '''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(__a    )



def     __lowercase(	UpperCAmelCase__ ):


      """simple docstring"""



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

def     __lowercase(	UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__=20   ,   **UpperCAmelCase__ ):


      """simple docstring"""



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

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

def     __lowercase(	):


      """simple docstring"""



      lowerCamelCase								=		HfArgumentParser(UpperCAmelCase__ )
      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=UpperCAmelCase__ )

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

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

      lowerCamelCase								=		complete_code(
          UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__   ,   UpperCAmelCase__   ,   n_tasks=UpperCAmelCase__   ,   batch_size=args.batch_size   ,   **UpperCAmelCase__   ,   )

      if accelerator.is_main_process:
            lowerCamelCase								=		[]

            for task in tqdm(range(UpperCAmelCase__ ) ):
                  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=UpperCAmelCase__   ,   predictions=UpperCAmelCase__   ,   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(UpperCAmelCase__   ,   UpperCAmelCase__ )


# 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()




from __future__ import annotations

from collections import deque
from collections.abc import Iterator
from dataclasses import dataclass





@dataclass
class    lowerCamelCase__  :
     """simple docstring"""



     _A       =			42
     _A       =			42





class    lowerCamelCase__  :
     """simple docstring"""



     def __init__(self				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		[[] for _ in range(__a    )]
           lowerCamelCase								=		size







     def __getitem__(self				,			__a    ):






           '''simple docstring'''



           return iter(self._graph[vertex]    )







     @property
     def 						_a      (self    ):






           '''simple docstring'''



           return self._size







     def 						_a      (self				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           if weight not in (0, 1):
                 raise ValueError("Edge weight must be either 0 or 1."    )

           if to_vertex < 0 or to_vertex >= self.size:
                 raise ValueError("Vertex indexes must be in [0; size)."    )

           self._graph[from_vertex].append(Edge(__a				,			__a    )    )







     def 						_a      (self				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		deque([start_vertex]    )
           lowerCamelCase								=		[None] * self.size
           lowerCamelCase								=		0

           while queue:
                 lowerCamelCase								=		queue.popleft()
                 lowerCamelCase								=		distances[current_vertex]
                 if current_distance is None:
                       continue

                 for edge in self[current_vertex]:
                       lowerCamelCase								=		current_distance + edge.weight
                       lowerCamelCase								=		distances[edge.destination_vertex]
                       if (
                           isinstance(__a				,			__a    )
                           and new_distance >= dest_vertex_distance
                       ):
                             continue
                       lowerCamelCase								=		new_distance
                       if edge.weight == 0:
                             queue.appendleft(edge.destination_vertex    )
                       else:
                             queue.append(edge.destination_vertex    )

           if distances[finish_vertex] is None:
                 raise ValueError("No path from start_vertex to finish_vertex."    )

           return distances[finish_vertex]


if __name__ == "__main__":
     import doctest

     doctest.testmod()