'''simple docstring'''





import collections
import inspect
import unittest

from transformers import FocalNetConfig
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_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import (
        FocalNetBackbone,
        FocalNetForImageClassification,
        FocalNetForMaskedImageModeling,
        FocalNetModel,
    )
    from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor




class 			__A :



    def __init__(self       :       str  ,     __a       :       Optional[int]  ,     __a       :       List[str]=13  ,     __a       :       Tuple=32  ,     __a       :       Optional[int]=2  ,     __a       :       Any=3  ,     __a       :       Union[str, Any]=16  ,     __a       :       Tuple=[32, 64, 128]  ,     __a       :       Tuple=[1, 2, 1]  ,     __a       :       Any=[2, 2, 4]  ,     __a       :       int=2  ,     __a       :       Tuple=2.0  ,     __a       :       Dict=True  ,     __a       :       List[str]=0.0  ,     __a       :       List[str]=0.0  ,     __a       :       int=0.1  ,     __a       :       List[Any]="gelu"  ,     __a       :       Optional[int]=False  ,     __a       :       Union[str, Any]=True  ,     __a       :       Optional[Any]=0.02  ,     __a       :       Tuple=1E-5  ,     __a       :       Tuple=True  ,     __a       :       List[Any]=None  ,     __a       :       List[str]=True  ,     __a       :       int=10  ,     __a       :       Optional[int]=8  ,     __a       :       List[Any]=["stage1", "stage2"]  ,     __a       :       Dict=[1, 2]  ,     ):
         UpperCAmelCase_			       =  parent
         UpperCAmelCase_			       =  batch_size
         UpperCAmelCase_			       =  image_size
         UpperCAmelCase_			       =  patch_size
         UpperCAmelCase_			       =  num_channels
         UpperCAmelCase_			       =  embed_dim
         UpperCAmelCase_			       =  hidden_sizes
         UpperCAmelCase_			       =  depths
         UpperCAmelCase_			       =  num_heads
         UpperCAmelCase_			       =  window_size
         UpperCAmelCase_			       =  mlp_ratio
         UpperCAmelCase_			       =  qkv_bias
         UpperCAmelCase_			       =  hidden_dropout_prob
         UpperCAmelCase_			       =  attention_probs_dropout_prob
         UpperCAmelCase_			       =  drop_path_rate
         UpperCAmelCase_			       =  hidden_act
         UpperCAmelCase_			       =  use_absolute_embeddings
         UpperCAmelCase_			       =  patch_norm
         UpperCAmelCase_			       =  layer_norm_eps
         UpperCAmelCase_			       =  initializer_range
         UpperCAmelCase_			       =  is_training
         UpperCAmelCase_			       =  scope
         UpperCAmelCase_			       =  use_labels
         UpperCAmelCase_			       =  type_sequence_label_size
         UpperCAmelCase_			       =  encoder_stride
         UpperCAmelCase_			       =  out_features
         UpperCAmelCase_			       =  out_indices



    def    _lowercase		(self       :       Tuple					):
         UpperCAmelCase_			       =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

         UpperCAmelCase_			       =  None
         if self.use_labels:
              UpperCAmelCase_			       =  ids_tensor([self.batch_size]  ,     self.type_sequence_label_size					)

         UpperCAmelCase_			       =  self.get_config()

         return config, pixel_values, labels



    def    _lowercase		(self       :       Union[str, Any]					):
         return FocalNetConfig(
             image_size=self.image_size  ,     patch_size=self.patch_size  ,     num_channels=self.num_channels  ,     embed_dim=self.embed_dim  ,     hidden_sizes=self.hidden_sizes  ,     depths=self.depths  ,     num_heads=self.num_heads  ,     window_size=self.window_size  ,     mlp_ratio=self.mlp_ratio  ,     qkv_bias=self.qkv_bias  ,     hidden_dropout_prob=self.hidden_dropout_prob  ,     attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,     drop_path_rate=self.drop_path_rate  ,     hidden_act=self.hidden_act  ,     use_absolute_embeddings=self.use_absolute_embeddings  ,     path_norm=self.patch_norm  ,     layer_norm_eps=self.layer_norm_eps  ,     initializer_range=self.initializer_range  ,     encoder_stride=self.encoder_stride  ,     out_features=self.out_features  ,     out_indices=self.out_indices  ,     )



    def    _lowercase		(self       :       Dict  ,     __a       :       str  ,     __a       :       int  ,     __a       :       Union[str, Any]					):
         UpperCAmelCase_			       =  FocalNetModel(config=__a					)
         model.to(__a					)
         model.eval()
         UpperCAmelCase_			       =  model(__a					)

         UpperCAmelCase_			       =  ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths					) - 1))
         UpperCAmelCase_			       =  int(config.embed_dim * 2 ** (len(config.depths					) - 1)					)

         self.parent.assertEqual(result.last_hidden_state.shape  ,     (self.batch_size, expected_seq_len, expected_dim)					)



    def    _lowercase		(self       :       str  ,     __a       :       List[str]  ,     __a       :       Optional[int]  ,     __a       :       List[str]					):
         UpperCAmelCase_			       =  FocalNetBackbone(config=__a					)
         model.to(__a					)
         model.eval()
         UpperCAmelCase_			       =  model(__a					)

         # verify feature maps
         self.parent.assertEqual(len(result.feature_maps					)  ,     len(config.out_features					)					)
         self.parent.assertListEqual(list(result.feature_maps[0].shape					)  ,     [self.batch_size, self.image_size, 8, 8]					)

         # verify channels
         self.parent.assertEqual(len(model.channels					)  ,     len(config.out_features					)					)
         self.parent.assertListEqual(model.channels  ,     config.hidden_sizes[:-1]					)

         # verify backbone works with out_features=None
         UpperCAmelCase_			       =  None
         UpperCAmelCase_			       =  FocalNetBackbone(config=__a					)
         model.to(__a					)
         model.eval()
         UpperCAmelCase_			       =  model(__a					)

         # verify feature maps
         self.parent.assertEqual(len(result.feature_maps					)  ,     1					)
         self.parent.assertListEqual(list(result.feature_maps[0].shape					)  ,     [self.batch_size, self.image_size * 2, 4, 4]					)

         # verify channels
         self.parent.assertEqual(len(model.channels					)  ,     1					)
         self.parent.assertListEqual(model.channels  ,     [config.hidden_sizes[-1]]					)



    def    _lowercase		(self       :       Tuple  ,     __a       :       Tuple  ,     __a       :       Optional[int]  ,     __a       :       int					):
         UpperCAmelCase_			       =  FocalNetForMaskedImageModeling(config=__a					)
         model.to(__a					)
         model.eval()
         UpperCAmelCase_			       =  model(__a					)
         self.parent.assertEqual(
             result.reconstruction.shape  ,     (self.batch_size, self.num_channels, self.image_size, self.image_size)					)

         # test greyscale images
         UpperCAmelCase_			       =  1
         UpperCAmelCase_			       =  FocalNetForMaskedImageModeling(__a					)
         model.to(__a					)
         model.eval()

         UpperCAmelCase_			       =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
         UpperCAmelCase_			       =  model(__a					)
         self.parent.assertEqual(result.reconstruction.shape  ,     (self.batch_size, 1, self.image_size, self.image_size)					)



    def    _lowercase		(self       :       Any  ,     __a       :       int  ,     __a       :       List[Any]  ,     __a       :       Tuple					):
         UpperCAmelCase_			       =  self.type_sequence_label_size
         UpperCAmelCase_			       =  FocalNetForImageClassification(__a					)
         model.to(__a					)
         model.eval()
         UpperCAmelCase_			       =  model(__a  ,     labels=__a					)
         self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.type_sequence_label_size)					)

         # test greyscale images
         UpperCAmelCase_			       =  1
         UpperCAmelCase_			       =  FocalNetForImageClassification(__a					)
         model.to(__a					)
         model.eval()

         UpperCAmelCase_			       =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
         UpperCAmelCase_			       =  model(__a					)
         self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.type_sequence_label_size)					)






    def    _lowercase		(self       :       List[str]					):
         UpperCAmelCase_			       =  self.prepare_config_and_inputs()

         UpperCAmelCase_       ,					UpperCAmelCase_       ,					UpperCAmelCase_			       =  config_and_inputs
         UpperCAmelCase_			       =  {"pixel_values": pixel_values}
         return config, inputs_dict




