


'''simple docstring'''


import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
						import os

						# The slow tests are often failing with OOM error on GPU
						# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
						# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
						lowerCamelCase__    	=     'platform'

						import jax
						import jax.numpy as jnp

						from transformers import BlenderbotTokenizer
						from transformers.models.blenderbot.modeling_flax_blenderbot import (
						    FlaxBlenderbotForConditionalGeneration,
						    FlaxBlenderbotModel,
						    shift_tokens_right,
						)






def  __lowerCAmelCase (__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase=None      , __lowerCAmelCase=None      , __lowerCAmelCase=None      , __lowerCAmelCase=None      , __lowerCAmelCase=None      , __lowerCAmelCase=None      , ):
		if attention_mask is None:
				_UpperCAmelCase		:       Optional[Any]								=      np.where(input_ids != config.pad_token_id      , 1      , 0   )
		if decoder_attention_mask is None:
				_UpperCAmelCase		:       Union[str, Any]								=      np.where(decoder_input_ids != config.pad_token_id      , 1      , 0   )
		if head_mask is None:
				_UpperCAmelCase		:       List[Any]								=      np.ones((config.encoder_layers, config.encoder_attention_heads)   )
		if decoder_head_mask is None:
				_UpperCAmelCase		:       Union[str, Any]								=      np.ones((config.decoder_layers, config.decoder_attention_heads)   )
		if cross_attn_head_mask is None:
				_UpperCAmelCase		:       str								=      np.ones((config.decoder_layers, config.decoder_attention_heads)   )
		return {
		    "input_ids": input_ids,
		    "decoder_input_ids": decoder_input_ids,
		    "attention_mask": attention_mask,
		    "decoder_attention_mask": attention_mask,
		}






