

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


UpperCAmelCase					:			Dict       =	['\nclass', '\ndef', '\n#', '\n@', '\nprint', '\nif']







class 		lowerCAmelCase__  ( a   ):



 """simple docstring"""



 def __init__(						self   :			Union[str, Any]   ,				__SCREAMING_SNAKE_CASE   :			str   ,				__SCREAMING_SNAKE_CASE   :			Optional[Any]   ,				__SCREAMING_SNAKE_CASE   :			List[str]=None   ,				__SCREAMING_SNAKE_CASE   :			Any=1     )    ->					Optional[int]:


      """simple docstring"""




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


 def __iter__(						self   :			Optional[int]     )    ->					Any:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = []
      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()     )
      __SCREAMING_SNAKE_CASE      = 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__  ( a   ):



 """simple docstring"""



 def __init__(						self   :			str   ,				__SCREAMING_SNAKE_CASE   :			Optional[int]   ,				__SCREAMING_SNAKE_CASE   :			Tuple   ,				__SCREAMING_SNAKE_CASE   :			str     )    ->					str:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = start_length
      __SCREAMING_SNAKE_CASE      = eof_strings
      __SCREAMING_SNAKE_CASE      = tokenizer


 def __call__(						self   :			str   ,				__SCREAMING_SNAKE_CASE   :			int   ,				__SCREAMING_SNAKE_CASE   :			str   ,				**__SCREAMING_SNAKE_CASE   :			Any     )    ->					List[str]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = self.tokenizer.batch_decode(input_ids[:, self.start_length :]     )
      __SCREAMING_SNAKE_CASE      = []
      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   a__			(    a__		):




     """simple docstring"""


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







def   a__			(    a__ ,       a__ ,       a__ ,       a__ ,       a__ ,       a__=20 ,       **a__		):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE      = defaultdict(a__		)  # dict of list of generated tokens
     for step, batch in tqdm(enumerate(a__		)		):
          with torch.no_grad():
               __SCREAMING_SNAKE_CASE      = batch["""ids"""].shape[-1]
               __SCREAMING_SNAKE_CASE      = accelerator.unwrap_model(a__		).generate(
                   input_ids=batch["""ids"""][:, : batch["""input_len"""]] ,       num_return_sequences=a__ ,       **a__		)
               # each task is generated batch_size times
               __SCREAMING_SNAKE_CASE      = batch["""task_id"""].repeat(a__		)
               __SCREAMING_SNAKE_CASE      = accelerator.pad_across_processes(
                   a__ ,       dim=1 ,       pad_index=tokenizer.pad_token_id		)

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

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

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







