





import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VideoMAEConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch
				from torch import nn

				from transformers import (
				    MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
				    VideoMAEForPreTraining,
				    VideoMAEForVideoClassification,
				    VideoMAEModel,
				)
				from transformers.models.videomae.modeling_videomae import VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
				from transformers import VideoMAEImageProcessor



class    _a     :






							"""simple docstring"""

							def __init__(      self				:     Tuple							,	a				:     Optional[int]							,	a				:     Optional[Any]=13							,	a				:     List[Any]=10							,	a				:     str=3							,	a				:     Tuple=2							,	a				:     List[str]=2							,	a				:     Union[str, Any]=2							,	a				:     Optional[int]=True							,	a				:     Optional[Any]=True							,	a				:     Optional[int]=32							,	a				:     Tuple=5							,	a				:     Any=4							,	a				:     Dict=37							,	a				:     Union[str, Any]="gelu"							,	a				:     Dict=0.1							,	a				:     Dict=0.1							,	a				:     int=10							,	a				:     int=0.02							,	a				:     Tuple=0.9							,	a				:     Union[str, Any]=None							,	)	->str:
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					parent
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					batch_size
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					image_size
													SCREAMING_SNAKE_CASE__     :				int   =					num_channels
													SCREAMING_SNAKE_CASE__     :				Dict   =					patch_size
													SCREAMING_SNAKE_CASE__     :				Any   =					tubelet_size
													SCREAMING_SNAKE_CASE__     :				str   =					num_frames
													SCREAMING_SNAKE_CASE__     :				Tuple   =					is_training
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					use_labels
													SCREAMING_SNAKE_CASE__     :				str   =					hidden_size
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					num_hidden_layers
													SCREAMING_SNAKE_CASE__     :				Any   =					num_attention_heads
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					intermediate_size
													SCREAMING_SNAKE_CASE__     :				Dict   =					hidden_act
													SCREAMING_SNAKE_CASE__     :				str   =					hidden_dropout_prob
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					attention_probs_dropout_prob
													SCREAMING_SNAKE_CASE__     :				Dict   =					type_sequence_label_size
													SCREAMING_SNAKE_CASE__     :				Tuple   =					initializer_range
													SCREAMING_SNAKE_CASE__     :				Tuple   =					mask_ratio
													SCREAMING_SNAKE_CASE__     :				List[str]   =					scope

													# in VideoMAE, the number of tokens equals num_frames/tubelet_size * num_patches per frame
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					(image_size // patch_size) ** 2
													SCREAMING_SNAKE_CASE__     :				Tuple   =					(num_frames // tubelet_size) * self.num_patches_per_frame

													# use this variable to define bool_masked_pos
													SCREAMING_SNAKE_CASE__     :				Tuple   =					int(mask_ratio * self.seq_length    )

							def      A_  (      self				:     Any    )	->List[str]:
													SCREAMING_SNAKE_CASE__     :				str   =					floats_tensor(
													    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]    )

													SCREAMING_SNAKE_CASE__     :				Tuple   =					None
													if self.use_labels:
																			SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					ids_tensor([self.batch_size]							,	self.type_sequence_label_size    )

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					self.get_config()

													return config, pixel_values, labels

							def      A_  (      self				:     str    )	->Optional[Any]:
													return VideoMAEConfig(
													    image_size=self.image_size							,	patch_size=self.patch_size							,	num_channels=self.num_channels							,	num_frames=self.num_frames							,	tubelet_size=self.tubelet_size							,	hidden_size=self.hidden_size							,	num_hidden_layers=self.num_hidden_layers							,	num_attention_heads=self.num_attention_heads							,	intermediate_size=self.intermediate_size							,	hidden_act=self.hidden_act							,	hidden_dropout_prob=self.hidden_dropout_prob							,	attention_probs_dropout_prob=self.attention_probs_dropout_prob							,	is_decoder=a							,	initializer_range=self.initializer_range							,	)

							def      A_  (      self				:     Dict							,	a				:     Optional[Any]							,	a				:     Dict							,	a				:     str    )	->int:
													SCREAMING_SNAKE_CASE__     :				Tuple   =					VideoMAEModel(config=a    )
													model.to(a    )
													model.eval()
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					model(a    )
													self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, self.seq_length, self.hidden_size)    )

							def      A_  (      self				:     str							,	a				:     Union[str, Any]							,	a				:     Tuple							,	a				:     Union[str, Any]    )	->str:
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					VideoMAEForPreTraining(a    )
													model.to(a    )
													model.eval()
													# important: each video needs to have the same number of masked patches
													# hence we define a single mask, which we then repeat for each example in the batch
													SCREAMING_SNAKE_CASE__     :				str   =					torch.ones((self.num_masks,)    )
													SCREAMING_SNAKE_CASE__     :				Any   =					torch.cat([mask, torch.zeros(self.seq_length - mask.size(0    )    )]    )
													SCREAMING_SNAKE_CASE__     :				Tuple   =					mask.expand(self.batch_size							,	-1    ).bool()

													SCREAMING_SNAKE_CASE__     :				Dict   =					model(a							,	a    )
													# model only returns predictions for masked patches
													SCREAMING_SNAKE_CASE__     :				Dict   =					mask.sum().item()
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					3 * self.tubelet_size * self.patch_size**2
													self.parent.assertEqual(result.logits.shape							,	(self.batch_size, num_masked_patches, decoder_num_labels)    )






							def      A_  (      self				:     Optional[Any]    )	->int:
													SCREAMING_SNAKE_CASE__     :				int   =					self.prepare_config_and_inputs()
													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Optional[int]   =					config_and_inputs
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					{"pixel_values": pixel_values}
													return config, inputs_dict







