


'''simple docstring'''


import inspect
import logging
import os
import random
import shutil
import tempfile
import unittest

import pytest
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_cuda
from accelerate.utils import ProjectConfiguration, set_seed


lowerCamelCase__    	=     logging.getLogger(__name__)






def  __lowerCAmelCase (__lowerCAmelCase=2      , __lowerCAmelCase=3      , __lowerCAmelCase=16      , __lowerCAmelCase = 10      , __lowerCAmelCase = 2   ):

		def get_dataset(__lowerCAmelCase   ):
				_UpperCAmelCase		:       Tuple								=      torch.randn(batch_size * n_batches      , 1   )
				return TensorDataset(__lowerCAmelCase      , a * x + b + 0.1 * torch.randn(batch_size * n_batches      , 1   )   )

		_UpperCAmelCase		:       Dict								=      get_dataset(__lowerCAmelCase   )
		_UpperCAmelCase		:       Optional[int]								=      get_dataset(__lowerCAmelCase   )
		_UpperCAmelCase		:       Any								=      DataLoader(__lowerCAmelCase      , shuffle=__lowerCAmelCase      , batch_size=__lowerCAmelCase      , num_workers=4   )
		_UpperCAmelCase		:       str								=      DataLoader(__lowerCAmelCase      , shuffle=__lowerCAmelCase      , batch_size=__lowerCAmelCase      , num_workers=4   )
		return (train_dataloader, valid_dataloader)






def  __lowerCAmelCase (__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase=None   ):
		_UpperCAmelCase		:       Optional[int]								=      []
		for epoch in range(__lowerCAmelCase   ):
				# Train quickly
				model.train()
				for batch in dataloader:
						_UpperCAmelCase					,    _UpperCAmelCase		:       Optional[int]								=      batch
						_UpperCAmelCase		:       Optional[int]								=      model(__lowerCAmelCase   )
						_UpperCAmelCase		:       Optional[int]								=      torch.nn.functional.mse_loss(__lowerCAmelCase      , __lowerCAmelCase   )
						accelerator.backward(__lowerCAmelCase   )
						optimizer.step()
						optimizer.zero_grad()
				rands.append(random.random()   )  # Introduce some randomness
				if scheduler is not None:
						scheduler.step()
		return rands






class 	lowerCAmelCase__			(			nn.Module	):




			def __init__(		self  :     Union[str, Any] )   ->int:



					'''simple docstring'''







					super().__init__()
					_UpperCAmelCase		:       Any								=      nn.Parameter(torch.randn(1 ) )
					_UpperCAmelCase		:       Tuple								=      nn.Parameter(torch.randn(1 ) )






			def 							lowerCAmelCase__    (		self  :     str	,  lowerCamelCase__  :     Optional[int] )   ->Union[str, Any]:



					'''simple docstring'''







					return x * self.a + self.b






