


'''simple docstring'''


import inspect
import unittest

from transformers import MobileViTVaConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, 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 transformers import MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation, MobileViTVaModel
						from transformers.models.mobilevitva.modeling_mobilevitva import (
						    MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
						    make_divisible,
						)


if is_vision_available():
						from PIL import Image

						from transformers import MobileViTImageProcessor






class 	lowerCAmelCase__			(			UpperCAmelCase__	):




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



			'''simple docstring'''







			_UpperCAmelCase		:       Any								=      self.config_class(**self.inputs_dict )
			self.parent.assertTrue(hasattr(lowerCamelCase__	,  "width_multiplier" ) )






class 	lowerCAmelCase__			:




	def __init__(		self  :     Optional[Any]	,  lowerCamelCase__  :     Dict	,  lowerCamelCase__  :     Any=13	,  lowerCamelCase__  :     int=64	,  lowerCamelCase__  :     Any=2	,  lowerCamelCase__  :     List[Any]=3	,  lowerCamelCase__  :     Optional[Any]="swish"	,  lowerCamelCase__  :     Union[str, Any]=3	,  lowerCamelCase__  :     Any=32	,  lowerCamelCase__  :     Tuple=0.1	,  lowerCamelCase__  :     Union[str, Any]=0.0_2	,  lowerCamelCase__  :     Union[str, Any]=True	,  lowerCamelCase__  :     Optional[Any]=True	,  lowerCamelCase__  :     List[Any]=10	,  lowerCamelCase__  :     Dict=None	,  lowerCamelCase__  :     Dict=0.2_5	,  lowerCamelCase__  :     Union[str, Any]=0.0	,  lowerCamelCase__  :     int=0.0	,  )   ->Dict:



			'''simple docstring'''







			_UpperCAmelCase		:       Tuple								=      parent
			_UpperCAmelCase		:       List[str]								=      batch_size
			_UpperCAmelCase		:       Union[str, Any]								=      image_size
			_UpperCAmelCase		:       Optional[int]								=      patch_size
			_UpperCAmelCase		:       List[str]								=      num_channels
			_UpperCAmelCase		:       List[str]								=      make_divisible(5_12 * width_multiplier	,  divisor=8 )
			_UpperCAmelCase		:       Any								=      hidden_act
			_UpperCAmelCase		:       Optional[int]								=      conv_kernel_size
			_UpperCAmelCase		:       List[Any]								=      output_stride
			_UpperCAmelCase		:       List[Any]								=      classifier_dropout_prob
			_UpperCAmelCase		:       Optional[Any]								=      use_labels
			_UpperCAmelCase		:       Any								=      is_training
			_UpperCAmelCase		:       Optional[Any]								=      num_labels
			_UpperCAmelCase		:       Optional[Any]								=      initializer_range
			_UpperCAmelCase		:       Optional[int]								=      scope
			_UpperCAmelCase		:       Dict								=      width_multiplier
			_UpperCAmelCase		:       str								=      ffn_dropout
			_UpperCAmelCase		:       Tuple								=      attn_dropout




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



			'''simple docstring'''







			_UpperCAmelCase		:       List[Any]								=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

			_UpperCAmelCase		:       int								=      None
			_UpperCAmelCase		:       Any								=      None
			if self.use_labels:
					_UpperCAmelCase		:       Optional[int]								=      ids_tensor([self.batch_size]	,  self.num_labels )
					_UpperCAmelCase		:       Any								=      ids_tensor([self.batch_size, self.image_size, self.image_size]	,  self.num_labels )

			_UpperCAmelCase		:       Tuple								=      self.get_config()

			return config, pixel_values, labels, pixel_labels




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



			'''simple docstring'''







			return MobileViTVaConfig(
			    image_size=self.image_size	,  patch_size=self.patch_size	,  num_channels=self.num_channels	,  hidden_act=self.hidden_act	,  conv_kernel_size=self.conv_kernel_size	,  output_stride=self.output_stride	,  classifier_dropout_prob=self.classifier_dropout_prob	,  initializer_range=self.initializer_range	,  width_multiplier=self.width_multiplier	,  ffn_dropout=self.ffn_dropout_prob	,  attn_dropout=self.attn_dropout_prob	,  )




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



			'''simple docstring'''







			_UpperCAmelCase		:       Optional[int]								=      MobileViTVaModel(config=lowerCamelCase__ )
			model.to(lowerCamelCase__ )
			model.eval()
			_UpperCAmelCase		:       Optional[int]								=      model(lowerCamelCase__ )
			self.parent.assertEqual(
			    result.last_hidden_state.shape	,  (
			        self.batch_size,
			        self.last_hidden_size,
			        self.image_size // self.output_stride,
			        self.image_size // self.output_stride,
			    )	,  )




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



			'''simple docstring'''







			_UpperCAmelCase		:       Any								=      self.num_labels
			_UpperCAmelCase		:       Optional[int]								=      MobileViTVaForImageClassification(lowerCamelCase__ )
			model.to(lowerCamelCase__ )
			model.eval()
			_UpperCAmelCase		:       Optional[int]								=      model(lowerCamelCase__	,  labels=lowerCamelCase__ )
			self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_labels) )




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



			'''simple docstring'''







			_UpperCAmelCase		:       List[Any]								=      self.num_labels
			_UpperCAmelCase		:       Optional[int]								=      MobileViTVaForSemanticSegmentation(lowerCamelCase__ )
			model.to(lowerCamelCase__ )
			model.eval()
			_UpperCAmelCase		:       Optional[Any]								=      model(lowerCamelCase__ )
			self.parent.assertEqual(
			    result.logits.shape	,  (
			        self.batch_size,
			        self.num_labels,
			        self.image_size // self.output_stride,
			        self.image_size // self.output_stride,
			    )	,  )
			_UpperCAmelCase		:       List[Any]								=      model(lowerCamelCase__	,  labels=lowerCamelCase__ )
			self.parent.assertEqual(
			    result.logits.shape	,  (
			        self.batch_size,
			        self.num_labels,
			        self.image_size // self.output_stride,
			        self.image_size // self.output_stride,
			    )	,  )






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



			'''simple docstring'''







			_UpperCAmelCase		:       Dict								=      self.prepare_config_and_inputs()
			_UpperCAmelCase		:       List[str]								=      config_and_inputs
			_UpperCAmelCase		:       Any								=      {"pixel_values": pixel_values}
			return config, inputs_dict