@require_torch
class    _a     (				lowercase__					,							lowercase__					,							unittest.TestCase							):






							"""simple docstring"""

							snake_case_							 =						(
							    (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
							)
							snake_case_							 =						(
							    {"feature-extraction": VideoMAEModel, "video-classification": VideoMAEForVideoClassification}
							    if is_torch_available()
							    else {}
							)

							snake_case_							 =						False
							snake_case_							 =						False
							snake_case_							 =						False
							snake_case_							 =						False

							def      A_  (      self				:     Dict    )	->int:
													SCREAMING_SNAKE_CASE__     :				Tuple   =					VideoMAEModelTester(self    )
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					ConfigTester(self							,	config_class=a							,	has_text_modality=a							,	hidden_size=37    )

							def      A_  (      self				:     Optional[int]							,	a				:     List[Any]							,	a				:     str							,	a				:     Dict=False    )	->Tuple:
													SCREAMING_SNAKE_CASE__     :				Any   =					copy.deepcopy(a    )

													if model_class == VideoMAEForPreTraining:
																			# important: each video needs to have the same number of masked patches
																			# hence we define a single mask, which we then repeat for each example in the batch
																			SCREAMING_SNAKE_CASE__     :				str   =					torch.ones((self.model_tester.num_masks,)    )
																			SCREAMING_SNAKE_CASE__     :				List[str]   =					torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0    )    )]    )
																			SCREAMING_SNAKE_CASE__     :				List[str]   =					mask.expand(self.model_tester.batch_size							,	-1    ).bool()
																			SCREAMING_SNAKE_CASE__     :				List[str]   =					bool_masked_pos.to(a    )

													if return_labels:
																			if model_class in [
																			    *get_values(a    ),
																			]:
																									SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					torch.zeros(
																									    self.model_tester.batch_size							,	dtype=torch.long							,	device=a    )

													return inputs_dict

							def      A_  (      self				:     Optional[int]    )	->List[str]:
													self.config_tester.run_common_tests()

							@unittest.skip(reason="VideoMAE does not use inputs_embeds"    )
							def      A_  (      self				:     List[Any]    )	->Optional[int]:
													pass

							def      A_  (      self				:     Optional[int]    )	->int:
													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				List[Any]   =					self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																			SCREAMING_SNAKE_CASE__     :				int   =					model_class(a    )
																			self.assertIsInstance(model.get_input_embeddings()							,	(nn.Module)    )
																			SCREAMING_SNAKE_CASE__     :				str   =					model.get_output_embeddings()
																			self.assertTrue(x is None or isinstance(a							,	nn.Linear    )    )

							def      A_  (      self				:     int    )	->Any:
													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Any   =					self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																			SCREAMING_SNAKE_CASE__     :				Optional[int]   =					model_class(a    )
																			SCREAMING_SNAKE_CASE__     :				List[str]   =					inspect.signature(model.forward    )
																			# signature.parameters is an OrderedDict => so arg_names order is deterministic
																			SCREAMING_SNAKE_CASE__     :				str   =					[*signature.parameters.keys()]

																			SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					["pixel_values"]
																			self.assertListEqual(arg_names[:1]							,	a    )

							def      A_  (      self				:     Optional[Any]    )	->Any:
													SCREAMING_SNAKE_CASE__     :				Dict   =					self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*a    )

							def      A_  (      self				:     Optional[Any]    )	->Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_pretraining(*a    )

							@slow
							def      A_  (      self				:     int    )	->Any:
													for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					VideoMAEModel.from_pretrained(a    )
																			self.assertIsNotNone(a    )

							def      A_  (      self				:     str    )	->Optional[Any]:
													if not self.has_attentions:
																			pass

													else:
																			SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Tuple   =					self.model_tester.prepare_config_and_inputs_for_common()
																			SCREAMING_SNAKE_CASE__     :				Any   =					True

																			for model_class in self.all_model_classes:
																									SCREAMING_SNAKE_CASE__     :				List[str]   =					self.model_tester.seq_length - self.model_tester.num_masks
																									SCREAMING_SNAKE_CASE__     :				Tuple   =					(
																									    num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
																									)

																									SCREAMING_SNAKE_CASE__     :				int   =					True
																									SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					False
																									SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					True
																									SCREAMING_SNAKE_CASE__     :				Dict   =					model_class(a    )
																									model.to(a    )
																									model.eval()
																									with torch.no_grad():
																															SCREAMING_SNAKE_CASE__     :				Tuple   =					model(**self._prepare_for_class(a							,	a    )    )
																									SCREAMING_SNAKE_CASE__     :				Dict   =					outputs.attentions
																									self.assertEqual(len(a    )							,	self.model_tester.num_hidden_layers    )

																									# check that output_attentions also work using config
																									del inputs_dict["output_attentions"]
																									SCREAMING_SNAKE_CASE__     :				int   =					True
																									SCREAMING_SNAKE_CASE__     :				str   =					model_class(a    )
																									model.to(a    )
																									model.eval()
																									with torch.no_grad():
																															SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					model(**self._prepare_for_class(a							,	a    )    )
																									SCREAMING_SNAKE_CASE__     :				Optional[int]   =					outputs.attentions
																									self.assertEqual(len(a    )							,	self.model_tester.num_hidden_layers    )

																									self.assertListEqual(
																									    list(attentions[0].shape[-3:]    )							,	[self.model_tester.num_attention_heads, seq_len, seq_len]							,	)
																									SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					len(a    )

																									# Check attention is always last and order is fine
																									SCREAMING_SNAKE_CASE__     :				List[str]   =					True
																									SCREAMING_SNAKE_CASE__     :				Dict   =					True
																									SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					model_class(a    )
																									model.to(a    )
																									model.eval()
																									with torch.no_grad():
																															SCREAMING_SNAKE_CASE__     :				int   =					model(**self._prepare_for_class(a							,	a    )    )

																									self.assertEqual(out_len + 1							,	len(a    )    )

																									SCREAMING_SNAKE_CASE__     :				Tuple   =					outputs.attentions

																									self.assertEqual(len(a    )							,	self.model_tester.num_hidden_layers    )
																									self.assertListEqual(
																									    list(self_attentions[0].shape[-3:]    )							,	[self.model_tester.num_attention_heads, seq_len, seq_len]							,	)

							def      A_  (      self				:     List[str]    )	->List[str]:
													def check_hidden_states_output(a				:     Optional[int]							,	a				:     Optional[Any]							,	a				:     Union[str, Any]    ):
																			SCREAMING_SNAKE_CASE__     :				Tuple   =					model_class(a    )
																			model.to(a    )
																			model.eval()

																			with torch.no_grad():
																									SCREAMING_SNAKE_CASE__     :				List[Any]   =					model(**self._prepare_for_class(a							,	a    )    )

																			SCREAMING_SNAKE_CASE__     :				List[str]   =					outputs.hidden_states
																			SCREAMING_SNAKE_CASE__     :				List[str]   =					self.model_tester.num_hidden_layers + 1
																			self.assertEqual(len(a    )							,	a    )

																			SCREAMING_SNAKE_CASE__     :				str   =					self.model_tester.seq_length - self.model_tester.num_masks
																			SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length

																			self.assertListEqual(
																			    list(hidden_states[0].shape[-2:]    )							,	[seq_length, self.model_tester.hidden_size]							,	)

													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Any   =					self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																			SCREAMING_SNAKE_CASE__     :				Optional[int]   =					True
																			check_hidden_states_output(a							,	a							,	a    )

																			# check that output_hidden_states also work using config
																			del inputs_dict["output_hidden_states"]
																			SCREAMING_SNAKE_CASE__     :				List[Any]   =					True

																			check_hidden_states_output(a							,	a							,	a    )






							@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."    )
							def      A_  (      self				:     str    )	->int:
													pass







