
import unittest

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

from diffusers import DDIMScheduler, LDMPipeline, UNetaDModel, VQModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()

class     _UpperCamelCase			(   unittest.TestCase					):






			'''simple docstring'''

			@property
			def  lowerCAmelCase__		(    self	:    Dict ):
							torch.manual_seed(0 )
							UpperCamelCase_:							int		=  UNetaDModel(
							    block_out_channels=(32, 64)  ,		layers_per_block=2  ,		sample_size=32  ,		in_channels=3  ,		out_channels=3  ,		down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")  ,		up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")  ,		)
							return model

			@property
			def  lowerCAmelCase__		(    self	:    List[str] ):
							torch.manual_seed(0 )
							UpperCamelCase_:							int		=  VQModel(
							    block_out_channels=[32, 64]  ,		in_channels=3  ,		out_channels=3  ,		down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]  ,		up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]  ,		latent_channels=3  ,		)
							return model

			@property
			def  lowerCAmelCase__		(    self	:    Any ):
							torch.manual_seed(0 )
							UpperCamelCase_:							List[str]		=  CLIPTextConfig(
							    bos_token_id=0  ,		eos_token_id=2  ,		hidden_size=32  ,		intermediate_size=37  ,		layer_norm_eps=1e-05  ,		num_attention_heads=4  ,		num_hidden_layers=5  ,		pad_token_id=1  ,		vocab_size=1000  ,		)
							return CLIPTextModel(snake_case_ )




			def  lowerCAmelCase__		(    self	:    str ):
							UpperCamelCase_:							Optional[Any]		=  self.dummy_uncond_unet
							UpperCamelCase_:							Optional[Any]		=  DDIMScheduler()
							UpperCamelCase_:							Optional[Any]		=  self.dummy_vq_model

							UpperCamelCase_:							Optional[Any]		=  LDMPipeline(unet=snake_case_  ,		vqvae=snake_case_  ,		scheduler=snake_case_ )
							ldm.to(snake_case_ )
							ldm.set_progress_bar_config(disable=snake_case_ )

							UpperCamelCase_:							str		=  torch.manual_seed(0 )
							UpperCamelCase_:							List[Any]		=  ldm(generator=snake_case_  ,		num_inference_steps=2  ,		output_type="""numpy""" ).images

							UpperCamelCase_:							Optional[Any]		=  torch.manual_seed(0 )
							UpperCamelCase_:							List[Any]		=  ldm(generator=snake_case_  ,		num_inference_steps=2  ,		output_type="""numpy"""  ,		return_dict=snake_case_ )[0]

							UpperCamelCase_:							str		=  image[0, -3:, -3:, -1]
							UpperCamelCase_:							Tuple		=  image_from_tuple[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							UpperCamelCase_:							Union[str, Any]		=  np.array([0.8512, 0.818, 0.6411, 0.6808, 0.4465, 0.5618, 0.46, 0.6231, 0.5172] )
							UpperCamelCase_:							Optional[int]		=  1e-2 if torch_device != """mps""" else 3e-2

							assert np.abs(image_slice.flatten() - expected_slice ).max() < tolerance
							assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < tolerance




@slow
@require_torch
class     _UpperCamelCase			(   unittest.TestCase					):






			'''simple docstring'''

			def  lowerCAmelCase__		(    self	:    List[str] ):
							UpperCamelCase_:							str		=  LDMPipeline.from_pretrained("""CompVis/ldm-celebahq-256""" )
							ldm.to(snake_case_ )
							ldm.set_progress_bar_config(disable=snake_case_ )

							UpperCamelCase_:							Union[str, Any]		=  torch.manual_seed(0 )
							UpperCamelCase_:							Tuple		=  ldm(generator=snake_case_  ,		num_inference_steps=5  ,		output_type="""numpy""" ).images

							UpperCamelCase_:							str		=  image[0, -3:, -3:, -1]

							assert image.shape == (1, 256, 256, 3)
							UpperCamelCase_:							List[str]		=  np.array([0.4399, 0.4_4975, 0.4_6825, 0.474, 0.4359, 0.4581, 0.4_5095, 0.4341, 0.4447] )
							UpperCamelCase_:							List[Any]		=  1e-2 if torch_device != """mps""" else 3e-2

							assert np.abs(image_slice.flatten() - expected_slice ).max() < tolerance








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


lowerCamelCase_  :	str 				=					["""\nclass""", """\ndef""", """\n#""", """\n@""", """\nprint""", """\nif"""]

class     _UpperCamelCase			(   _A					):






			'''simple docstring'''

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




			def __iter__(    self	:    Tuple ):
							UpperCamelCase_:							Optional[int]		=  []
							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_:							List[str]		=  self.tokenizer(snake_case_  ,		padding=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     _UpperCamelCase			(   _A					):






			'''simple docstring'''

			def __init__(    self	:    Any  ,		snake_case_	:    List[str]  ,		snake_case_	:    Dict  ,		snake_case_	:    int ):
							UpperCamelCase_:							Union[str, Any]		=  start_length
							UpperCamelCase_:							Dict		=  eof_strings
							UpperCamelCase_:							List[str]		=  tokenizer




			def __call__(    self	:    Tuple  ,		snake_case_	:    List[str]  ,		snake_case_	:    Optional[Any]  ,		**snake_case_	:    int ):
							UpperCamelCase_:							Union[str, Any]		=  self.tokenizer.batch_decode(input_ids[:, self.start_length :] )
							UpperCamelCase_:							Dict		=  []
							for decoded_generation in decoded_generations:
											done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings ) )
							return all(snake_case_ )







