




"""simple docstring"""







from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block

@dataclass
class 				snake_case_			(					lowerCamelCase_			):







   """simple docstring"""




   A_				      =       42

class 				snake_case_			(					nn.Module			):







   """simple docstring"""




   def __init__(		self       ,					lowerCamelCase_=3       ,					lowerCamelCase_=3       ,					lowerCamelCase_=("DownEncoderBlock2D",)       ,					lowerCamelCase_=(6_4,)       ,					lowerCamelCase_=2       ,					lowerCamelCase_=3_2       ,					lowerCamelCase_="silu"       ,					lowerCamelCase_=True       ,					)				->   Tuple:
       super().__init__()
       UpperCamelCase              =	layers_per_block

       UpperCamelCase              =	torch.nn.Convad(
           lowerCamelCase_       ,					block_out_channels[0]       ,					kernel_size=3       ,					stride=1       ,					padding=1       ,					)

       UpperCamelCase              =	None
       UpperCamelCase              =	nn.ModuleList([])

       # down
       UpperCamelCase              =	block_out_channels[0]
       for i, down_block_type in enumerate(lowerCamelCase_):
           UpperCamelCase              =	output_channel
           UpperCamelCase              =	block_out_channels[i]
           UpperCamelCase              =	i == len(lowerCamelCase_) - 1

           UpperCamelCase              =	get_down_block(
               lowerCamelCase_       ,					num_layers=self.layers_per_block       ,					in_channels=lowerCamelCase_       ,					out_channels=lowerCamelCase_       ,					add_downsample=not is_final_block       ,					resnet_eps=1e-6       ,					downsample_padding=0       ,					resnet_act_fn=lowerCamelCase_       ,					resnet_groups=lowerCamelCase_       ,					attention_head_dim=lowerCamelCase_       ,					temb_channels=lowerCamelCase_       ,					)
           self.down_blocks.append(lowerCamelCase_)

       # mid
       UpperCamelCase              =	UNetMidBlockaD(
           in_channels=block_out_channels[-1]       ,					resnet_eps=1e-6       ,					resnet_act_fn=lowerCamelCase_       ,					output_scale_factor=1       ,					resnet_time_scale_shift='''default'''       ,					attention_head_dim=block_out_channels[-1]       ,					resnet_groups=lowerCamelCase_       ,					temb_channels=lowerCamelCase_       ,					)

       # out
       UpperCamelCase              =	nn.GroupNorm(num_channels=block_out_channels[-1]       ,					num_groups=lowerCamelCase_       ,					eps=1e-6)
       UpperCamelCase              =	nn.SiLU()

       UpperCamelCase              =	2 * out_channels if double_z else out_channels
       UpperCamelCase              =	nn.Convad(block_out_channels[-1]       ,					lowerCamelCase_       ,					3       ,					padding=1)

       UpperCamelCase              =	False







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_)				->   Any:
       UpperCamelCase              =	x
       UpperCamelCase              =	self.conv_in(lowerCamelCase_)

       if self.training and self.gradient_checkpointing:

           def create_custom_forward(lowerCamelCase_):
               def custom_forward(*lowerCamelCase_):
                   return module(*lowerCamelCase_)

               return custom_forward

           # down
           if is_torch_version('''>='''       ,					'''1.11.0'''):
               for down_block in self.down_blocks:
                   UpperCamelCase              =	torch.utils.checkpoint.checkpoint(
                       create_custom_forward(lowerCamelCase_)       ,					lowerCamelCase_       ,					use_reentrant=lowerCamelCase_)
               # middle
               UpperCamelCase              =	torch.utils.checkpoint.checkpoint(
                   create_custom_forward(self.mid_block)       ,					lowerCamelCase_       ,					use_reentrant=lowerCamelCase_)
           else:
               for down_block in self.down_blocks:
                   UpperCamelCase              =	torch.utils.checkpoint.checkpoint(create_custom_forward(lowerCamelCase_)       ,					lowerCamelCase_)
               # middle
               UpperCamelCase              =	torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block)       ,					lowerCamelCase_)

       else:
           # down
           for down_block in self.down_blocks:
               UpperCamelCase              =	down_block(lowerCamelCase_)

           # middle
           UpperCamelCase              =	self.mid_block(lowerCamelCase_)

       # post-process
       UpperCamelCase              =	self.conv_norm_out(lowerCamelCase_)
       UpperCamelCase              =	self.conv_act(lowerCamelCase_)
       UpperCamelCase              =	self.conv_out(lowerCamelCase_)

       return sample