@require_torch
class 			__A (					UpperCamelCase__						,      UpperCamelCase__						,      unittest.TestCase      ):
    a__			:		Tuple			  =    (
        (
            FocalNetModel,
            FocalNetForImageClassification,
            FocalNetForMaskedImageModeling,
            FocalNetBackbone,
        )
        if is_torch_available()
        else ()
    )
    a__			:		List[str]			  =    (
        {"""feature-extraction""": FocalNetModel, """image-classification""": FocalNetForImageClassification}
        if is_torch_available()
        else {}
    )
    a__			:		Union[str, Any]			  =    False

    a__			:		str			  =    False
    a__			:		List[str]			  =    False
    a__			:		int			  =    False
    a__			:		Optional[Any]			  =    False



    def    _lowercase		(self       :       List[Any]					):
         UpperCAmelCase_			       =  FocalNetModelTester(self					)
         UpperCAmelCase_			       =  ConfigTester(self  ,     config_class=__a  ,     embed_dim=37  ,     has_text_modality=__a					)



    def    _lowercase		(self       :       List[str]					):
         self.create_and_test_config_common_properties()
         self.config_tester.create_and_test_config_to_json_string()
         self.config_tester.create_and_test_config_to_json_file()
         self.config_tester.create_and_test_config_from_and_save_pretrained()
         self.config_tester.create_and_test_config_with_num_labels()
         self.config_tester.check_config_can_be_init_without_params()
         self.config_tester.check_config_arguments_init()



    def    _lowercase		(self       :       Optional[Any]					):
         return



    def    _lowercase		(self       :       Any					):
         UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__a					)



    def    _lowercase		(self       :       Dict					):
         UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_backbone(*__a					)



    def    _lowercase		(self       :       Tuple					):
         UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_image_modeling(*__a					)



    def    _lowercase		(self       :       Optional[int]					):
         UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__a					)



    @unittest.skip(reason="FocalNet does not use inputs_embeds"					)
    def    _lowercase		(self       :       str					):
         pass



    @unittest.skip(reason="FocalNet does not use feedforward chunking"					)
    def    _lowercase		(self       :       Dict					):
         pass



    def    _lowercase		(self       :       Dict					):
         UpperCAmelCase_       ,					UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes[:-1]:
              UpperCAmelCase_			       =  model_class(__a					)
              self.assertIsInstance(model.get_input_embeddings()  ,     (nn.Module)					)
              UpperCAmelCase_			       =  model.get_output_embeddings()
              self.assertTrue(x is None or isinstance(__a  ,     nn.Linear					)					)



    def    _lowercase		(self       :       str					):
         UpperCAmelCase_       ,					UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes[:-1]:
              UpperCAmelCase_			       =  model_class(__a					)
              UpperCAmelCase_			       =  inspect.signature(model.forward					)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              UpperCAmelCase_			       =  [*signature.parameters.keys()]

              UpperCAmelCase_			       =  ["pixel_values"]
              self.assertListEqual(arg_names[:1]  ,     __a					)



    def    _lowercase		(self       :       Union[str, Any]  ,     __a       :       Dict  ,     __a       :       List[str]  ,     __a       :       Tuple  ,     __a       :       Optional[int]					):
         UpperCAmelCase_			       =  model_class(__a					)
         model.to(__a					)
         model.eval()

         with torch.no_grad():
              UpperCAmelCase_			       =  model(**self._prepare_for_class(__a  ,     __a					)					)

         UpperCAmelCase_			       =  outputs.hidden_states

         UpperCAmelCase_			       =  getattr(
             self.model_tester  ,     "expected_num_hidden_layers"  ,     len(self.model_tester.depths					) + 1					)
         self.assertEqual(len(__a					)  ,     __a					)

         # FocalNet has a different seq_length
         UpperCAmelCase_			       =  (
             config.patch_size
             if isinstance(config.patch_size  ,     collections.abc.Iterable					)
             else (config.patch_size, config.patch_size)
         )

         UpperCAmelCase_			       =  (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

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

         UpperCAmelCase_			       =  outputs.reshaped_hidden_states
         self.assertEqual(len(__a					)  ,     __a					)

         UpperCAmelCase_       ,					UpperCAmelCase_       ,					UpperCAmelCase_       ,					UpperCAmelCase_			       =  reshaped_hidden_states[0].shape
         UpperCAmelCase_			       =  (
             reshaped_hidden_states[0].view(__a  ,     __a  ,     height * width					).permute(0  ,     2  ,     1					)
         )
         self.assertListEqual(
             list(reshaped_hidden_states.shape[-2:]					)  ,     [num_patches, self.model_tester.embed_dim]  ,     )



    def    _lowercase		(self       :       Union[str, Any]					):
         UpperCAmelCase_       ,					UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs_for_common()

         UpperCAmelCase_			       =  (
             self.model_tester.image_size
             if isinstance(self.model_tester.image_size  ,     collections.abc.Iterable					)
             else (self.model_tester.image_size, self.model_tester.image_size)
         )

         for model_class in self.all_model_classes[:-1]:
              UpperCAmelCase_			       =  True
              self.check_hidden_states_output(__a  ,     __a  ,     __a  ,     __a					)

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              UpperCAmelCase_			       =  True

              self.check_hidden_states_output(__a  ,     __a  ,     __a  ,     __a					)



    def    _lowercase		(self       :       Optional[int]					):
         UpperCAmelCase_       ,					UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase_			       =  3

         UpperCAmelCase_			       =  (
             self.model_tester.image_size
             if isinstance(self.model_tester.image_size  ,     collections.abc.Iterable					)
             else (self.model_tester.image_size, self.model_tester.image_size)
         )
         UpperCAmelCase_			       =  (
             config.patch_size
             if isinstance(config.patch_size  ,     collections.abc.Iterable					)
             else (config.patch_size, config.patch_size)
         )

         UpperCAmelCase_			       =  image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
         UpperCAmelCase_			       =  image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

         for model_class in self.all_model_classes[:-1]:
              UpperCAmelCase_			       =  True
              self.check_hidden_states_output(__a  ,     __a  ,     __a  ,     (padded_height, padded_width)					)

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              UpperCAmelCase_			       =  True
              self.check_hidden_states_output(__a  ,     __a  ,     __a  ,     (padded_height, padded_width)					)



    @slow
    def    _lowercase		(self       :       Optional[int]					):
         for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              UpperCAmelCase_			       =  FocalNetModel.from_pretrained(__a					)
              self.assertIsNotNone(__a					)






    def    _lowercase		(self       :       List[Any]					):
         UpperCAmelCase_       ,					UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs_for_common()

         UpperCAmelCase_			       =  _config_zero_init(__a					)
         for model_class in self.all_model_classes:
              UpperCAmelCase_			       =  model_class(config=__a					)
              for name, param in model.named_parameters():
                   if "embeddings" not in name and param.requires_grad:
                        self.assertIn(
                            ((param.data.mean() * 1E9).round() / 1E9).item()  ,     [0.0, 1.0]  ,     msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""  ,     )




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



    @cached_property
    def    _lowercase		(self       :       Optional[int]					):
         # TODO update organization
         return AutoImageProcessor.from_pretrained("microsoft/focalnet-tiny"					) if is_vision_available() else None






    @slow
    def    _lowercase		(self       :       Tuple					):
         UpperCAmelCase_			       =  FocalNetForImageClassification.from_pretrained("microsoft/focalnet-tiny"					).to(__a					)
         UpperCAmelCase_			       =  self.default_image_processor

         UpperCAmelCase_			       =  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"					)
         UpperCAmelCase_			       =  image_processor(images=__a  ,     return_tensors="pt"					).to(__a					)

         # forward pass
         with torch.no_grad():
              UpperCAmelCase_			       =  model(**__a					)

         # verify the logits
         UpperCAmelCase_			       =  torch.Size((1, 1000)					)
         self.assertEqual(outputs.logits.shape  ,     __a					)
         UpperCAmelCase_			       =  torch.tensor([0.21_66, -0.43_68, 0.21_91]					).to(__a					)
         self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,     __a  ,     atol=1E-4					)					)
         self.assertTrue(outputs.logits.argmax(dim=-1					).item()  ,     281					)