class 	lowerCAmelCase__			:




			def __init__(		self  :     Optional[int]	,  lowerCamelCase__  :     Dict	,  lowerCamelCase__  :     List[Any]=13	,  lowerCamelCase__  :     int=7	,  lowerCamelCase__  :     str=True	,  lowerCamelCase__  :     Union[str, Any]=False	,  lowerCamelCase__  :     Optional[Any]=99	,  lowerCamelCase__  :     str=16	,  lowerCamelCase__  :     str=2	,  lowerCamelCase__  :     Tuple=4	,  lowerCamelCase__  :     str=4	,  lowerCamelCase__  :     Optional[Any]="gelu"	,  lowerCamelCase__  :     List[str]=0.1	,  lowerCamelCase__  :     Any=0.1	,  lowerCamelCase__  :     List[Any]=32	,  lowerCamelCase__  :     str=2	,  lowerCamelCase__  :     str=1	,  lowerCamelCase__  :     List[str]=0	,  lowerCamelCase__  :     Optional[Any]=0.0_2	,  )   ->int:



					'''simple docstring'''







					_UpperCAmelCase		:       int								=      parent
					_UpperCAmelCase		:       List[str]								=      batch_size
					_UpperCAmelCase		:       int								=      seq_length
					_UpperCAmelCase		:       List[str]								=      is_training
					_UpperCAmelCase		:       List[str]								=      use_labels
					_UpperCAmelCase		:       str								=      vocab_size
					_UpperCAmelCase		:       Tuple								=      hidden_size
					_UpperCAmelCase		:       Tuple								=      num_hidden_layers
					_UpperCAmelCase		:       List[Any]								=      num_attention_heads
					_UpperCAmelCase		:       Optional[int]								=      intermediate_size
					_UpperCAmelCase		:       Any								=      hidden_act
					_UpperCAmelCase		:       Optional[Any]								=      hidden_dropout_prob
					_UpperCAmelCase		:       Tuple								=      attention_probs_dropout_prob
					_UpperCAmelCase		:       Optional[int]								=      max_position_embeddings
					_UpperCAmelCase		:       Optional[Any]								=      eos_token_id
					_UpperCAmelCase		:       Optional[Any]								=      pad_token_id
					_UpperCAmelCase		:       str								=      bos_token_id
					_UpperCAmelCase		:       Dict								=      initializer_range




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



					'''simple docstring'''







					_UpperCAmelCase		:       Optional[Any]								=      np.clip(ids_tensor([self.batch_size, self.seq_length - 1]	,  self.vocab_size )	,  3	,  self.vocab_size )
					_UpperCAmelCase		:       str								=      np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)	,  dtype=np.intaa ))	,  -1 )

					_UpperCAmelCase		:       List[str]								=      shift_tokens_right(lowerCamelCase__	,  1	,  2 )

					_UpperCAmelCase		:       Tuple								=      BlenderbotConfig(
					    vocab_size=self.vocab_size	,  d_model=self.hidden_size	,  encoder_layers=self.num_hidden_layers	,  decoder_layers=self.num_hidden_layers	,  encoder_attention_heads=self.num_attention_heads	,  decoder_attention_heads=self.num_attention_heads	,  encoder_ffn_dim=self.intermediate_size	,  decoder_ffn_dim=self.intermediate_size	,  dropout=self.hidden_dropout_prob	,  attention_dropout=self.attention_probs_dropout_prob	,  max_position_embeddings=self.max_position_embeddings	,  eos_token_id=self.eos_token_id	,  bos_token_id=self.bos_token_id	,  pad_token_id=self.pad_token_id	,  initializer_range=self.initializer_range	,  use_cache=lowerCamelCase__	,  )
					_UpperCAmelCase		:       Optional[Any]								=      prepare_blenderbot_inputs_dict(lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )
					return config, inputs_dict




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



					'''simple docstring'''







					_UpperCAmelCase					,    _UpperCAmelCase		:       Any								=      self.prepare_config_and_inputs()
					return config, inputs_dict




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



					'''simple docstring'''







					_UpperCAmelCase		:       List[str]								=      20
					_UpperCAmelCase		:       Optional[int]								=      model_class_name(lowerCamelCase__ )

					_UpperCAmelCase		:       int								=      model.encode(inputs_dict["input_ids"] )

					_UpperCAmelCase					,    _UpperCAmelCase		:       Union[str, Any]								=      (
					    inputs_dict["decoder_input_ids"],
					    inputs_dict["decoder_attention_mask"],
					)

					_UpperCAmelCase		:       Optional[int]								=      model.init_cache(decoder_input_ids.shape[0]	,  lowerCamelCase__	,  lowerCamelCase__ )
					_UpperCAmelCase		:       Union[str, Any]								=      jnp.ones((decoder_input_ids.shape[0], max_decoder_length)	,  dtype="i4" )

					_UpperCAmelCase		:       Any								=      jnp.broadcast_to(
					    jnp.arange(decoder_input_ids.shape[-1] - 1 )[None, :]	,  (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)	,  )
					_UpperCAmelCase		:       int								=      model.decode(
					    decoder_input_ids[:, :-1]	,  lowerCamelCase__	,  decoder_attention_mask=lowerCamelCase__	,  past_key_values=lowerCamelCase__	,  decoder_position_ids=lowerCamelCase__	,  )

					_UpperCAmelCase		:       Any								=      jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]	,  dtype="i4" )
					_UpperCAmelCase		:       List[Any]								=      model.decode(
					    decoder_input_ids[:, -1:]	,  lowerCamelCase__	,  decoder_attention_mask=lowerCamelCase__	,  past_key_values=outputs_cache.past_key_values	,  decoder_position_ids=lowerCamelCase__	,  )

					_UpperCAmelCase		:       Any								=      model.decode(lowerCamelCase__	,  lowerCamelCase__ )

					_UpperCAmelCase		:       Optional[int]								=      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
					self.parent.assertTrue(diff < 1E-3	,  msg=F"""Max diff is {diff}""" )






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



					'''simple docstring'''







					_UpperCAmelCase		:       Any								=      20
					_UpperCAmelCase		:       List[Any]								=      model_class_name(lowerCamelCase__ )

					_UpperCAmelCase		:       Union[str, Any]								=      model.encode(inputs_dict["input_ids"] )

					_UpperCAmelCase					,    _UpperCAmelCase		:       Optional[int]								=      (
					    inputs_dict["decoder_input_ids"],
					    inputs_dict["decoder_attention_mask"],
					)

					_UpperCAmelCase		:       Dict								=      jnp.concatenate(
					    [
					        decoder_attention_mask,
					        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1]) ),
					    ]	,  axis=-1	,  )

					_UpperCAmelCase		:       Tuple								=      model.init_cache(decoder_input_ids.shape[0]	,  lowerCamelCase__	,  lowerCamelCase__ )
					_UpperCAmelCase		:       List[str]								=      jnp.broadcast_to(
					    jnp.arange(decoder_input_ids.shape[-1] - 1 )[None, :]	,  (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)	,  )

					_UpperCAmelCase		:       List[str]								=      model.decode(
					    decoder_input_ids[:, :-1]	,  lowerCamelCase__	,  decoder_attention_mask=lowerCamelCase__	,  past_key_values=lowerCamelCase__	,  decoder_position_ids=lowerCamelCase__	,  )
					_UpperCAmelCase		:       str								=      jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]	,  dtype="i4" )
					_UpperCAmelCase		:       Optional[int]								=      model.decode(
					    decoder_input_ids[:, -1:]	,  lowerCamelCase__	,  past_key_values=outputs_cache.past_key_values	,  decoder_attention_mask=lowerCamelCase__	,  decoder_position_ids=lowerCamelCase__	,  )

					_UpperCAmelCase		:       Tuple								=      model.decode(lowerCamelCase__	,  lowerCamelCase__	,  decoder_attention_mask=lowerCamelCase__ )

					_UpperCAmelCase		:       Optional[Any]								=      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
					self.parent.assertTrue(diff < 1E-3	,  msg=F"""Max diff is {diff}""" )






