



'''simple docstring'''





import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class     UpperCAmelCase_       (       _a		,		_a		,		unittest.TestCase	):
    """simple docstring"""







    lowercase					=		StableDiffusionXLImgaImgPipeline
    lowercase					=		TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
    lowercase					=		PipelineTesterMixin.required_optional_params - {"latents"}
    lowercase					=		TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    lowercase					=		IMAGE_TO_IMAGE_IMAGE_PARAMS
    lowercase					=		IMAGE_TO_IMAGE_IMAGE_PARAMS




    def lowerCamelCase       (    self					:     List[Any]						):
          torch.manual_seed(0						)
          snake_case__  :    List[str]    						=      UNetaDConditionModel(
              block_out_channels=(32, 64) ,     layers_per_block=2 ,     sample_size=32 ,     in_channels=4 ,     out_channels=4 ,     down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""") ,     up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""") ,     attention_head_dim=(2, 4) ,     use_linear_projection=snake_case_ ,     addition_embed_type="""text_time""" ,     addition_time_embed_dim=8 ,     transformer_layers_per_block=(1, 2) ,     projection_class_embeddings_input_dim=80 ,     cross_attention_dim=64 ,     )
          snake_case__  :    Optional[int]    						=      EulerDiscreteScheduler(
              beta_start=0.00085 ,     beta_end=0.012 ,     steps_offset=1 ,     beta_schedule="""scaled_linear""" ,     timestep_spacing="""leading""" ,     )
          torch.manual_seed(0						)
          snake_case__  :    Dict    						=      AutoencoderKL(
              block_out_channels=[32, 64] ,     in_channels=3 ,     out_channels=3 ,     down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""] ,     up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""] ,     latent_channels=4 ,     sample_size=128 ,     )
          torch.manual_seed(0						)
          snake_case__  :    int    						=      CLIPTextConfig(
              bos_token_id=0 ,     eos_token_id=2 ,     hidden_size=32 ,     intermediate_size=37 ,     layer_norm_eps=1E-0_5 ,     num_attention_heads=4 ,     num_hidden_layers=5 ,     pad_token_id=1 ,     vocab_size=1_000 ,     hidden_act="""gelu""" ,     projection_dim=32 ,     )
          snake_case__  :    Optional[Any]    						=      CLIPTextModel(snake_case_						)
          snake_case__  :    Optional[int]    						=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" ,     local_files_only=snake_case_						)

          snake_case__  :    Dict    						=      CLIPTextModelWithProjection(snake_case_						)
          snake_case__  :    List[Any]    						=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" ,     local_files_only=snake_case_						)

          snake_case__  :    List[str]    						=      {
              """unet""": unet,
              """scheduler""": scheduler,
              """vae""": vae,
              """text_encoder""": text_encoder,
              """tokenizer""": tokenizer,
              """text_encoder_2""": text_encoder_a,
              """tokenizer_2""": tokenizer_a,
              # "safety_checker": None,
              # "feature_extractor": None,
          }
          return components




    def lowerCamelCase       (    self					:     Dict ,     snake_case_					:     str ,     snake_case_					:     Tuple=0						):
          snake_case__  :    Optional[int]    						=      floats_tensor((1, 3, 32, 32) ,     rng=random.Random(snake_case_						)						).to(snake_case_						)
          snake_case__  :    Union[str, Any]    						=      image / 2 + 0.5
          if str(snake_case_						).startswith("""mps"""						):
                snake_case__  :    List[str]    						=      torch.manual_seed(snake_case_						)
          else:
                snake_case__  :    List[Any]    						=      torch.Generator(device=snake_case_						).manual_seed(snake_case_						)
          snake_case__  :    Optional[Any]    						=      {
              """prompt""": """A painting of a squirrel eating a burger""",
              """image""": image,
              """generator""": generator,
              """num_inference_steps""": 2,
              """guidance_scale""": 5.0,
              """output_type""": """numpy""",
              """strength""": 0.75,
          }
          return inputs




    def lowerCamelCase       (    self					:     Dict						):
          snake_case__  :    List[Any]    						=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
          snake_case__  :    Optional[Any]    						=      self.get_dummy_components()
          snake_case__  :    Any    						=      StableDiffusionXLImgaImgPipeline(**snake_case_						)
          snake_case__  :    Optional[int]    						=      sd_pipe.to(snake_case_						)
          sd_pipe.set_progress_bar_config(disable=snake_case_						)

          snake_case__  :    Optional[Any]    						=      self.get_dummy_inputs(snake_case_						)
          snake_case__  :    Any    						=      sd_pipe(**snake_case_						).images
          snake_case__  :    Optional[Any]    						=      image[0, -3:, -3:, -1]

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

          snake_case__  :    Dict    						=      np.array([0.4656, 0.4840, 0.4439, 0.6698, 0.5574, 0.4524, 0.5799, 0.5943, 0.5165]						)

          assert np.abs(image_slice.flatten() - expected_slice						).max() < 1E-2




    def lowerCamelCase       (    self					:     Union[str, Any]						):
          super().test_attention_slicing_forward_pass(expected_max_diff=3E-3						)




    def lowerCamelCase       (    self					:     Any						):
          super().test_inference_batch_single_identical(expected_max_diff=3E-3						)




    def lowerCamelCase       (    self					:     Any						):
          pass




    def lowerCamelCase       (    self					:     List[str]						):
          snake_case__  :    Union[str, Any]    						=      self.get_dummy_components()
          snake_case__  :    str    						=      StableDiffusionXLImgaImgPipeline(**snake_case_						)
          snake_case__  :    Optional[int]    						=      sd_pipe.to(snake_case_						)
          snake_case__  :    Optional[Any]    						=      sd_pipe.to(snake_case_						)
          sd_pipe.set_progress_bar_config(disable=snake_case_						)

          # forward without prompt embeds
          snake_case__  :    Union[str, Any]    						=      self.get_dummy_inputs(snake_case_						)
          snake_case__  :    List[Any]    						=      3 * ["""this is a negative prompt"""]
          snake_case__  :    Union[str, Any]    						=      negative_prompt
          snake_case__  :    Optional[Any]    						=      3 * [inputs["""prompt"""]]

          snake_case__  :    str    						=      sd_pipe(**snake_case_						)
          snake_case__  :    str    						=      output.images[0, -3:, -3:, -1]

          # forward with prompt embeds
          snake_case__  :    Union[str, Any]    						=      self.get_dummy_inputs(snake_case_						)
          snake_case__  :    Any    						=      3 * ["""this is a negative prompt"""]
          snake_case__  :    Any    						=      3 * [inputs.pop("""prompt"""						)]

          (
              (
              snake_case__
          )						,				(
              snake_case__
          )						,				(
              snake_case__
          )						,				(
              snake_case__
          )						,				
          )  :    Union[str, Any]    						=      sd_pipe.encode_prompt(snake_case_ ,     negative_prompt=snake_case_						)

          snake_case__  :    List[str]    						=      sd_pipe(
              **snake_case_ ,     prompt_embeds=snake_case_ ,     negative_prompt_embeds=snake_case_ ,     pooled_prompt_embeds=snake_case_ ,     negative_pooled_prompt_embeds=snake_case_ ,     )
          snake_case__  :    Optional[int]    						=      output.images[0, -3:, -3:, -1]

          # make sure that it's equal
          assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()						).max() < 1E-4