class 				snake_case_			(					nn.Module			):







   """simple docstring"""




   def __init__(		self       ,					lowerCamelCase_=3       ,					lowerCamelCase_=3       ,					lowerCamelCase_=("UpDecoderBlock2D",)       ,					lowerCamelCase_=(6_4,)       ,					lowerCamelCase_=2       ,					lowerCamelCase_=3_2       ,					lowerCamelCase_="silu"       ,					lowerCamelCase_="group"       ,					)				->   Optional[int]:
       super().__init__()
       UpperCamelCase              =	layers_per_block

       UpperCamelCase              =	nn.Convad(
           lowerCamelCase_       ,					block_out_channels[-1]       ,					kernel_size=3       ,					stride=1       ,					padding=1       ,					)

       UpperCamelCase              =	None
       UpperCamelCase              =	nn.ModuleList([])

       UpperCamelCase              =	in_channels if norm_type == '''spatial''' else None

       # mid
       UpperCamelCase              =	UNetMidBlockaD(
           in_channels=block_out_channels[-1]       ,					resnet_eps=1e-6       ,					resnet_act_fn=lowerCamelCase_       ,					output_scale_factor=1       ,					resnet_time_scale_shift='''default''' if norm_type == '''group''' else norm_type       ,					attention_head_dim=block_out_channels[-1]       ,					resnet_groups=lowerCamelCase_       ,					temb_channels=lowerCamelCase_       ,					)

       # up
       UpperCamelCase              =	list(reversed(lowerCamelCase_))
       UpperCamelCase              =	reversed_block_out_channels[0]
       for i, up_block_type in enumerate(lowerCamelCase_):
           UpperCamelCase              =	output_channel
           UpperCamelCase              =	reversed_block_out_channels[i]

           UpperCamelCase              =	i == len(lowerCamelCase_) - 1

           UpperCamelCase              =	get_up_block(
               lowerCamelCase_       ,					num_layers=self.layers_per_block + 1       ,					in_channels=lowerCamelCase_       ,					out_channels=lowerCamelCase_       ,					prev_output_channel=lowerCamelCase_       ,					add_upsample=not is_final_block       ,					resnet_eps=1e-6       ,					resnet_act_fn=lowerCamelCase_       ,					resnet_groups=lowerCamelCase_       ,					attention_head_dim=lowerCamelCase_       ,					temb_channels=lowerCamelCase_       ,					resnet_time_scale_shift=lowerCamelCase_       ,					)
           self.up_blocks.append(lowerCamelCase_)
           UpperCamelCase              =	output_channel

       # out
       if norm_type == "spatial":
           UpperCamelCase              =	SpatialNorm(block_out_channels[0]       ,					lowerCamelCase_)
       else:
           UpperCamelCase              =	nn.GroupNorm(num_channels=block_out_channels[0]       ,					num_groups=lowerCamelCase_       ,					eps=1e-6)
       UpperCamelCase              =	nn.SiLU()
       UpperCamelCase              =	nn.Convad(block_out_channels[0]       ,					lowerCamelCase_       ,					3       ,					padding=1)

       UpperCamelCase              =	False







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_=None)				->   str:
       UpperCamelCase              =	z
       UpperCamelCase              =	self.conv_in(lowerCamelCase_)

       UpperCamelCase              =	next(iter(self.up_blocks.parameters())).dtype
       if self.training and self.gradient_checkpointing:

           def create_custom_forward(lowerCamelCase_):
               def custom_forward(*lowerCamelCase_):
                   return module(*lowerCamelCase_)

               return custom_forward

           if is_torch_version('''>='''       ,					'''1.11.0'''):
               # middle
               UpperCamelCase              =	torch.utils.checkpoint.checkpoint(
                   create_custom_forward(self.mid_block)       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					use_reentrant=lowerCamelCase_)
               UpperCamelCase              =	sample.to(lowerCamelCase_)

               # up
               for up_block in self.up_blocks:
                   UpperCamelCase              =	torch.utils.checkpoint.checkpoint(
                       create_custom_forward(lowerCamelCase_)       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					use_reentrant=lowerCamelCase_)
           else:
               # middle
               UpperCamelCase              =	torch.utils.checkpoint.checkpoint(
                   create_custom_forward(self.mid_block)       ,					lowerCamelCase_       ,					lowerCamelCase_)
               UpperCamelCase              =	sample.to(lowerCamelCase_)

               # up
               for up_block in self.up_blocks:
                   UpperCamelCase              =	torch.utils.checkpoint.checkpoint(create_custom_forward(lowerCamelCase_)       ,					lowerCamelCase_       ,					lowerCamelCase_)
       else:
           # middle
           UpperCamelCase              =	self.mid_block(lowerCamelCase_       ,					lowerCamelCase_)
           UpperCamelCase              =	sample.to(lowerCamelCase_)

           # up
           for up_block in self.up_blocks:
               UpperCamelCase              =	up_block(lowerCamelCase_       ,					lowerCamelCase_)

        # post-process
       if latent_embeds is None:
           UpperCamelCase              =	self.conv_norm_out(lowerCamelCase_)
       else:
           UpperCamelCase              =	self.conv_norm_out(lowerCamelCase_       ,					lowerCamelCase_)
       UpperCamelCase              =	self.conv_act(lowerCamelCase_)
       UpperCamelCase              =	self.conv_out(lowerCamelCase_)

       return sample

