


'''simple docstring'''


from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)






@flax.struct.dataclass
class 	lowerCAmelCase__			(			UpperCAmelCase__	):
			lowerCAmelCase :	jnp.ndarray
			lowerCAmelCase :	jnp.ndarray






class 	lowerCAmelCase__			(			nn.Module	):
			lowerCAmelCase :	int
			lowerCAmelCase :	Tuple[int]          =     (16, 32, 96, 256)
			lowerCAmelCase :	jnp.dtype          =     jnp.floataa




			def 							lowerCAmelCase__    (		self  :     Tuple )   ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase		:       Optional[Any]								=      nn.Conv(
					    self.block_out_channels[0]	,  kernel_size=(3, 3)	,  padding=((1, 1), (1, 1))	,  dtype=self.dtype	,  )

					_UpperCAmelCase		:       Tuple								=      []
					for i in range(len(self.block_out_channels ) - 1 ):
							_UpperCAmelCase		:       Union[str, Any]								=      self.block_out_channels[i]
							_UpperCAmelCase		:       List[Any]								=      self.block_out_channels[i + 1]
							_UpperCAmelCase		:       Optional[int]								=      nn.Conv(
							    lowerCamelCase__	,  kernel_size=(3, 3)	,  padding=((1, 1), (1, 1))	,  dtype=self.dtype	,  )
							blocks.append(lowerCamelCase__ )
							_UpperCAmelCase		:       Tuple								=      nn.Conv(
							    lowerCamelCase__	,  kernel_size=(3, 3)	,  strides=(2, 2)	,  padding=((1, 1), (1, 1))	,  dtype=self.dtype	,  )
							blocks.append(lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[int]								=      blocks

					_UpperCAmelCase		:       Optional[Any]								=      nn.Conv(
					    self.conditioning_embedding_channels	,  kernel_size=(3, 3)	,  padding=((1, 1), (1, 1))	,  kernel_init=nn.initializers.zeros_init()	,  bias_init=nn.initializers.zeros_init()	,  dtype=self.dtype	,  )






			def __call__(		self  :     Optional[Any]	,  lowerCamelCase__  :     Tuple )   ->int:



					'''simple docstring'''







					_UpperCAmelCase		:       Union[str, Any]								=      self.conv_in(lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[Any]								=      nn.silu(lowerCamelCase__ )

					for block in self.blocks:
							_UpperCAmelCase		:       str								=      block(lowerCamelCase__ )
							_UpperCAmelCase		:       Union[str, Any]								=      nn.silu(lowerCamelCase__ )

					_UpperCAmelCase		:       Union[str, Any]								=      self.conv_out(lowerCamelCase__ )

					return embedding






@flax_register_to_config
class 	lowerCAmelCase__			(			nn.Module    ,       UpperCAmelCase__    ,       UpperCAmelCase__	):
			lowerCAmelCase :	int          =     32
			lowerCAmelCase :	int          =     4
			lowerCAmelCase :	Tuple[str]          =     (
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "DownBlock2D",
			)
			lowerCAmelCase :	Union[bool, Tuple[bool]]          =     False
			lowerCAmelCase :	Tuple[int]          =     (320, 640, 1_280, 1_280)
			lowerCAmelCase :	int          =     2
			lowerCAmelCase :	Union[int, Tuple[int]]          =     8
			lowerCAmelCase :	Optional[Union[int, Tuple[int]]]          =     None
			lowerCAmelCase :	int          =     1_280
			lowerCAmelCase :	float          =     0.0
			lowerCAmelCase :	bool          =     False
			lowerCAmelCase :	jnp.dtype          =     jnp.floataa
			lowerCAmelCase :	bool          =     True
			lowerCAmelCase :	int          =     0
			lowerCAmelCase :	str          =     "rgb"
			lowerCAmelCase :	Tuple[int]          =     (16, 32, 96, 256)




			def 							lowerCAmelCase__    (		self  :     Tuple	,  lowerCamelCase__  :     jax.random.KeyArray )   ->FrozenDict:



					'''simple docstring'''







					_UpperCAmelCase		:       str								=      (1, self.in_channels, self.sample_size, self.sample_size)
					_UpperCAmelCase		:       Optional[Any]								=      jnp.zeros(lowerCamelCase__	,  dtype=jnp.floataa )
					_UpperCAmelCase		:       Tuple								=      jnp.ones((1,)	,  dtype=jnp.intaa )
					_UpperCAmelCase		:       Any								=      jnp.zeros((1, 1, self.cross_attention_dim)	,  dtype=jnp.floataa )
					_UpperCAmelCase		:       Optional[Any]								=      (1, 3, self.sample_size * 8, self.sample_size * 8)
					_UpperCAmelCase		:       int								=      jnp.zeros(lowerCamelCase__	,  dtype=jnp.floataa )

					_UpperCAmelCase					,    _UpperCAmelCase		:       List[Any]								=      jax.random.split(lowerCamelCase__ )
					_UpperCAmelCase		:       Union[str, Any]								=      {"params": params_rng, "dropout": dropout_rng}

					return self.init(lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )["params"]




			def 							lowerCAmelCase__    (		self  :     Dict )   ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       str								=      self.block_out_channels
					_UpperCAmelCase		:       List[Any]								=      block_out_channels[0] * 4

					# If `num_attention_heads` is not defined (which is the case for most models)
					# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
					# The reason for this behavior is to correct for incorrectly named variables that were introduced
					# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
					# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
					# which is why we correct for the naming here.
					_UpperCAmelCase		:       Union[str, Any]								=      self.num_attention_heads or self.attention_head_dim

					# input
					_UpperCAmelCase		:       Optional[Any]								=      nn.Conv(
					    block_out_channels[0]	,  kernel_size=(3, 3)	,  strides=(1, 1)	,  padding=((1, 1), (1, 1))	,  dtype=self.dtype	,  )

					# time
					_UpperCAmelCase		:       Any								=      FlaxTimesteps(
					    block_out_channels[0]	,  flip_sin_to_cos=self.flip_sin_to_cos	,  freq_shift=self.config.freq_shift )
					_UpperCAmelCase		:       Optional[Any]								=      FlaxTimestepEmbedding(lowerCamelCase__	,  dtype=self.dtype )

					_UpperCAmelCase		:       Union[str, Any]								=      FlaxControlNetConditioningEmbedding(
					    conditioning_embedding_channels=block_out_channels[0]	,  block_out_channels=self.conditioning_embedding_out_channels	,  )

					_UpperCAmelCase		:       int								=      self.only_cross_attention
					if isinstance(lowerCamelCase__	,  lowerCamelCase__ ):
							_UpperCAmelCase		:       str								=      (only_cross_attention,) * len(self.down_block_types )

					if isinstance(lowerCamelCase__	,  lowerCamelCase__ ):
							_UpperCAmelCase		:       Optional[int]								=      (num_attention_heads,) * len(self.down_block_types )

					# down
					_UpperCAmelCase		:       Dict								=      []
					_UpperCAmelCase		:       List[str]								=      []

					_UpperCAmelCase		:       Dict								=      block_out_channels[0]

					_UpperCAmelCase		:       List[Any]								=      nn.Conv(
					    lowerCamelCase__	,  kernel_size=(1, 1)	,  padding="VALID"	,  kernel_init=nn.initializers.zeros_init()	,  bias_init=nn.initializers.zeros_init()	,  dtype=self.dtype	,  )
					controlnet_down_blocks.append(lowerCamelCase__ )

					for i, down_block_type in enumerate(self.down_block_types ):
							_UpperCAmelCase		:       Dict								=      output_channel
							_UpperCAmelCase		:       Optional[Any]								=      block_out_channels[i]
							_UpperCAmelCase		:       int								=      i == len(lowerCamelCase__ ) - 1

							if down_block_type == "CrossAttnDownBlock2D":
									_UpperCAmelCase		:       Union[str, Any]								=      FlaxCrossAttnDownBlockaD(
									    in_channels=lowerCamelCase__	,  out_channels=lowerCamelCase__	,  dropout=self.dropout	,  num_layers=self.layers_per_block	,  num_attention_heads=num_attention_heads[i]	,  add_downsample=not is_final_block	,  use_linear_projection=self.use_linear_projection	,  only_cross_attention=only_cross_attention[i]	,  dtype=self.dtype	,  )
							else:
									_UpperCAmelCase		:       Union[str, Any]								=      FlaxDownBlockaD(
									    in_channels=lowerCamelCase__	,  out_channels=lowerCamelCase__	,  dropout=self.dropout	,  num_layers=self.layers_per_block	,  add_downsample=not is_final_block	,  dtype=self.dtype	,  )

							down_blocks.append(lowerCamelCase__ )

							for _ in range(self.layers_per_block ):
									_UpperCAmelCase		:       Dict								=      nn.Conv(
									    lowerCamelCase__	,  kernel_size=(1, 1)	,  padding="VALID"	,  kernel_init=nn.initializers.zeros_init()	,  bias_init=nn.initializers.zeros_init()	,  dtype=self.dtype	,  )
									controlnet_down_blocks.append(lowerCamelCase__ )

							if not is_final_block:
									_UpperCAmelCase		:       List[Any]								=      nn.Conv(
									    lowerCamelCase__	,  kernel_size=(1, 1)	,  padding="VALID"	,  kernel_init=nn.initializers.zeros_init()	,  bias_init=nn.initializers.zeros_init()	,  dtype=self.dtype	,  )
									controlnet_down_blocks.append(lowerCamelCase__ )

					_UpperCAmelCase		:       List[str]								=      down_blocks
					_UpperCAmelCase		:       List[Any]								=      controlnet_down_blocks

					# mid
					_UpperCAmelCase		:       Tuple								=      block_out_channels[-1]
					_UpperCAmelCase		:       str								=      FlaxUNetMidBlockaDCrossAttn(
					    in_channels=lowerCamelCase__	,  dropout=self.dropout	,  num_attention_heads=num_attention_heads[-1]	,  use_linear_projection=self.use_linear_projection	,  dtype=self.dtype	,  )

					_UpperCAmelCase		:       Optional[int]								=      nn.Conv(
					    lowerCamelCase__	,  kernel_size=(1, 1)	,  padding="VALID"	,  kernel_init=nn.initializers.zeros_init()	,  bias_init=nn.initializers.zeros_init()	,  dtype=self.dtype	,  )






			def __call__(		self  :     Union[str, Any]	,  lowerCamelCase__  :     List[str]	,  lowerCamelCase__  :     List[Any]	,  lowerCamelCase__  :     List[Any]	,  lowerCamelCase__  :     int	,  lowerCamelCase__  :     float = 1.0	,  lowerCamelCase__  :     bool = True	,  lowerCamelCase__  :     bool = False	,  )   ->Union[FlaxControlNetOutput, Tuple]:



					'''simple docstring'''







					_UpperCAmelCase		:       List[Any]								=      self.controlnet_conditioning_channel_order
					if channel_order == "bgr":
							_UpperCAmelCase		:       Optional[Any]								=      jnp.flip(lowerCamelCase__	,  axis=1 )

					# 1. time
					if not isinstance(lowerCamelCase__	,  jnp.ndarray ):
							_UpperCAmelCase		:       Union[str, Any]								=      jnp.array([timesteps]	,  dtype=jnp.intaa )
					elif isinstance(lowerCamelCase__	,  jnp.ndarray ) and len(timesteps.shape ) == 0:
							_UpperCAmelCase		:       List[str]								=      timesteps.astype(dtype=jnp.floataa )
							_UpperCAmelCase		:       Optional[int]								=      jnp.expand_dims(lowerCamelCase__	,  0 )

					_UpperCAmelCase		:       Any								=      self.time_proj(lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[Any]								=      self.time_embedding(lowerCamelCase__ )

					# 2. pre-process
					_UpperCAmelCase		:       List[Any]								=      jnp.transpose(lowerCamelCase__	,  (0, 2, 3, 1) )
					_UpperCAmelCase		:       Dict								=      self.conv_in(lowerCamelCase__ )

					_UpperCAmelCase		:       List[Any]								=      jnp.transpose(lowerCamelCase__	,  (0, 2, 3, 1) )
					_UpperCAmelCase		:       Optional[Any]								=      self.controlnet_cond_embedding(lowerCamelCase__ )
					sample += controlnet_cond

					# 3. down
					_UpperCAmelCase		:       int								=      (sample,)
					for down_block in self.down_blocks:
							if isinstance(lowerCamelCase__	,  lowerCamelCase__ ):
									_UpperCAmelCase					,    _UpperCAmelCase		:       str								=      down_block(lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  deterministic=not train )
							else:
									_UpperCAmelCase					,    _UpperCAmelCase		:       List[str]								=      down_block(lowerCamelCase__	,  lowerCamelCase__	,  deterministic=not train )
							down_block_res_samples += res_samples

					# 4. mid
					_UpperCAmelCase		:       Optional[int]								=      self.mid_block(lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  deterministic=not train )

					# 5. contronet blocks
					_UpperCAmelCase		:       List[str]								=      ()
					for down_block_res_sample, controlnet_block in zip(lowerCamelCase__	,  self.controlnet_down_blocks ):
							_UpperCAmelCase		:       Dict								=      controlnet_block(lowerCamelCase__ )
							controlnet_down_block_res_samples += (down_block_res_sample,)

					_UpperCAmelCase		:       Dict								=      controlnet_down_block_res_samples

					_UpperCAmelCase		:       Dict								=      self.controlnet_mid_block(lowerCamelCase__ )

					# 6. scaling
					_UpperCAmelCase		:       int								=      [sample * conditioning_scale for sample in down_block_res_samples]
					mid_block_res_sample *= conditioning_scale

					if not return_dict:
							return (down_block_res_samples, mid_block_res_sample)

					return FlaxControlNetOutput(
					    down_block_res_samples=lowerCamelCase__	,  mid_block_res_sample=lowerCamelCase__ )







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


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






class 	lowerCAmelCase__			(			UpperCAmelCase__	):




			def __init__(		self  :     Tuple	,  lowerCamelCase__  :     List[Any]	,  lowerCamelCase__  :     Optional[int]	,  lowerCamelCase__  :     Optional[Any]=None	,  lowerCamelCase__  :     int=1 )   ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase		:       List[Any]								=      tokenizer
					_UpperCAmelCase		:       Tuple								=      dataset
					_UpperCAmelCase		:       Union[str, Any]								=      len(lowerCamelCase__ ) if n_tasks is None else n_tasks
					_UpperCAmelCase		:       Any								=      n_copies






			def __iter__(		self  :     Any )   ->Dict:



					'''simple docstring'''







					_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		:       Optional[Any]								=      self.tokenizer(lowerCamelCase__	,  padding=lowerCamelCase__	,  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__			(			UpperCAmelCase__	):




			def __init__(		self  :     Optional[int]	,  lowerCamelCase__  :     Tuple	,  lowerCamelCase__  :     Any	,  lowerCamelCase__  :     Dict )   ->Any:



					'''simple docstring'''







					_UpperCAmelCase		:       List[str]								=      start_length
					_UpperCAmelCase		:       Union[str, Any]								=      eof_strings
					_UpperCAmelCase		:       Union[str, Any]								=      tokenizer






			def __call__(		self  :     Tuple	,  lowerCamelCase__  :     Union[str, Any]	,  lowerCamelCase__  :     List[str]	,  **lowerCamelCase__  :     Optional[int] )   ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       Tuple								=      self.tokenizer.batch_decode(input_ids[:, self.start_length :] )
					_UpperCAmelCase		:       Optional[int]								=      []
					for decoded_generation in decoded_generations:
							done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings ) )
					return all(lowerCamelCase__ )






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






def  __lowerCAmelCase (__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase=20      , **__lowerCAmelCase   ):
		_UpperCAmelCase		:       Tuple								=      defaultdict(__lowerCAmelCase   )  # dict of list of generated tokens
		for step, batch in tqdm(enumerate(__lowerCAmelCase   )   ):
				with torch.no_grad():
						_UpperCAmelCase		:       Tuple								=      batch["ids"].shape[-1]
						_UpperCAmelCase		:       Optional[int]								=      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		:       str								=      batch["task_id"].repeat(__lowerCAmelCase   )
						_UpperCAmelCase		:       str								=      accelerator.pad_across_processes(
						    __lowerCAmelCase      , dim=1      , pad_index=tokenizer.pad_token_id   )

						_UpperCAmelCase					,    _UpperCAmelCase		:       int								=      accelerator.gather((generated_tokens, generated_tasks)   )
						_UpperCAmelCase		:       Dict								=      generated_tokens.cpu().numpy()
						_UpperCAmelCase		:       Dict								=      generated_tasks.cpu().numpy()

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

		_UpperCAmelCase		:       int								=      [[] for _ in range(__lowerCAmelCase   )]
		for task, generated_tokens in gen_token_dict.items():
				for s in generated_tokens:
						_UpperCAmelCase		:       List[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  __lowerCAmelCase ():
		# Setup configuration
		_UpperCAmelCase		:       List[str]								=      HfArgumentParser(__lowerCAmelCase   )
		_UpperCAmelCase		:       Tuple								=      parser.parse_args()

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

		if args.num_workers is None:
				_UpperCAmelCase		:       List[str]								=      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		:       Tuple								=      AutoTokenizer.from_pretrained(args.model_ckpt   )
		_UpperCAmelCase		:       List[str]								=      tokenizer.eos_token
		_UpperCAmelCase		:       str								=      AutoModelForCausalLM.from_pretrained(args.model_ckpt   )

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

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

		_UpperCAmelCase		:       Tuple								=      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		:       List[str]								=      DataLoader(__lowerCAmelCase      , batch_size=1   )

		# Run a quick test to see if code evaluation is enabled
		try:
				_UpperCAmelCase		:       Optional[int]								=      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		:       Any								=      accelerator.prepare(__lowerCAmelCase      , __lowerCAmelCase   )

		_UpperCAmelCase		:       Dict								=      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		:       str								=      human_eval["test"][task]["test"]
						_UpperCAmelCase		:       Union[str, Any]								=      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()