def   a__			(    ):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE      = HfArgumentParser(a__		)
     __SCREAMING_SNAKE_CASE      = parser.parse_args()

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

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

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

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

     # Generation settings
     __SCREAMING_SNAKE_CASE      = {
         """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 ,       a__ ,       a__		)]		),
     }

     # Load evaluation dataset and metric
     __SCREAMING_SNAKE_CASE      = load_dataset("""openai_humaneval"""		)
     __SCREAMING_SNAKE_CASE      = load_metric("""code_eval"""		)

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

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

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

     __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = accelerator.prepare(a__ ,       a__		)

     __SCREAMING_SNAKE_CASE      = complete_code(
         a__ ,       a__ ,       a__ ,       a__ ,       n_tasks=a__ ,       batch_size=args.batch_size ,       **a__ ,       )

     if accelerator.is_main_process:
          __SCREAMING_SNAKE_CASE      = []

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

          # Evaluate completions with "code_eval" metric
          __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = code_eval_metric.compute(
              references=a__ ,       predictions=a__ ,       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(a__ ,       a__		)


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



import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference







class 		lowerCAmelCase__  ( a					,   unittest.TestCase   ):



 """simple docstring"""



 lowerCAmelCase__       = ShapEImgaImgPipeline
 lowerCAmelCase__       = ["image"]
 lowerCAmelCase__       = ["image"]
 lowerCAmelCase__       = [
     "num_images_per_prompt",
     "num_inference_steps",
     "generator",
     "latents",
     "guidance_scale",
     "frame_size",
     "output_type",
     "return_dict",
 ]
 lowerCAmelCase__       = False
 @property
 def  UpperCAmelCase__						(						self   :			Tuple     )    ->					Tuple:


      """simple docstring"""




      return 32
 @property
 def  UpperCAmelCase__						(						self   :			str     )    ->					int:


      """simple docstring"""




      return 32
 @property
 def  UpperCAmelCase__						(						self   :			List[str]     )    ->					Union[str, Any]:


      """simple docstring"""




      return self.time_input_dim * 4
 @property
 def  UpperCAmelCase__						(						self   :			Dict     )    ->					str:


      """simple docstring"""




      return 8
 @property
 def  UpperCAmelCase__						(						self   :			Optional[int]     )    ->					Optional[Any]:


      """simple docstring"""




      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = CLIPVisionConfig(
          hidden_size=self.text_embedder_hidden_size   ,				image_size=64   ,				projection_dim=self.text_embedder_hidden_size   ,				intermediate_size=37   ,				num_attention_heads=4   ,				num_channels=3   ,				num_hidden_layers=5   ,				patch_size=1   ,				)

      __SCREAMING_SNAKE_CASE      = CLIPVisionModel(__SCREAMING_SNAKE_CASE     )
      return model
 @property
 def  UpperCAmelCase__						(						self   :			Any     )    ->					Optional[int]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = CLIPImageProcessor(
          crop_size=224   ,				do_center_crop=__SCREAMING_SNAKE_CASE   ,				do_normalize=__SCREAMING_SNAKE_CASE   ,				do_resize=__SCREAMING_SNAKE_CASE   ,				image_mean=[0.48145466, 0.4578275, 0.40821073]   ,				image_std=[0.26862954, 0.26130258, 0.27577711]   ,				resample=3   ,				size=224   ,				)

      return image_processor
 @property
 def  UpperCAmelCase__						(						self   :			Union[str, Any]     )    ->					Tuple:


      """simple docstring"""




      torch.manual_seed(0     )

      __SCREAMING_SNAKE_CASE      = {
          """num_attention_heads""": 2,
          """attention_head_dim""": 16,
          """embedding_dim""": self.time_input_dim,
          """num_embeddings""": 32,
          """embedding_proj_dim""": self.text_embedder_hidden_size,
          """time_embed_dim""": self.time_embed_dim,
          """num_layers""": 1,
          """clip_embed_dim""": self.time_input_dim * 2,
          """additional_embeddings""": 0,
          """time_embed_act_fn""": """gelu""",
          """norm_in_type""": """layer""",
          """embedding_proj_norm_type""": """layer""",
          """encoder_hid_proj_type""": None,
          """added_emb_type""": None,
      }

      __SCREAMING_SNAKE_CASE      = PriorTransformer(**__SCREAMING_SNAKE_CASE     )
      return model
 @property
 def  UpperCAmelCase__						(						self   :			str     )    ->					List[str]:


      """simple docstring"""




      torch.manual_seed(0     )

      __SCREAMING_SNAKE_CASE      = {
          """param_shapes""": (
              (self.renderer_dim, 93),
              (self.renderer_dim, 8),
              (self.renderer_dim, 8),
              (self.renderer_dim, 8),
          ),
          """d_latent""": self.time_input_dim,
          """d_hidden""": self.renderer_dim,
          """n_output""": 12,
          """background""": (
              0.1,
              0.1,
              0.1,
          ),
      }
      __SCREAMING_SNAKE_CASE      = ShapERenderer(**__SCREAMING_SNAKE_CASE     )
      return model
 def  UpperCAmelCase__						(						self   :			str     )    ->					Tuple:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = self.dummy_prior
      __SCREAMING_SNAKE_CASE      = self.dummy_image_encoder
      __SCREAMING_SNAKE_CASE      = self.dummy_image_processor
      __SCREAMING_SNAKE_CASE      = self.dummy_renderer

      __SCREAMING_SNAKE_CASE      = HeunDiscreteScheduler(
          beta_schedule="""exp"""   ,				num_train_timesteps=1_024   ,				prediction_type="""sample"""   ,				use_karras_sigmas=__SCREAMING_SNAKE_CASE   ,				clip_sample=__SCREAMING_SNAKE_CASE   ,				clip_sample_range=1.0   ,				)
      __SCREAMING_SNAKE_CASE      = {
          """prior""": prior,
          """image_encoder""": image_encoder,
          """image_processor""": image_processor,
          """renderer""": renderer,
          """scheduler""": scheduler,
      }

      return components
 def  UpperCAmelCase__						(						self   :			Dict   ,				__SCREAMING_SNAKE_CASE   :			Tuple   ,				__SCREAMING_SNAKE_CASE   :			str=0     )    ->					Tuple:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = floats_tensor((1, 3, 64, 64)   ,				rng=random.Random(__SCREAMING_SNAKE_CASE     )     ).to(__SCREAMING_SNAKE_CASE     )

      if str(__SCREAMING_SNAKE_CASE     ).startswith("""mps"""     ):
           __SCREAMING_SNAKE_CASE      = torch.manual_seed(__SCREAMING_SNAKE_CASE     )
      else:
           __SCREAMING_SNAKE_CASE      = torch.Generator(device=__SCREAMING_SNAKE_CASE     ).manual_seed(__SCREAMING_SNAKE_CASE     )
      __SCREAMING_SNAKE_CASE      = {
          """image""": input_image,
          """generator""": generator,
          """num_inference_steps""": 1,
          """frame_size""": 32,
          """output_type""": """np""",
      }
      return inputs
 def  UpperCAmelCase__						(						self   :			Union[str, Any]     )    ->					int:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = """cpu"""

      __SCREAMING_SNAKE_CASE      = self.get_dummy_components()

      __SCREAMING_SNAKE_CASE      = self.pipeline_class(**__SCREAMING_SNAKE_CASE     )
      __SCREAMING_SNAKE_CASE      = pipe.to(__SCREAMING_SNAKE_CASE     )

      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

      __SCREAMING_SNAKE_CASE      = pipe(**self.get_dummy_inputs(__SCREAMING_SNAKE_CASE     )     )
      __SCREAMING_SNAKE_CASE      = output.images[0]
      __SCREAMING_SNAKE_CASE      = image[0, -3:, -3:, -1]

      assert image.shape == (20, 32, 32, 3)

      __SCREAMING_SNAKE_CASE      = np.array(
          [
              0.00039216,
              0.00039216,
              0.00039216,
              0.00039216,
              0.00039216,
              0.00039216,
              0.00039216,
              0.00039216,
              0.00039216,
          ]     )

      assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-2
 def  UpperCAmelCase__						(						self   :			str     )    ->					Dict:


      """simple docstring"""




      self._test_inference_batch_consistent(batch_sizes=[1, 2]     )
 def  UpperCAmelCase__						(						self   :			Dict     )    ->					str:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = torch_device == """cpu"""
      __SCREAMING_SNAKE_CASE      = True
      self._test_inference_batch_single_identical(
          batch_size=2   ,				test_max_difference=__SCREAMING_SNAKE_CASE   ,				relax_max_difference=__SCREAMING_SNAKE_CASE   ,				)


 def  UpperCAmelCase__						(						self   :			Tuple     )    ->					Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = self.get_dummy_components()
      __SCREAMING_SNAKE_CASE      = self.pipeline_class(**__SCREAMING_SNAKE_CASE     )
      __SCREAMING_SNAKE_CASE      = pipe.to(__SCREAMING_SNAKE_CASE     )
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

      __SCREAMING_SNAKE_CASE      = 1
      __SCREAMING_SNAKE_CASE      = 2

      __SCREAMING_SNAKE_CASE      = self.get_dummy_inputs(__SCREAMING_SNAKE_CASE     )

      for key in inputs.keys():
           if key in self.batch_params:
                __SCREAMING_SNAKE_CASE      = batch_size * [inputs[key]]

      __SCREAMING_SNAKE_CASE      = pipe(**__SCREAMING_SNAKE_CASE   ,				num_images_per_prompt=__SCREAMING_SNAKE_CASE     )[0]

      assert images.shape[0] == batch_size * num_images_per_prompt







@slow
@require_torch_gpu
class 		lowerCAmelCase__  ( unittest.TestCase   ):



 """simple docstring"""



 def  UpperCAmelCase__						(						self   :			int     )    ->					Dict:


      """simple docstring"""




      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()


 def  UpperCAmelCase__						(						self   :			Optional[Any]     )    ->					Dict:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/shap_e/corgi.png"""     )
      __SCREAMING_SNAKE_CASE      = load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/shap_e/test_shap_e_img2img_out.npy"""     )
      __SCREAMING_SNAKE_CASE      = ShapEImgaImgPipeline.from_pretrained("""openai/shap-e-img2img"""     )
      __SCREAMING_SNAKE_CASE      = pipe.to(__SCREAMING_SNAKE_CASE     )
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

      __SCREAMING_SNAKE_CASE      = torch.Generator(device=__SCREAMING_SNAKE_CASE     ).manual_seed(0     )

      __SCREAMING_SNAKE_CASE      = pipe(
          __SCREAMING_SNAKE_CASE   ,				generator=__SCREAMING_SNAKE_CASE   ,				guidance_scale=3.0   ,				num_inference_steps=64   ,				frame_size=64   ,				output_type="""np"""   ,				).images[0]

      assert images.shape == (20, 64, 64, 3)

      assert_mean_pixel_difference(__SCREAMING_SNAKE_CASE   ,				__SCREAMING_SNAKE_CASE     )






