






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


SCREAMING_SNAKE_CASE__     =       ['\nclass', '\ndef', '\n#', '\n@', '\nprint', '\nif']






class     a_						(      lowerCamelCase    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE=None      ,		_SCREAMING_SNAKE_CASE=1				)  ->				Any:



        """simple docstring"""






        UpperCamelCase      = tokenizer
        UpperCamelCase      = dataset
        UpperCamelCase      = len(_SCREAMING_SNAKE_CASE				) if n_tasks is None else n_tasks
        UpperCamelCase      = n_copies


  def __iter__(				self				)  ->				Any:



        """simple docstring"""






        UpperCamelCase      = []
        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()				)
        UpperCamelCase      = 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     a_						(      lowerCamelCase    ):

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)  ->				int:



        """simple docstring"""






        UpperCamelCase      = start_length
        UpperCamelCase      = eof_strings
        UpperCamelCase      = tokenizer


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



        """simple docstring"""






        UpperCamelCase      = self.tokenizer.batch_decode(input_ids[:, self.start_length :]				)
        UpperCamelCase      = []
        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    lowercase__			(				__UpperCamelCase   )->			Optional[Any]:
      UpperCamelCase      = 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   )->			List[str]:
      UpperCamelCase      = defaultdict(__UpperCamelCase   )  # dict of list of generated tokens
      for step, batch in tqdm(enumerate(__UpperCamelCase   )   ):
            with torch.no_grad():
                  UpperCamelCase      = batch["""ids"""].shape[-1]
                  UpperCamelCase      = accelerator.unwrap_model(__UpperCamelCase   ).generate(
                      input_ids=batch["""ids"""][:, : batch["""input_len"""]]		,						num_return_sequences=__UpperCamelCase		,						**__UpperCamelCase   )
                  # each task is generated batch_size times
                  UpperCamelCase      = batch["""task_id"""].repeat(__UpperCamelCase   )
                  UpperCamelCase      = accelerator.pad_across_processes(
                      __UpperCamelCase		,						dim=1		,						pad_index=tokenizer.pad_token_id   )

                  UpperCamelCase  ,UpperCamelCase      = accelerator.gather((generated_tokens, generated_tasks)   )
                  UpperCamelCase      = generated_tokens.cpu().numpy()
                  UpperCamelCase      = generated_tasks.cpu().numpy()

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

      UpperCamelCase      = [[] for _ in range(__UpperCamelCase   )]
      for task, generated_tokens in gen_token_dict.items():
            for s in generated_tokens:
                  UpperCamelCase      = 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__			(				)->			Any:
      # Setup configuration
      UpperCamelCase      = HfArgumentParser(__UpperCamelCase   )
      UpperCamelCase      = parser.parse_args()

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

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

      # Use dataset load to feed to accelerate
      UpperCamelCase      = Accelerator()
      set_seed(args.seed		,						device_specific=__UpperCamelCase   )

      # Load model and tokenizer
      UpperCamelCase      = AutoTokenizer.from_pretrained(args.model_ckpt   )
      UpperCamelCase      = tokenizer.eos_token
      UpperCamelCase      = AutoModelForCausalLM.from_pretrained(args.model_ckpt   )

      # Generation settings
      UpperCamelCase      = {
          """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
      UpperCamelCase      = load_dataset("""openai_humaneval"""   )
      UpperCamelCase      = load_metric("""code_eval"""   )

      UpperCamelCase      = args.num_tasks if args.num_tasks is not None else len(human_eval["""test"""]   )
      UpperCamelCase      = args.n_samples // args.batch_size

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

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

      UpperCamelCase  ,UpperCamelCase      = accelerator.prepare(__UpperCamelCase		,						__UpperCamelCase   )

      UpperCamelCase      = complete_code(
          __UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase		,						n_tasks=__UpperCamelCase		,						batch_size=args.batch_size		,						**__UpperCamelCase		,						)

      if accelerator.is_main_process:
            UpperCamelCase      = []

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

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








'''simple docstring'''

from typing import Dict, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import flip_channel_order, resize, to_channel_dimension_format, to_pil_image
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_pytesseract_available, is_vision_available, logging, requires_backends


if is_vision_available():
     import PIL

# soft dependency
if is_pytesseract_available():
     import pytesseract

SCREAMING_SNAKE_CASE__     =       logging.get_logger(__name__)







def    lowercase__			(				__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase   )->			int:
      return [
          int(1000 * (box[0] / width)   ),
          int(1000 * (box[1] / height)   ),
          int(1000 * (box[2] / width)   ),
          int(1000 * (box[3] / height)   ),
      ]







def    lowercase__			(				__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase = None   )->			Dict:
      UpperCamelCase      = tesseract_config if tesseract_config is not None else """"""

      # apply OCR
      UpperCamelCase      = to_pil_image(__UpperCamelCase   )
      UpperCamelCase  ,UpperCamelCase      = pil_image.size
      UpperCamelCase      = pytesseract.image_to_data(__UpperCamelCase		,						lang=__UpperCamelCase		,						output_type="""dict"""		,						config=__UpperCamelCase   )
      UpperCamelCase  ,UpperCamelCase  ,UpperCamelCase  ,UpperCamelCase  ,UpperCamelCase      = data["""text"""], data["""left"""], data["""top"""], data["""width"""], data["""height"""]

      # filter empty words and corresponding coordinates
      UpperCamelCase      = [idx for idx, word in enumerate(__UpperCamelCase   ) if not word.strip()]
      UpperCamelCase      = [word for idx, word in enumerate(__UpperCamelCase   ) if idx not in irrelevant_indices]
      UpperCamelCase      = [coord for idx, coord in enumerate(__UpperCamelCase   ) if idx not in irrelevant_indices]
      UpperCamelCase      = [coord for idx, coord in enumerate(__UpperCamelCase   ) if idx not in irrelevant_indices]
      UpperCamelCase      = [coord for idx, coord in enumerate(__UpperCamelCase   ) if idx not in irrelevant_indices]
      UpperCamelCase      = [coord for idx, coord in enumerate(__UpperCamelCase   ) if idx not in irrelevant_indices]

      # turn coordinates into (left, top, left+width, top+height) format
      UpperCamelCase      = []
      for x, y, w, h in zip(__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase   ):
            UpperCamelCase      = [x, y, x + w, y + h]
            actual_boxes.append(__UpperCamelCase   )

      # finally, normalize the bounding boxes
      UpperCamelCase      = []
      for box in actual_boxes:
            normalized_boxes.append(normalize_box(__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase   )   )

      assert len(__UpperCamelCase   ) == len(__UpperCamelCase   ), "Not as many words as there are bounding boxes"

      return words, normalized_boxes






class     a_						(      lowerCamelCase    ):
  lowercase							=				["""pixel_values"""]

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE = True      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = PILImageResampling.BILINEAR      ,		_SCREAMING_SNAKE_CASE = True      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = ""      ,		**_SCREAMING_SNAKE_CASE      ,		)  ->				None:



        """simple docstring"""






        super().__init__(**_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = size if size is not None else {"""height""": 224, """width""": 224}
        UpperCamelCase      = get_size_dict(_SCREAMING_SNAKE_CASE				)

        UpperCamelCase      = do_resize
        UpperCamelCase      = size
        UpperCamelCase      = resample
        UpperCamelCase      = apply_ocr
        UpperCamelCase      = ocr_lang
        UpperCamelCase      = tesseract_config

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE = PILImageResampling.BILINEAR      ,		_SCREAMING_SNAKE_CASE = None      ,		**_SCREAMING_SNAKE_CASE      ,		)  ->				np.ndarray:



        """simple docstring"""






        UpperCamelCase      = get_size_dict(_SCREAMING_SNAKE_CASE				)
        if "height" not in size or "width" not in size:
              raise ValueError(F"The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}"				)
        UpperCamelCase      = (size["""height"""], size["""width"""])
        return resize(_SCREAMING_SNAKE_CASE      ,		size=_SCREAMING_SNAKE_CASE      ,		resample=_SCREAMING_SNAKE_CASE      ,		data_format=_SCREAMING_SNAKE_CASE      ,		**_SCREAMING_SNAKE_CASE				)


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = ChannelDimension.FIRST      ,		**_SCREAMING_SNAKE_CASE      ,		)  ->				PIL.Image.Image:



        """simple docstring"""






        UpperCamelCase      = do_resize if do_resize is not None else self.do_resize
        UpperCamelCase      = size if size is not None else self.size
        UpperCamelCase      = get_size_dict(_SCREAMING_SNAKE_CASE				)
        UpperCamelCase      = resample if resample is not None else self.resample
        UpperCamelCase      = apply_ocr if apply_ocr is not None else self.apply_ocr
        UpperCamelCase      = ocr_lang if ocr_lang is not None else self.ocr_lang
        UpperCamelCase      = tesseract_config if tesseract_config is not None else self.tesseract_config

        UpperCamelCase      = make_list_of_images(_SCREAMING_SNAKE_CASE				)

        if not valid_images(_SCREAMING_SNAKE_CASE				):
              raise ValueError(
                  """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
                  """torch.Tensor, tf.Tensor or jax.ndarray."""				)

        if do_resize and size is None:
              raise ValueError("""Size must be specified if do_resize is True."""				)

        # All transformations expect numpy arrays.
        UpperCamelCase      = [to_numpy_array(_SCREAMING_SNAKE_CASE				) for image in images]

        if apply_ocr:
              requires_backends(self      ,		"""pytesseract"""				)
              UpperCamelCase      = []
              UpperCamelCase      = []
              for image in images:
                    UpperCamelCase  ,UpperCamelCase      = apply_tesseract(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)
                    words_batch.append(_SCREAMING_SNAKE_CASE				)
                    boxes_batch.append(_SCREAMING_SNAKE_CASE				)

        if do_resize:
              UpperCamelCase      = [self.resize(image=_SCREAMING_SNAKE_CASE      ,		size=_SCREAMING_SNAKE_CASE      ,		resample=_SCREAMING_SNAKE_CASE				) for image in images]

        # flip color channels from RGB to BGR (as Detectron2 requires this)
        UpperCamelCase      = [flip_channel_order(_SCREAMING_SNAKE_CASE				) for image in images]
        UpperCamelCase      = [to_channel_dimension_format(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				) for image in images]

        UpperCamelCase      = BatchFeature(data={"""pixel_values""": images}      ,		tensor_type=_SCREAMING_SNAKE_CASE				)

        if apply_ocr:
              UpperCamelCase      = words_batch
              UpperCamelCase      = boxes_batch
        return data