class 				snake_case_			(					nn.Module			):







   """simple docstring"""




   def __init__(		self       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_=None       ,					lowerCamelCase_="random"       ,					lowerCamelCase_=False       ,					lowerCamelCase_=True)				->   Any:
       super().__init__()
       UpperCamelCase              =	n_e
       UpperCamelCase              =	vq_embed_dim
       UpperCamelCase              =	beta
       UpperCamelCase              =	legacy

       UpperCamelCase              =	nn.Embedding(self.n_e       ,					self.vq_embed_dim)
       self.embedding.weight.data.uniform_(-1.0 / self.n_e       ,					1.0 / self.n_e)

       UpperCamelCase              =	remap
       if self.remap is not None:
           self.register_buffer('''used'''       ,					torch.tensor(np.load(self.remap)))
           UpperCamelCase              =	self.used.shape[0]
           UpperCamelCase              =	unknown_index  # "random" or "extra" or integer
           if self.unknown_index == "extra":
               UpperCamelCase              =	self.re_embed
               UpperCamelCase              =	self.re_embed + 1
           print(
               F'Remapping {self.n_e} indices to {self.re_embed} indices. '
               F'Using {self.unknown_index} for unknown indices.')
       else:
           UpperCamelCase              =	n_e

       UpperCamelCase              =	sane_index_shape







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_)				->   List[str]:
       UpperCamelCase              =	inds.shape
       assert len(lowerCamelCase_) > 1
       UpperCamelCase              =	inds.reshape(ishape[0]       ,					-1)
       UpperCamelCase              =	self.used.to(lowerCamelCase_)
       UpperCamelCase              =	(inds[:, :, None] == used[None, None, ...]).long()
       UpperCamelCase              =	match.argmax(-1)
       UpperCamelCase              =	match.sum(2) < 1
       if self.unknown_index == "random":
           UpperCamelCase              =	torch.randint(0       ,					self.re_embed       ,					size=new[unknown].shape).to(device=new.device)
       else:
           UpperCamelCase              =	self.unknown_index
       return new.reshape(lowerCamelCase_)







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_)				->   int:
       UpperCamelCase              =	inds.shape
       assert len(lowerCamelCase_) > 1
       UpperCamelCase              =	inds.reshape(ishape[0]       ,					-1)
       UpperCamelCase              =	self.used.to(lowerCamelCase_)
       if self.re_embed > self.used.shape[0]:  # extra token
           UpperCamelCase              =	0  # simply set to zero
       UpperCamelCase              =	torch.gather(used[None, :][inds.shape[0] * [0], :]       ,					1       ,					lowerCamelCase_)
       return back.reshape(lowerCamelCase_)







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_)				->   List[str]:
       # reshape z -> (batch, height, width, channel) and flatten
       UpperCamelCase              =	z.permute(0       ,					2       ,					3       ,					1).contiguous()
       UpperCamelCase              =	z.view(-1       ,					self.vq_embed_dim)

       # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
       UpperCamelCase              =	torch.argmin(torch.cdist(lowerCamelCase_       ,					self.embedding.weight)       ,					dim=1)

       UpperCamelCase              =	self.embedding(lowerCamelCase_).view(z.shape)
       UpperCamelCase              =	None
       UpperCamelCase              =	None

       # compute loss for embedding
       if not self.legacy:
           UpperCamelCase              =	self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean((z_q - z.detach()) ** 2)
       else:
           UpperCamelCase              =	torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean((z_q - z.detach()) ** 2)

       # preserve gradients
       UpperCamelCase              =	z + (z_q - z).detach()

       # reshape back to match original input shape
       UpperCamelCase              =	z_q.permute(0       ,					3       ,					1       ,					2).contiguous()

       if self.remap is not None:
           UpperCamelCase              =	min_encoding_indices.reshape(z.shape[0]       ,					-1)  # add batch axis
           UpperCamelCase              =	self.remap_to_used(lowerCamelCase_)
           UpperCamelCase              =	min_encoding_indices.reshape(-1       ,					1)  # flatten

       if self.sane_index_shape:
           UpperCamelCase              =	min_encoding_indices.reshape(z_q.shape[0]       ,					z_q.shape[2]       ,					z_q.shape[3])

       return z_q, loss, (perplexity, min_encodings, min_encoding_indices)







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_)				->   List[str]:
       # shape specifying (batch, height, width, channel)
       if self.remap is not None:
           UpperCamelCase              =	indices.reshape(shape[0]       ,					-1)  # add batch axis
           UpperCamelCase              =	self.unmap_to_all(lowerCamelCase_)
           UpperCamelCase              =	indices.reshape(-1)  # flatten again

       # get quantized latent vectors
       UpperCamelCase              =	self.embedding(lowerCamelCase_)

       if shape is not None:
           UpperCamelCase              =	z_q.view(lowerCamelCase_)
           # reshape back to match original input shape
           UpperCamelCase              =	z_q.permute(0       ,					3       ,					1       ,					2).contiguous()

       return z_q