@require_torch
class 			__A (					UpperCamelCase__						,      unittest.TestCase      ):
    a__			:		Tuple			  =    (FocalNetBackbone,) if is_torch_available() else ()
    a__			:		int			  =    FocalNetConfig

    a__			:		List[Any]			  =    False



    def    _lowercase		(self       :       Dict					):
         UpperCAmelCase_			       =  FocalNetModelTester(self					)





import json
import multiprocessing
import os
import re
from collections import defaultdict

import torch
from accelerate import Accelerator
from accelerate.utils import set_seed
from arguments import HumanEvalArguments
from datasets import load_dataset, load_metric
from torch.utils.data import IterableDataset
from torch.utils.data.dataloader import DataLoader
from tqdm import tqdm

import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, StoppingCriteria, StoppingCriteriaList


SCREAMING_SNAKE_CASE		:     Optional[int]         =	["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif"]

class        UpperCamelCase						(				__a				):





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





						def __iter__(self						)				->				Any:
												UpperCamelCase_		:							int  =		[]
												for task in range(self.n_tasks						):
																		# without strip, the model generate commented codes ...
																		prompts.append(self.tokenizer.eos_token + self.dataset[task]["""prompt"""].strip()						)
												UpperCamelCase_		:							Any  =		self.tokenizer(__UpperCamelCase			,     padding=__UpperCamelCase			,     return_tensors="""pt"""						)
												for task in range(self.n_tasks						):
																		for _ in range(self.n_copies						):
																								yield {
																								    "ids": outputs.input_ids[task],
																								    "task_id": task,
																								    "input_len": outputs.attention_mask[task].sum(),
																								}