class 	lowerCAmelCase__			(			unittest.TestCase	):




			def 							lowerCAmelCase__    (		self  :     Any )   ->Any:



					'''simple docstring'''







					with tempfile.TemporaryDirectory() as tmpdir:
							set_seed(42 )
							_UpperCAmelCase		:       List[Any]								=      DummyModel()
							_UpperCAmelCase		:       List[str]								=      torch.optim.Adam(params=model.parameters()	,  lr=1E-3 )
							_UpperCAmelCase					,    _UpperCAmelCase		:       Union[str, Any]								=      dummy_dataloaders()
							_UpperCAmelCase		:       Tuple								=      ProjectConfiguration(total_limit=1	,  project_dir=lowerCamelCase__	,  automatic_checkpoint_naming=lowerCamelCase__ )
							# Train baseline
							_UpperCAmelCase		:       List[str]								=      Accelerator(project_config=lowerCamelCase__ )
							_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase		:       Optional[Any]								=      accelerator.prepare(
							    lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							# Save initial
							accelerator.save_state()

							# Save second state
							accelerator.save_state()
							self.assertEqual(len(os.listdir(accelerator.project_dir ) )	,  1 )




			def 							lowerCAmelCase__    (		self  :     Any )   ->Tuple:



					'''simple docstring'''







					with tempfile.TemporaryDirectory() as tmpdir:
							set_seed(42 )
							_UpperCAmelCase		:       List[Any]								=      DummyModel()
							_UpperCAmelCase		:       Dict								=      torch.optim.Adam(params=model.parameters()	,  lr=1E-3 )
							_UpperCAmelCase					,    _UpperCAmelCase		:       Union[str, Any]								=      dummy_dataloaders()
							# Train baseline
							_UpperCAmelCase		:       List[str]								=      Accelerator()
							_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase		:       Optional[Any]								=      accelerator.prepare(
							    lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							# Save initial
							_UpperCAmelCase		:       Any								=      os.path.join(lowerCamelCase__	,  "initial" )
							accelerator.save_state(lowerCamelCase__ )
							((_UpperCAmelCase)					,    (_UpperCAmelCase))		:       Union[str, Any]								=      model.a.item(), model.b.item()
							_UpperCAmelCase		:       str								=      optimizer.state_dict()
							_UpperCAmelCase		:       Tuple								=      train(3	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							((_UpperCAmelCase)					,    (_UpperCAmelCase))		:       Optional[int]								=      model.a.item(), model.b.item()
							_UpperCAmelCase		:       List[str]								=      optimizer.state_dict()

							# Train partially
							set_seed(42 )
							_UpperCAmelCase		:       Optional[Any]								=      DummyModel()
							_UpperCAmelCase		:       int								=      torch.optim.Adam(params=model.parameters()	,  lr=1E-3 )
							_UpperCAmelCase					,    _UpperCAmelCase		:       List[Any]								=      dummy_dataloaders()
							_UpperCAmelCase		:       Union[str, Any]								=      Accelerator()
							_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase		:       Optional[Any]								=      accelerator.prepare(
							    lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							accelerator.load_state(lowerCamelCase__ )
							((_UpperCAmelCase)					,    (_UpperCAmelCase))		:       Tuple								=      model.a.item(), model.b.item()
							_UpperCAmelCase		:       Union[str, Any]								=      optimizer.state_dict()
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )

							_UpperCAmelCase		:       Optional[Any]								=      train(2	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							# Save everything
							_UpperCAmelCase		:       Optional[Any]								=      os.path.join(lowerCamelCase__	,  "checkpoint" )
							accelerator.save_state(lowerCamelCase__ )

							# Load everything back in and make sure all states work
							accelerator.load_state(lowerCamelCase__ )
							test_rands += train(1	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							((_UpperCAmelCase)					,    (_UpperCAmelCase))		:       Union[str, Any]								=      model.a.item(), model.b.item()
							_UpperCAmelCase		:       Dict								=      optimizer.state_dict()
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )




			def 							lowerCAmelCase__    (		self  :     Tuple )   ->Any:



					'''simple docstring'''







					with tempfile.TemporaryDirectory() as tmpdir:
							set_seed(42 )
							_UpperCAmelCase		:       Optional[Any]								=      DummyModel()
							_UpperCAmelCase		:       int								=      torch.optim.Adam(params=model.parameters()	,  lr=1E-3 )
							_UpperCAmelCase					,    _UpperCAmelCase		:       Any								=      dummy_dataloaders()
							_UpperCAmelCase		:       Union[str, Any]								=      ProjectConfiguration(automatic_checkpoint_naming=lowerCamelCase__ )

							# Train baseline
							_UpperCAmelCase		:       str								=      Accelerator(project_dir=lowerCamelCase__	,  project_config=lowerCamelCase__ )
							_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase		:       Optional[Any]								=      accelerator.prepare(
							    lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							# Save initial
							accelerator.save_state()
							((_UpperCAmelCase)					,    (_UpperCAmelCase))		:       Optional[Any]								=      model.a.item(), model.b.item()
							_UpperCAmelCase		:       Union[str, Any]								=      optimizer.state_dict()
							_UpperCAmelCase		:       Optional[Any]								=      train(3	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							((_UpperCAmelCase)					,    (_UpperCAmelCase))		:       Dict								=      model.a.item(), model.b.item()
							_UpperCAmelCase		:       Optional[Any]								=      optimizer.state_dict()

							# Train partially
							set_seed(42 )
							_UpperCAmelCase		:       Optional[Any]								=      DummyModel()
							_UpperCAmelCase		:       str								=      torch.optim.Adam(params=model.parameters()	,  lr=1E-3 )
							_UpperCAmelCase					,    _UpperCAmelCase		:       Dict								=      dummy_dataloaders()
							_UpperCAmelCase		:       Optional[Any]								=      ProjectConfiguration(iteration=1	,  automatic_checkpoint_naming=lowerCamelCase__ )
							_UpperCAmelCase		:       Optional[Any]								=      Accelerator(project_dir=lowerCamelCase__	,  project_config=lowerCamelCase__ )
							_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase		:       Union[str, Any]								=      accelerator.prepare(
							    lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							accelerator.load_state(os.path.join(lowerCamelCase__	,  "checkpoints"	,  "checkpoint_0" ) )
							((_UpperCAmelCase)					,    (_UpperCAmelCase))		:       int								=      model.a.item(), model.b.item()
							_UpperCAmelCase		:       int								=      optimizer.state_dict()
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )

							_UpperCAmelCase		:       str								=      train(2	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							# Save everything
							accelerator.save_state()

							# Load everything back in and make sure all states work
							accelerator.load_state(os.path.join(lowerCamelCase__	,  "checkpoints"	,  "checkpoint_1" ) )
							test_rands += train(1	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							((_UpperCAmelCase)					,    (_UpperCAmelCase))		:       Dict								=      model.a.item(), model.b.item()
							_UpperCAmelCase		:       Dict								=      optimizer.state_dict()
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )
							self.assertEqual(lowerCamelCase__	,  lowerCamelCase__ )




			def 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->str:



					'''simple docstring'''







					_UpperCAmelCase		:       int								=      torch.tensor([1, 2, 3] )
					_UpperCAmelCase		:       Any								=      torch.tensor([2, 3, 4] )
					_UpperCAmelCase		:       Optional[int]								=      DummyModel()
					_UpperCAmelCase		:       List[Any]								=      torch.optim.Adam(net.parameters() )
					_UpperCAmelCase		:       Union[str, Any]								=      Accelerator()
					with self.assertRaises(lowerCamelCase__ ) as ve:
							accelerator.register_for_checkpointing(lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
					_UpperCAmelCase		:       Tuple								=      str(ve.exception )
					self.assertTrue("Item at index 0" in message )
					self.assertTrue("Item at index 1" in message )
					self.assertFalse("Item at index 2" in message )
					self.assertFalse("Item at index 3" in message )




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



					'''simple docstring'''







					with tempfile.TemporaryDirectory() as tmpdir:
							set_seed(42 )
							_UpperCAmelCase		:       Optional[int]								=      DummyModel()
							_UpperCAmelCase		:       int								=      torch.optim.Adam(params=model.parameters()	,  lr=1E-3 )
							_UpperCAmelCase		:       int								=      torch.optim.lr_scheduler.StepLR(lowerCamelCase__	,  step_size=1	,  gamma=0.9_9 )
							_UpperCAmelCase					,    _UpperCAmelCase		:       Any								=      dummy_dataloaders()
							_UpperCAmelCase		:       Any								=      ProjectConfiguration(automatic_checkpoint_naming=lowerCamelCase__ )
							# Train baseline
							_UpperCAmelCase		:       Optional[Any]								=      Accelerator(project_dir=lowerCamelCase__	,  project_config=lowerCamelCase__ )
							_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase		:       int								=      accelerator.prepare(
							    lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							# Save initial
							accelerator.save_state()
							_UpperCAmelCase		:       Union[str, Any]								=      scheduler.state_dict()
							train(3	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
							self.assertNotEqual(lowerCamelCase__	,  scheduler.state_dict() )

							# Load everything back in and make sure all states work
							accelerator.load_state(os.path.join(lowerCamelCase__	,  "checkpoints"	,  "checkpoint_0" ) )
							self.assertEqual(lowerCamelCase__	,  scheduler.state_dict() )




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



					'''simple docstring'''







					with tempfile.TemporaryDirectory() as tmpdir:
							set_seed(42 )
							_UpperCAmelCase		:       Dict								=      DummyModel()
							_UpperCAmelCase		:       int								=      ProjectConfiguration(automatic_checkpoint_naming=lowerCamelCase__	,  total_limit=2 )
							# Train baseline
							_UpperCAmelCase		:       Optional[int]								=      Accelerator(project_dir=lowerCamelCase__	,  project_config=lowerCamelCase__ )
							_UpperCAmelCase		:       Union[str, Any]								=      accelerator.prepare(lowerCamelCase__ )
							# Save 3 states:
							for _ in range(11 ):
									accelerator.save_state()
							self.assertTrue(not os.path.exists(os.path.join(lowerCamelCase__	,  "checkpoints"	,  "checkpoint_0" ) ) )
							self.assertTrue(os.path.exists(os.path.join(lowerCamelCase__	,  "checkpoints"	,  "checkpoint_9" ) ) )
							self.assertTrue(os.path.exists(os.path.join(lowerCamelCase__	,  "checkpoints"	,  "checkpoint_10" ) ) )






			@require_cuda
			def 							lowerCAmelCase__    (		self  :     List[Any] )   ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       str								=      ["torchrun", F"""--nproc_per_node={torch.cuda.device_count()}""", inspect.getfile(self.__class__ )]
					execute_subprocess_async(lowerCamelCase__	,  env=os.environ.copy() )


if __name__ == "__main__":
						lowerCamelCase__    	=     '/tmp/accelerate/state_checkpointing'
						lowerCamelCase__    	=     DummyModel()
						lowerCamelCase__    	=     torch.optim.Adam(params=model.parameters(), lr=1e-3)
						lowerCamelCase__    	=     torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.99)
						lowerCamelCase__	,lowerCamelCase__    	=     dummy_dataloaders()
						lowerCamelCase__    	=     ProjectConfiguration(automatic_checkpoint_naming=True)
						# Train baseline
						lowerCamelCase__    	=     Accelerator(project_dir=savedir, project_config=project_config, mixed_precision='no')
						if accelerator.process_index == 0:
												if os.path.exists(savedir):
																		shutil.rmtree(savedir)
												os.makedirs(savedir)
						lowerCamelCase__	,lowerCamelCase__	,lowerCamelCase__	,lowerCamelCase__	,lowerCamelCase__    	=     accelerator.prepare(
						    model, optimizer, train_dataloader, valid_dataloader, scheduler
						)
						lowerCamelCase__	,lowerCamelCase__    	=     accelerator.prepare(model, optimizer)
						train(3, model, train_dataloader, optimizer, accelerator, scheduler)
						# Check that the intial optimizer is loaded on the GPU
						for group in optimizer.param_groups:
												lowerCamelCase__    	=     group['params'][0].device
												break
						assert param_device.type == accelerator.device.type
						lowerCamelCase__    	=     model.cpu()
						accelerator.wait_for_everyone()
						accelerator.save_state()
						accelerator.wait_for_everyone()

						# Check CPU state
						accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='cpu')
						for group in optimizer.param_groups:
												lowerCamelCase__    	=     group['params'][0].device
												break
						assert (
						    param_device.type == torch.device('cpu').type
						), F"Loaded optimizer states did not match, expected to be loaded on the CPU but got {param_device}"

						# Check device state
						model.to(accelerator.device)
						accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='on_device')
						for group in optimizer.param_groups:
												lowerCamelCase__    	=     group['params'][0].device
												break
						assert (
						    param_device.type == accelerator.device.type
						), F"Loaded optimizer states did not match, expected to be loaded on {accelerator.device} but got {param_device}"

						# Check error
						with pytest.raises(TypeError, match='Unsupported optimizer map location passed'):
												accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='invalid')
						accelerator.wait_for_everyone()
						if accelerator.process_index == 0:
												shutil.rmtree(savedir)
						accelerator.wait_for_everyone()







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