class 				snake_case_			(					lowerCamelCase_			):







   """simple docstring"""




   def __init__(		self       ,					lowerCamelCase_       ,					lowerCamelCase_=False)				->   str:
       UpperCamelCase              =	parameters
       UpperCamelCase       ,	UpperCamelCase              =	torch.chunk(lowerCamelCase_       ,					2       ,					dim=1)
       UpperCamelCase              =	torch.clamp(self.logvar       ,					-30.0       ,					20.0)
       UpperCamelCase              =	deterministic
       UpperCamelCase              =	torch.exp(0.5 * self.logvar)
       UpperCamelCase              =	torch.exp(self.logvar)
       if self.deterministic:
           UpperCamelCase              =	UpperCamelCase              =	torch.zeros_like(
               self.mean       ,					device=self.parameters.device       ,					dtype=self.parameters.dtype)







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_ = None)				->   torch.FloatTensor:
       # make sure sample is on the same device as the parameters and has same dtype
       UpperCamelCase              =	randn_tensor(
           self.mean.shape       ,					generator=lowerCamelCase_       ,					device=self.parameters.device       ,					dtype=self.parameters.dtype)
       UpperCamelCase              =	self.mean + self.std * sample
       return x







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_=None)				->   int:
       if self.deterministic:
           return torch.Tensor([0.0])
       else:
           if other is None:
               return 0.5 * torch.sum(torch.pow(self.mean       ,					2) + self.var - 1.0 - self.logvar       ,					dim=[1, 2, 3])
           else:
               return 0.5 * torch.sum(
                   torch.pow(self.mean - other.mean       ,					2) / other.var
                   + self.var / other.var
                   - 1.0
                   - self.logvar
                   + other.logvar       ,					dim=[1, 2, 3]       ,					)







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_=[1, 2, 3])				->   Optional[int]:
       if self.deterministic:
           return torch.Tensor([0.0])
       UpperCamelCase              =	np.log(2.0 * np.pi)
       return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean       ,					2) / self.var       ,					dim=lowerCamelCase_)







   def   UpperCAmelCase__			(		self)				->   str:
       return self.mean
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		:     Optional[int]         =	["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif"]