@require_flax
class 	lowerCAmelCase__			(			unittest.TestCase	):
			lowerCAmelCase :	Tuple          =     99




			def 							lowerCAmelCase__    (		self  :     Dict )   ->Any:



					'''simple docstring'''







					_UpperCAmelCase		:       int								=      np.array(
					    [
					        [71, 82, 18, 33, 46, 91, 2],
					        [68, 34, 26, 58, 30, 82, 2],
					        [5, 97, 17, 39, 94, 40, 2],
					        [76, 83, 94, 25, 70, 78, 2],
					        [87, 59, 41, 35, 48, 66, 2],
					        [55, 13, 16, 58, 5, 2, 1],  # note padding
					        [64, 27, 31, 51, 12, 75, 2],
					        [52, 64, 86, 17, 83, 39, 2],
					        [48, 61, 9, 24, 71, 82, 2],
					        [26, 1, 60, 48, 22, 13, 2],
					        [21, 5, 62, 28, 14, 76, 2],
					        [45, 98, 37, 86, 59, 48, 2],
					        [70, 70, 50, 9, 28, 0, 2],
					    ]	,  dtype=np.intaa	,  )

					_UpperCAmelCase		:       Optional[Any]								=      input_ids.shape[0]
					_UpperCAmelCase		:       Optional[Any]								=      BlenderbotConfig(
					    vocab_size=self.vocab_size	,  d_model=24	,  encoder_layers=2	,  decoder_layers=2	,  encoder_attention_heads=2	,  decoder_attention_heads=2	,  encoder_ffn_dim=32	,  decoder_ffn_dim=32	,  max_position_embeddings=48	,  eos_token_id=2	,  pad_token_id=1	,  bos_token_id=0	,  )
					return config, input_ids, batch_size




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



					'''simple docstring'''







					_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase		:       Optional[Any]								=      self._get_config_and_data()
					_UpperCAmelCase		:       Dict								=      FlaxBlenderbotForConditionalGeneration(lowerCamelCase__ )
					_UpperCAmelCase		:       List[str]								=      lm_model(input_ids=lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[Any]								=      (batch_size, input_ids.shape[1], config.vocab_size)
					self.assertEqual(outputs["logits"].shape	,  lowerCamelCase__ )




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



					'''simple docstring'''







					_UpperCAmelCase		:       Any								=      BlenderbotConfig(
					    vocab_size=self.vocab_size	,  d_model=14	,  encoder_layers=2	,  decoder_layers=2	,  encoder_attention_heads=2	,  decoder_attention_heads=2	,  encoder_ffn_dim=8	,  decoder_ffn_dim=8	,  max_position_embeddings=48	,  )
					_UpperCAmelCase		:       List[str]								=      FlaxBlenderbotForConditionalGeneration(lowerCamelCase__ )
					_UpperCAmelCase		:       Union[str, Any]								=      np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]	,  dtype=np.intaa )
					_UpperCAmelCase		:       Union[str, Any]								=      np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]	,  dtype=np.intaa )
					_UpperCAmelCase		:       Union[str, Any]								=      lm_model(input_ids=lowerCamelCase__	,  decoder_input_ids=lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[Any]								=      (*summary.shape, config.vocab_size)
					self.assertEqual(outputs["logits"].shape	,  lowerCamelCase__ )






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



					'''simple docstring'''







					_UpperCAmelCase		:       Optional[Any]								=      np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]	,  dtype=np.intaa )
					_UpperCAmelCase		:       Any								=      shift_tokens_right(lowerCamelCase__	,  1	,  2 )
					_UpperCAmelCase		:       Union[str, Any]								=      np.equal(lowerCamelCase__	,  1 ).astype(np.floataa ).sum()
					_UpperCAmelCase		:       Optional[int]								=      np.equal(lowerCamelCase__	,  1 ).astype(np.floataa ).sum()
					self.assertEqual(shifted.shape	,  input_ids.shape )
					self.assertEqual(lowerCamelCase__	,  n_pad_before - 1 )
					self.assertTrue(np.equal(shifted[:, 0]	,  2 ).all() )