@slow
@require_torch_gpu
class     UpperCAmelCase_       (       unittest.TestCase	):
    """simple docstring"""







    def lowerCamelCase       (    self					:     Union[str, Any]						):
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()




    def lowerCamelCase       (    self					:     Optional[Any] ,     snake_case_					:     Optional[int] ,     snake_case_					:     Optional[int]="cpu" ,     snake_case_					:     Dict=torch.floataa ,     snake_case_					:     Optional[int]=0						):
          snake_case__  :    Tuple    						=      torch.Generator(device=snake_case_						).manual_seed(snake_case_						)
          snake_case__  :    List[Any]    						=      np.random.RandomState(snake_case_						).standard_normal((1, 4, 64, 64)						)
          snake_case__  :    str    						=      torch.from_numpy(snake_case_						).to(device=snake_case_ ,     dtype=snake_case_						)
          snake_case__  :    List[Any]    						=      {
              """prompt""": """a photograph of an astronaut riding a horse""",
              """latents""": latents,
              """generator""": generator,
              """num_inference_steps""": 3,
              """guidance_scale""": 7.5,
              """output_type""": """numpy""",
          }
          return inputs




    def lowerCamelCase       (    self					:     List[Any]						):
          snake_case__  :    Optional[Any]    						=      DiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-base"""						)
          pipe.to(snake_case_						)
          pipe.set_progress_bar_config(disable=snake_case_						)

          snake_case__  :    Optional[Any]    						=      self.get_inputs(snake_case_						)
          snake_case__  :    List[str]    						=      pipe(**snake_case_						).images
          snake_case__  :    Dict    						=      image[0, -3:, -3:, -1].flatten()

          assert image.shape == (1, 512, 512, 3)
          snake_case__  :    int    						=      np.array([0.49493, 0.47896, 0.40798, 0.54214, 0.53212, 0.48202, 0.47656, 0.46329, 0.48506]						)
          assert np.abs(image_slice - expected_slice						).max() < 7E-3








"""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					:				Union[str, Any]					    =						['''\nclass''', '''\ndef''', '''\n#''', '''\n@''', '''\nprint''', '''\nif''']






class      a						(					__snake_case		):






    def __init__(       self     :							Dict       ,						__SCREAMING_SNAKE_CASE     :							Any       ,						__SCREAMING_SNAKE_CASE     :							Optional[Any]       ,						__SCREAMING_SNAKE_CASE     :							Any=None       ,						__SCREAMING_SNAKE_CASE     :							Tuple=1					)     ->		str:
      lowerCamelCase_         =						tokenizer
      lowerCamelCase_         =						dataset
      lowerCamelCase_         =						len(__SCREAMING_SNAKE_CASE					) if n_tasks is None else n_tasks
      lowerCamelCase_         =						n_copies






    def __iter__(       self     :							Dict					)     ->		Any:
      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      a						(					__snake_case		):






    def __init__(       self     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							str       ,						__SCREAMING_SNAKE_CASE     :							int       ,						__SCREAMING_SNAKE_CASE     :							Any					)     ->		List[Any]:
      lowerCamelCase_         =						start_length
      lowerCamelCase_         =						eof_strings
      lowerCamelCase_         =						tokenizer






    def __call__(       self     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							Optional[int]       ,						__SCREAMING_SNAKE_CASE     :							Dict       ,						**__SCREAMING_SNAKE_CASE     :							Union[str, Any]					)     ->		Optional[Any]:
      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__       (							_lowerCamelCase       :		List[Any] )      ->				Tuple:
  lowerCamelCase_         =						re.split('(%s)' % '|'.join(_lowerCamelCase )		,				_lowerCamelCase )
  # last string should be ""
  return "".join(string_list[:-2] )







def      lowerCamelCase__       (							_lowerCamelCase       :		str		,				_lowerCamelCase       :		str		,				_lowerCamelCase       :		Optional[Any]		,				_lowerCamelCase       :		Optional[Any]		,				_lowerCamelCase       :		str		,				_lowerCamelCase       :		Dict=20		,				**_lowerCamelCase       :		Dict )      ->				List[str]:
  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:
  # Setup configuration
  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()