class        UpperCamelCase						(				__a				):





						def __init__(self			,     __UpperCamelCase			,     __UpperCamelCase			,     __UpperCamelCase=None			,     __UpperCamelCase=1						)				->				List[str]:
												UpperCamelCase_		:							Dict  =		tokenizer
												UpperCamelCase_		:							Dict  =		dataset
												UpperCamelCase_		:							str  =		len(__UpperCamelCase						) if n_tasks is None else n_tasks
												UpperCamelCase_		:							int  =		n_copies





						def __iter__(self						)				->				Any:
												UpperCamelCase_		:							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_		:							Any  =		self.tokenizer(__UpperCamelCase			,     padding=__UpperCamelCase			,     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				):





						def __init__(self			,     __UpperCamelCase			,     __UpperCamelCase			,     __UpperCamelCase						)				->				Dict:
												UpperCamelCase_		:							List[Any]  =		start_length
												UpperCamelCase_		:							List[str]  =		eof_strings
												UpperCamelCase_		:							str  =		tokenizer





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






def     lowerCAmelCase_			(   _SCREAMING_SNAKE_CASE     :				Dict			):
						UpperCamelCase_		:							int  =		re.split("""(%s)""" % """|""".join(_SCREAMING_SNAKE_CASE			)		, _SCREAMING_SNAKE_CASE			)
						# last string should be ""
						return "".join(string_list[:-2]			)






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

																		UpperCamelCase_,UpperCamelCase_		:							Any  =		accelerator.gather((generated_tokens, generated_tasks)			)
																		UpperCamelCase_		:							str  =		generated_tokens.cpu().numpy()
																		UpperCamelCase_		:							List[Any]  =		generated_tasks.cpu().numpy()

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

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






def     lowerCAmelCase_			(   ):
						# Setup configuration
						UpperCamelCase_		:							List[Any]  =		HfArgumentParser(_SCREAMING_SNAKE_CASE			)
						UpperCamelCase_		:							Optional[Any]  =		parser.parse_args()

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

						if args.num_workers is None:
												UpperCamelCase_		:							List[str]  =		multiprocessing.cpu_count()

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

						# Load model and tokenizer
						UpperCamelCase_		:							str  =		AutoTokenizer.from_pretrained(args.model_ckpt			)
						UpperCamelCase_		:							Any  =		tokenizer.eos_token
						UpperCamelCase_		:							Tuple  =		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		, _SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE			)]			),
						}

						# Load evaluation dataset and metric
						UpperCamelCase_		:							Optional[int]  =		load_dataset("""openai_humaneval"""			)
						UpperCamelCase_		:							Any  =		load_metric("""code_eval"""			)

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

						UpperCamelCase_		:							Tuple  =		TokenizedDataset(_SCREAMING_SNAKE_CASE		, human_eval["""test"""]		, n_copies=_SCREAMING_SNAKE_CASE		, n_tasks=_SCREAMING_SNAKE_CASE			)
						# do not confuse args.batch_size, which is actually the num_return_sequences
						UpperCamelCase_		:							List[str]  =		DataLoader(_SCREAMING_SNAKE_CASE		, batch_size=1			)

						# Run a quick test to see if code evaluation is enabled
						try:
												UpperCamelCase_		:							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_		:							int  =		accelerator.prepare(_SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE			)

						UpperCamelCase_		:							Optional[int]  =		complete_code(
						    _SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE		, n_tasks=_SCREAMING_SNAKE_CASE		, batch_size=args.batch_size		, **_SCREAMING_SNAKE_CASE		, )

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

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

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


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