def       UpperCAmelCase							(     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					hf_hub_download(
						    repo_id="hf-internal-testing/spaghetti-video"			,   filename="eating_spaghetti.npy"			,   repo_type="dataset"     )
						SCREAMING_SNAKE_CASE__     :				List[Any]   =					np.load(_lowerCamelCase     )
						return list(_lowerCamelCase     )



@require_torch
@require_vision
class    _a     (				unittest.TestCase							):






							"""simple docstring"""

							@cached_property
							def      A_  (      self				:     Optional[Any]    )	->str:
													# logits were tested with a different mean and std, so we use the same here
													return (
													    VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]							,	image_std=[0.5, 0.5, 0.5]    )
													    if is_vision_available()
													    else None
													)

							@slow
							def      A_  (      self				:     Optional[Any]    )	->List[Any]:
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					VideoMAEForVideoClassification.from_pretrained("MCG-NJU/videomae-base-finetuned-kinetics"    ).to(
													    a    )

													SCREAMING_SNAKE_CASE__     :				List[str]   =					self.default_image_processor
													SCREAMING_SNAKE_CASE__     :				str   =					prepare_video()
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					image_processor(a							,	return_tensors="pt"    ).to(a    )

													# forward pass
													with torch.no_grad():
																			SCREAMING_SNAKE_CASE__     :				Any   =					model(**a    )

													# verify the logits
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					torch.Size((1, 4_00)    )
													self.assertEqual(outputs.logits.shape							,	a    )

													SCREAMING_SNAKE_CASE__     :				Dict   =					torch.tensor([0.3669, -0.0688, -0.2421]    ).to(a    )

													self.assertTrue(torch.allclose(outputs.logits[0, :3]							,	a							,	atol=1E-4    )    )






							@slow
							def      A_  (      self				:     Union[str, Any]    )	->Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					VideoMAEForPreTraining.from_pretrained("MCG-NJU/videomae-base-short"    ).to(a    )

													SCREAMING_SNAKE_CASE__     :				List[Any]   =					self.default_image_processor
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					prepare_video()
													SCREAMING_SNAKE_CASE__     :				List[str]   =					image_processor(a							,	return_tensors="pt"    ).to(a    )

													# add boolean mask, indicating which patches to mask
													SCREAMING_SNAKE_CASE__     :				Any   =					hf_hub_download(repo_id="hf-internal-testing/bool-masked-pos"							,	filename="bool_masked_pos.pt"    )
													SCREAMING_SNAKE_CASE__     :				str   =					torch.load(a    )

													# forward pass
													with torch.no_grad():
																			SCREAMING_SNAKE_CASE__     :				Optional[int]   =					model(**a    )

													# verify the logits
													SCREAMING_SNAKE_CASE__     :				List[str]   =					torch.Size([1, 14_08, 15_36]    )
													SCREAMING_SNAKE_CASE__     :				str   =					torch.tensor(
													    [[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]]							,	device=a    )
													self.assertEqual(outputs.logits.shape							,	a    )
													self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]							,	a							,	atol=1E-4    )    )

													# verify the loss (`config.norm_pix_loss` = `True`)
													SCREAMING_SNAKE_CASE__     :				Tuple   =					torch.tensor([0.5142]							,	device=a    )
													self.assertTrue(torch.allclose(outputs.loss							,	a							,	atol=1E-4    )    )

													# verify the loss (`config.norm_pix_loss` = `False`)
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					VideoMAEForPreTraining.from_pretrained("MCG-NJU/videomae-base-short"							,	norm_pix_loss=a    ).to(
													    a    )

													with torch.no_grad():
																			SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					model(**a    )

													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					torch.tensor(torch.tensor([0.6469]    )							,	device=a    )
													self.assertTrue(torch.allclose(outputs.loss							,	a							,	atol=1E-4    )    )






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


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