class        UpperCamelCase						(				__a				):





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





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






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






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

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

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

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






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

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

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

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

						# Load model and tokenizer
						UpperCamelCase_		:							str  =		AutoTokenizer.from_pretrained(args.model_ckpt			)
						UpperCamelCase_		:							Any  =		tokenizer.eos_token
						UpperCamelCase_		:							Tuple  =		AutoModelForCausalLM.from_pretrained(args.model_ckpt			)

						# Generation settings
						UpperCamelCase_		:							Union[str, Any]  =		{
						    """do_sample""": args.do_sample,
						    """temperature""": args.temperature,
						    """max_new_tokens""": args.max_new_tokens,
						    """top_p""": args.top_p,
						    """top_k""": args.top_k,
						    """stopping_criteria""": StoppingCriteriaList([EndOfFunctionCriteria(0		, _SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE			)]			),
						}

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

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

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

						# Run a quick test to see if code evaluation is enabled
						try:
												UpperCamelCase_		:							str  =		code_eval_metric.compute(references=[""""""]		, predictions=[[""""""]]			)
						except ValueError as exception:
												print(
												    """Code evaluation not enabled. Read the warning below carefully and then use `--HF_ALLOW_CODE_EVAL=\"1\"`"""
												    """ flag to enable code evaluation."""			)
												raise exception

						UpperCamelCase_,UpperCamelCase_		:							int  =		accelerator.prepare(_SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE			)

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

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

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

												# Evaluate completions with "code_eval" metric
												UpperCamelCase_,UpperCamelCase_		:							Any  =		code_eval_metric.compute(
												    references=_SCREAMING_SNAKE_CASE		, predictions=_SCREAMING_SNAKE_CASE		, num_workers=args.num_workers			)
												print(f'''Results: {pass_at_k}'''			)

												# Save results to json file
												with open(args.output_file		, """w"""			) as fp:
																		json.dump(_SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE			)


# For some reason the folliwng seems to be necessary sometimes for code_eval to work nice with multiprocessing
# https://stackoverflow.com/questions/60804599/python-multiprocessing-keeps-spawning-the-whole-script
if __name__ == "__main__":
	main()