def     A__   (  lowerCamelCase  )					->		Optional[int]:
				UpperCamelCase_:							str		=  re.split("""(%s)""" % """|""".join(lowerCamelCase  )       ,      lowerCamelCase  )
				# last string should be ""
				return "".join(string_list[:-2]  )







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

												UpperCamelCase_,				UpperCamelCase_:							Tuple		=  accelerator.gather((generated_tokens, generated_tasks)  )
												UpperCamelCase_:							Tuple		=  generated_tokens.cpu().numpy()
												UpperCamelCase_:							Any		=  generated_tasks.cpu().numpy()

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

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







def     A__   (  )					->		Union[str, Any]:
				# Setup configuration
				UpperCamelCase_:							Optional[Any]		=  HfArgumentParser(lowerCamelCase  )
				UpperCamelCase_:							str		=  parser.parse_args()

				transformers.logging.set_verbosity_error()
				# enables code execution in code_eval metric
				UpperCamelCase_:							List[str]		=  args.HF_ALLOW_CODE_EVAL
				# make sure tokenizer plays nice with multiprocessing
				UpperCamelCase_:							Union[str, Any]		=  """false"""

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

				# Use dataset load to feed to accelerate
				UpperCamelCase_:							List[Any]		=  Accelerator()
				set_seed(args.seed       ,      device_specific=lowerCamelCase  )

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

				# Generation settings
				UpperCamelCase_:							Union[str, Any]		=  {
				    """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
				UpperCamelCase_:							Any		=  load_dataset("""openai_humaneval"""  )
				UpperCamelCase_:							Union[str, Any]		=  load_metric("""code_eval"""  )

				UpperCamelCase_:							Union[str, Any]		=  args.num_tasks if args.num_tasks is not None else len(human_eval["""test"""]  )
				UpperCamelCase_:							List[Any]		=  args.n_samples // args.batch_size

				UpperCamelCase_:							Any		=  TokenizedDataset(lowerCamelCase       ,      human_eval["""test"""]       ,      n_copies=lowerCamelCase       ,      n_tasks=lowerCamelCase  )
				# do not confuse args.batch_size, which is actually the num_return_sequences
				UpperCamelCase_:							Optional[int]		=  DataLoader(lowerCamelCase       ,      batch_size=1  )

				# Run a quick test to see if code evaluation is enabled
				try:
								UpperCamelCase_:							List[str]		=  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_:							Dict		=  accelerator.prepare(lowerCamelCase       ,      lowerCamelCase  )

				UpperCamelCase_:							Union[str, Any]		=  complete_code(
				    lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase       ,      n_tasks=lowerCamelCase       ,      batch_size=args.batch_size       ,      **lowerCamelCase       ,      )

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

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

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