class    _a     (				lowercase__							):






							"""simple docstring"""

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






							def __iter__(      self				:     str    )	->Tuple:
													SCREAMING_SNAKE_CASE__     :				str   =					[]
													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()    )
													SCREAMING_SNAKE_CASE__     :				int   =					self.tokenizer(a							,	padding=a							,	return_tensors="pt"    )
													for task in range(self.n_tasks    ):
																			for _ in range(self.n_copies    ):
																									yield {
																									    "ids": outputs.input_ids[task],
																									    "task_id": task,
																									    "input_len": outputs.attention_mask[task].sum(),
																									}







class    _a     (				lowercase__							):






							"""simple docstring"""

							def __init__(      self				:     Dict							,	a				:     int							,	a				:     int							,	a				:     Tuple    )	->Dict:
													SCREAMING_SNAKE_CASE__     :				Dict   =					start_length
													SCREAMING_SNAKE_CASE__     :				Any   =					eof_strings
													SCREAMING_SNAKE_CASE__     :				Any   =					tokenizer






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







def       UpperCAmelCase							(     _lowerCamelCase      :  Dict     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				List[Any]   =					re.split("(%s)" % "|".join(_lowerCamelCase     )			,   _lowerCamelCase     )
						# last string should be ""
						return "".join(string_list[:-2]     )







def       UpperCAmelCase							(     _lowerCamelCase      :  Dict			,   _lowerCamelCase      :  Optional[Any]			,   _lowerCamelCase      :  Tuple			,   _lowerCamelCase      :  str			,   _lowerCamelCase      :  int			,   _lowerCamelCase      :  str=20			,   **_lowerCamelCase      :  Dict     ):





						'''simple docstring'''


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

																		SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Any   =					accelerator.gather((generated_tokens, generated_tasks)     )
																		SCREAMING_SNAKE_CASE__     :				Dict   =					generated_tokens.cpu().numpy()
																		SCREAMING_SNAKE_CASE__     :				Any   =					generated_tasks.cpu().numpy()

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

						SCREAMING_SNAKE_CASE__     :				List[Any]   =					[[] for _ in range(_lowerCamelCase     )]
						for task, generated_tokens in gen_token_dict.items():
												for s in generated_tokens:
																		SCREAMING_SNAKE_CASE__     :				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       UpperCAmelCase							(     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					HfArgumentParser(_lowerCamelCase     )
						SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					parser.parse_args()

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

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

						# Use dataset load to feed to accelerate
						SCREAMING_SNAKE_CASE__     :				Dict   =					Accelerator()
						set_seed(args.seed			,   device_specific=_lowerCamelCase     )

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

						# Generation settings
						SCREAMING_SNAKE_CASE__     :				List[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
						SCREAMING_SNAKE_CASE__     :				str   =					load_dataset("openai_humaneval"     )
						SCREAMING_SNAKE_CASE__     :				Any   =					load_metric("code_eval"     )

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

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

						# Run a quick test to see if code evaluation is enabled
						try:
												SCREAMING_SNAKE_CASE__     :				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

						SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Optional[int]   =					accelerator.prepare(_lowerCamelCase			,   _lowerCamelCase     )

						SCREAMING_SNAKE_CASE__     :				Tuple   =					complete_code(
						    _lowerCamelCase			,   _lowerCamelCase			,   _lowerCamelCase			,   _lowerCamelCase			,   n_tasks=_lowerCamelCase			,   batch_size=args.batch_size			,   **_lowerCamelCase			,   )

						if accelerator.is_main_process:
												SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					[]

												for task in tqdm(range(_lowerCamelCase     )     ):
																		SCREAMING_SNAKE_CASE__     :				List[Any]   =					human_eval["test"][task]["test"]
																		SCREAMING_SNAKE_CASE__     :				List[Any]   =					f"""check({human_eval['test'][task]['entry_point']})"""
																		references.append("\n" + test_func + "\n" + entry_point     )

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