@require_torch
class 	lowerCAmelCase__			(			UpperCAmelCase__    ,       UpperCAmelCase__    ,       unittest.TestCase	):
	lowerCAmelCase :	Dict          =     (
	    (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
	    if is_torch_available()
	    else ()
	)

	lowerCAmelCase :	Optional[Any]          =     (
	    {
	        "feature-extraction": MobileViTVaModel,
	        "image-classification": MobileViTVaForImageClassification,
	        "image-segmentation": MobileViTVaForSemanticSegmentation,
	    }
	    if is_torch_available()
	    else {}
	)

	lowerCAmelCase :	Optional[int]          =     False
	lowerCAmelCase :	Optional[Any]          =     False
	lowerCAmelCase :	Dict          =     False
	lowerCAmelCase :	str          =     False




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



			'''simple docstring'''







			_UpperCAmelCase		:       Dict								=      MobileViTVaModelTester(self )
			_UpperCAmelCase		:       List[Any]								=      MobileViTVaConfigTester(self	,  config_class=lowerCamelCase__	,  has_text_modality=lowerCamelCase__ )




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



			'''simple docstring'''







			self.config_tester.run_common_tests()




	@unittest.skip(reason="MobileViTV2 does not use inputs_embeds" )
	def 							lowerCAmelCase__    (		self  :     int )   ->Dict:



			'''simple docstring'''







			pass




	@unittest.skip(reason="MobileViTV2 does not support input and output embeddings" )
	def 							lowerCAmelCase__    (		self  :     Tuple )   ->Dict:



			'''simple docstring'''







			pass




	@unittest.skip(reason="MobileViTV2 does not output attentions" )
	def 							lowerCAmelCase__    (		self  :     Union[str, Any] )   ->Any:



			'''simple docstring'''







			pass




	@require_torch_multi_gpu
	@unittest.skip(reason="Got `CUDA error: misaligned address` for tests after this one being run." )
	def 							lowerCAmelCase__    (		self  :     int )   ->Any:



			'''simple docstring'''







			pass




	@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests." )
	def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->Tuple:



			'''simple docstring'''







			pass




	def 							lowerCAmelCase__    (		self  :     Tuple )   ->Union[str, Any]:



			'''simple docstring'''







			_UpperCAmelCase		:       List[str]								=      self.model_tester.prepare_config_and_inputs_for_common()

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

					_UpperCAmelCase		:       Any								=      ["pixel_values"]
					self.assertListEqual(arg_names[:1]	,  lowerCamelCase__ )




	def 							lowerCAmelCase__    (		self  :     Tuple )   ->str:



			'''simple docstring'''







			_UpperCAmelCase		:       Any								=      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*lowerCamelCase__ )




	def 							lowerCAmelCase__    (		self  :     Union[str, Any] )   ->Tuple:



			'''simple docstring'''







			def check_hidden_states_output(lowerCamelCase__  :     Union[str, Any]	,  lowerCamelCase__  :     Dict	,  lowerCamelCase__  :     int ):
					_UpperCAmelCase		:       str								=      model_class(lowerCamelCase__ )
					model.to(lowerCamelCase__ )
					model.eval()

					with torch.no_grad():
							_UpperCAmelCase		:       Tuple								=      model(**self._prepare_for_class(lowerCamelCase__	,  lowerCamelCase__ ) )

					_UpperCAmelCase		:       Optional[Any]								=      outputs.hidden_states

					_UpperCAmelCase		:       Tuple								=      5
					self.assertEqual(len(lowerCamelCase__ )	,  lowerCamelCase__ )

					# MobileViTV2's feature maps are of shape (batch_size, num_channels, height, width)
					# with the width and height being successively divided by 2.
					_UpperCAmelCase		:       Optional[Any]								=      2
					for i in range(len(lowerCamelCase__ ) ):
							self.assertListEqual(
							    list(hidden_states[i].shape[-2:] )	,  [self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]	,  )
							divisor *= 2

					self.assertEqual(self.model_tester.output_stride	,  divisor // 2 )

			_UpperCAmelCase		:       List[Any]								=      self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					_UpperCAmelCase		:       Union[str, Any]								=      True
					check_hidden_states_output(lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )

					# check that output_hidden_states also work using config
					del inputs_dict["output_hidden_states"]
					_UpperCAmelCase		:       int								=      True

					check_hidden_states_output(lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )




	def 							lowerCAmelCase__    (		self  :     str )   ->Tuple:



			'''simple docstring'''







			_UpperCAmelCase		:       Tuple								=      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_image_classification(*lowerCamelCase__ )




	def 							lowerCAmelCase__    (		self  :     str )   ->Dict:



			'''simple docstring'''







			_UpperCAmelCase		:       Tuple								=      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_semantic_segmentation(*lowerCamelCase__ )






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



			'''simple docstring'''







			for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					_UpperCAmelCase		:       int								=      MobileViTVaModel.from_pretrained(lowerCamelCase__ )
					self.assertIsNotNone(lowerCamelCase__ )






def  __lowerCAmelCase ():
		_UpperCAmelCase		:       Union[str, Any]								=      Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"   )
		return image






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




	@cached_property
	def 							lowerCAmelCase__    (		self  :     str )   ->Any:



			'''simple docstring'''







			return (
			    MobileViTImageProcessor.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256" )
			    if is_vision_available()
			    else None
			)




	@slow
	def 							lowerCAmelCase__    (		self  :     List[Any] )   ->int:



			'''simple docstring'''







			_UpperCAmelCase		:       Any								=      MobileViTVaForImageClassification.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256" ).to(
			    lowerCamelCase__ )

			_UpperCAmelCase		:       str								=      self.default_image_processor
			_UpperCAmelCase		:       str								=      prepare_img()
			_UpperCAmelCase		:       Union[str, Any]								=      image_processor(images=lowerCamelCase__	,  return_tensors="pt" ).to(lowerCamelCase__ )

			# forward pass
			with torch.no_grad():
					_UpperCAmelCase		:       Optional[Any]								=      model(**lowerCamelCase__ )

			# verify the logits
			_UpperCAmelCase		:       List[Any]								=      torch.Size((1, 10_00) )
			self.assertEqual(outputs.logits.shape	,  lowerCamelCase__ )

			_UpperCAmelCase		:       Tuple								=      torch.tensor([-1.6336E00, -7.3204E-02, -5.1883E-01] ).to(lowerCamelCase__ )

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




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



			'''simple docstring'''







			_UpperCAmelCase		:       str								=      MobileViTVaForSemanticSegmentation.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3" )
			_UpperCAmelCase		:       Union[str, Any]								=      model.to(lowerCamelCase__ )

			_UpperCAmelCase		:       Optional[Any]								=      MobileViTImageProcessor.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3" )

			_UpperCAmelCase		:       Union[str, Any]								=      prepare_img()
			_UpperCAmelCase		:       List[Any]								=      image_processor(images=lowerCamelCase__	,  return_tensors="pt" ).to(lowerCamelCase__ )

			# forward pass
			with torch.no_grad():
					_UpperCAmelCase		:       List[Any]								=      model(**lowerCamelCase__ )
			_UpperCAmelCase		:       List[Any]								=      outputs.logits

			# verify the logits
			_UpperCAmelCase		:       Optional[Any]								=      torch.Size((1, 21, 32, 32) )
			self.assertEqual(logits.shape	,  lowerCamelCase__ )

			_UpperCAmelCase		:       int								=      torch.tensor(
			    [
			        [[7.0_8_6_3, 7.1_5_2_5, 6.8_2_0_1], [6.6_9_3_1, 6.8_7_7_0, 6.8_9_3_3], [6.2_9_7_8, 7.0_3_6_6, 6.9_6_3_6]],
			        [[-3.7_1_3_4, -3.6_7_1_2, -3.6_6_7_5], [-3.5_8_2_5, -3.3_5_4_9, -3.4_7_7_7], [-3.3_4_3_5, -3.3_9_7_9, -3.2_8_5_7]],
			        [[-2.9_3_2_9, -2.8_0_0_3, -2.7_3_6_9], [-3.0_5_6_4, -2.4_7_8_0, -2.0_2_0_7], [-2.6_8_8_9, -1.9_2_9_8, -1.7_6_4_0]],
			    ]	,  device=lowerCamelCase__	,  )

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






	@slow
	def 							lowerCAmelCase__    (		self  :     Any )   ->Any:



			'''simple docstring'''







			_UpperCAmelCase		:       Any								=      MobileViTVaForSemanticSegmentation.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3" )
			_UpperCAmelCase		:       int								=      model.to(lowerCamelCase__ )

			_UpperCAmelCase		:       Any								=      MobileViTImageProcessor.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3" )

			_UpperCAmelCase		:       List[str]								=      prepare_img()
			_UpperCAmelCase		:       Tuple								=      image_processor(images=lowerCamelCase__	,  return_tensors="pt" ).to(lowerCamelCase__ )

			# forward pass
			with torch.no_grad():
					_UpperCAmelCase		:       Optional[Any]								=      model(**lowerCamelCase__ )

			_UpperCAmelCase		:       List[Any]								=      outputs.logits.detach().cpu()

			_UpperCAmelCase		:       Optional[Any]								=      image_processor.post_process_semantic_segmentation(outputs=lowerCamelCase__	,  target_sizes=[(50, 60)] )
			_UpperCAmelCase		:       Tuple								=      torch.Size((50, 60) )
			self.assertEqual(segmentation[0].shape	,  lowerCamelCase__ )

			_UpperCAmelCase		:       List[Any]								=      image_processor.post_process_semantic_segmentation(outputs=lowerCamelCase__ )
			_UpperCAmelCase		:       Tuple								=      torch.Size((32, 32) )
			self.assertEqual(segmentation[0].shape	,  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()