@require_flax
class 	lowerCAmelCase__			(			UpperCAmelCase__    ,       unittest.TestCase    ,       UpperCAmelCase__	):
			lowerCAmelCase :	int          =     True
			lowerCAmelCase :	Tuple          =     (
			    (
			        FlaxBlenderbotModel,
			        FlaxBlenderbotForConditionalGeneration,
			    )
			    if is_flax_available()
			    else ()
			)
			lowerCAmelCase :	List[str]          =     (FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()




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



					'''simple docstring'''







					_UpperCAmelCase		:       str								=      FlaxBlenderbotModelTester(self )




			def 							lowerCAmelCase__    (		self  :     Any )   ->int:



					'''simple docstring'''







					_UpperCAmelCase					,    _UpperCAmelCase		:       Optional[int]								=      self.model_tester.prepare_config_and_inputs()
					for model_class in self.all_model_classes:
							self.model_tester.check_use_cache_forward(lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )




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



					'''simple docstring'''







					_UpperCAmelCase					,    _UpperCAmelCase		:       Optional[Any]								=      self.model_tester.prepare_config_and_inputs()
					for model_class in self.all_model_classes:
							self.model_tester.check_use_cache_forward_with_attn_mask(lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )




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



					'''simple docstring'''







					_UpperCAmelCase					,    _UpperCAmelCase		:       Tuple								=      self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							with self.subTest(model_class.__name__ ):
									_UpperCAmelCase		:       List[str]								=      self._prepare_for_class(lowerCamelCase__	,  lowerCamelCase__ )
									_UpperCAmelCase		:       List[Any]								=      model_class(lowerCamelCase__ )

									@jax.jit
									def encode_jitted(lowerCamelCase__  :     Dict	,  lowerCamelCase__  :     str=None	,  **lowerCamelCase__  :     List[str] ):
											return model.encode(input_ids=lowerCamelCase__	,  attention_mask=lowerCamelCase__ )

									with self.subTest("JIT Enabled" ):
											_UpperCAmelCase		:       Any								=      encode_jitted(**lowerCamelCase__ ).to_tuple()

									with self.subTest("JIT Disabled" ):
											with jax.disable_jit():
													_UpperCAmelCase		:       Dict								=      encode_jitted(**lowerCamelCase__ ).to_tuple()

									self.assertEqual(len(lowerCamelCase__ )	,  len(lowerCamelCase__ ) )
									for jitted_output, output in zip(lowerCamelCase__	,  lowerCamelCase__ ):
											self.assertEqual(jitted_output.shape	,  output.shape )




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



					'''simple docstring'''







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

					for model_class in self.all_model_classes:
							with self.subTest(model_class.__name__ ):
									_UpperCAmelCase		:       Tuple								=      model_class(lowerCamelCase__ )
									_UpperCAmelCase		:       Dict								=      model.encode(inputs_dict["input_ids"]	,  inputs_dict["attention_mask"] )

									_UpperCAmelCase		:       str								=      {
									    "decoder_input_ids": inputs_dict["decoder_input_ids"],
									    "decoder_attention_mask": inputs_dict["decoder_attention_mask"],
									    "encoder_outputs": encoder_outputs,
									}

									@jax.jit
									def decode_jitted(lowerCamelCase__  :     Tuple	,  lowerCamelCase__  :     Union[str, Any]	,  lowerCamelCase__  :     int ):
											return model.decode(
											    decoder_input_ids=lowerCamelCase__	,  decoder_attention_mask=lowerCamelCase__	,  encoder_outputs=lowerCamelCase__	,  )

									with self.subTest("JIT Enabled" ):
											_UpperCAmelCase		:       Any								=      decode_jitted(**lowerCamelCase__ ).to_tuple()

									with self.subTest("JIT Disabled" ):
											with jax.disable_jit():
													_UpperCAmelCase		:       int								=      decode_jitted(**lowerCamelCase__ ).to_tuple()

									self.assertEqual(len(lowerCamelCase__ )	,  len(lowerCamelCase__ ) )
									for jitted_output, output in zip(lowerCamelCase__	,  lowerCamelCase__ ):
											self.assertEqual(jitted_output.shape	,  output.shape )




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



					'''simple docstring'''







					for model_class_name in self.all_model_classes:
							_UpperCAmelCase		:       Optional[Any]								=      model_class_name.from_pretrained("facebook/blenderbot-400M-distill" )
							# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
							_UpperCAmelCase		:       Optional[int]								=      np.ones((1, 1) ) * model.config.eos_token_id
							_UpperCAmelCase		:       int								=      model(lowerCamelCase__ )
							self.assertIsNotNone(lowerCamelCase__ )






			@unittest.skipUnless(jax_device != "cpu"	,  "3B test too slow on CPU." )
			@slow
			def 							lowerCAmelCase__    (		self  :     str )   ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       Any								=      {"num_beams": 1, "early_stopping": True, "min_length": 15, "max_length": 25}
					_UpperCAmelCase		:       Any								=      {"skip_special_tokens": True, "clean_up_tokenization_spaces": True}

					_UpperCAmelCase		:       Optional[int]								=      FlaxBlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-3B"	,  from_pt=lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[Any]								=      BlenderbotTokenizer.from_pretrained("facebook/blenderbot-3B" )

					_UpperCAmelCase		:       Tuple								=      ["Sam"]
					_UpperCAmelCase		:       Dict								=      tokenizer(lowerCamelCase__	,  return_tensors="jax" )

					_UpperCAmelCase		:       Tuple								=      model.generate(**lowerCamelCase__	,  **lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[Any]								=      "Sam is a great name. It means \"sun\" in Gaelic."

					_UpperCAmelCase		:       str								=      tokenizer.batch_decode(lowerCamelCase__	,  **lowerCamelCase__ )
					assert generated_txt[0].strip() == tgt_text







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



