
import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        GPTNeoXForCausalLM,
        GPTNeoXForQuestionAnswering,
        GPTNeoXForSequenceClassification,
        GPTNeoXForTokenClassification,
        GPTNeoXModel,
    )




class        UpperCAmelCase_			:


  '''simple docstring'''







  def __init__( self			: Union[str, Any]			,							UpperCamelCase__			: Dict			,							UpperCamelCase__			: List[Any]=13			,							UpperCamelCase__			: int=7			,							UpperCamelCase__			: Tuple=True			,							UpperCamelCase__			: int=True			,							UpperCamelCase__			: Optional[int]=True			,							UpperCamelCase__			: List[str]=True			,							UpperCamelCase__			: int=99			,							UpperCamelCase__			: List[Any]=64			,							UpperCamelCase__			: List[Any]=5			,							UpperCamelCase__			: int=4			,							UpperCamelCase__			: Tuple=37			,							UpperCamelCase__			: Optional[int]="gelu"			,							UpperCamelCase__			: Optional[Any]=0.1			,							UpperCamelCase__			: List[str]=0.1			,							UpperCamelCase__			: int=512			,							UpperCamelCase__			: Any=16			,							UpperCamelCase__			: Optional[Any]=2			,							UpperCamelCase__			: List[str]=0.02			,							UpperCamelCase__			: Dict=3			,							UpperCamelCase__			: List[Any]=4			,							UpperCamelCase__			: str=None			,							)			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       parent
    __magic_name__				     =       batch_size
    __magic_name__				     =       seq_length
    __magic_name__				     =       is_training
    __magic_name__				     =       use_input_mask
    __magic_name__				     =       use_token_type_ids
    __magic_name__				     =       use_labels
    __magic_name__				     =       vocab_size
    __magic_name__				     =       hidden_size
    __magic_name__				     =       num_hidden_layers
    __magic_name__				     =       num_attention_heads
    __magic_name__				     =       intermediate_size
    __magic_name__				     =       hidden_act
    __magic_name__				     =       hidden_dropout_prob
    __magic_name__				     =       attention_probs_dropout_prob
    __magic_name__				     =       max_position_embeddings
    __magic_name__				     =       type_vocab_size
    __magic_name__				     =       type_sequence_label_size
    __magic_name__				     =       initializer_range
    __magic_name__				     =       num_labels
    __magic_name__				     =       num_choices
    __magic_name__				     =       scope
    __magic_name__				     =       vocab_size - 1
  def 							_lowercase							( self			: str    )			->	str:







    """simple docstring"""






    __magic_name__				     =       ids_tensor([self.batch_size, self.seq_length]			,							self.vocab_size    )

    __magic_name__				     =       None
    if self.use_input_mask:
      __magic_name__				     =       random_attention_mask([self.batch_size, self.seq_length]    )

    __magic_name__				     =       None
    if self.use_labels:
      __magic_name__				     =       ids_tensor([self.batch_size, self.seq_length]			,							self.num_labels    )

    __magic_name__				     =       self.get_config()

    return config, input_ids, input_mask, token_labels
  def 							_lowercase							( self			: int    )			->	int:







    """simple docstring"""






    return GPTNeoXConfig(
        vocab_size=self.vocab_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			,							max_position_embeddings=self.max_position_embeddings			,							type_vocab_size=self.type_vocab_size			,							is_decoder=UpperCamelCase__			,							initializer_range=self.initializer_range			,							pad_token_id=self.pad_token_id			,							)
  def 							_lowercase							( self			: Any    )			->	Tuple:







    """simple docstring"""






    __magic_name__     ,    __magic_name__     ,    __magic_name__     ,    __magic_name__				     =       self.prepare_config_and_inputs()

    __magic_name__				     =       True

    return config, input_ids, input_mask, token_labels
  def 							_lowercase							( self			: Any			,							UpperCamelCase__			: Optional[int]			,							UpperCamelCase__			: Dict			,							UpperCamelCase__			: Optional[Any]    )			->	Optional[Any]:







    """simple docstring"""






    __magic_name__				     =       GPTNeoXModel(config=UpperCamelCase__    )
    model.to(UpperCamelCase__    )
    model.eval()
    __magic_name__				     =       model(UpperCamelCase__			,							attention_mask=UpperCamelCase__    )
    __magic_name__				     =       model(UpperCamelCase__    )
    self.parent.assertEqual(result.last_hidden_state.shape			,							(self.batch_size, self.seq_length, self.hidden_size)    )
  def 							_lowercase							( self			: List[Any]			,							UpperCamelCase__			: int			,							UpperCamelCase__			: int			,							UpperCamelCase__			: List[Any]    )			->	Any:







    """simple docstring"""






    __magic_name__				     =       True
    __magic_name__				     =       GPTNeoXModel(UpperCamelCase__    )
    model.to(UpperCamelCase__    )
    model.eval()
    __magic_name__				     =       model(UpperCamelCase__			,							attention_mask=UpperCamelCase__    )
    self.parent.assertEqual(result.last_hidden_state.shape			,							(self.batch_size, self.seq_length, self.hidden_size)    )
  def 							_lowercase							( self			: int			,							UpperCamelCase__			: Union[str, Any]			,							UpperCamelCase__			: List[Any]			,							UpperCamelCase__			: Any			,							UpperCamelCase__			: Dict    )			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       GPTNeoXForCausalLM(config=UpperCamelCase__    )
    model.to(UpperCamelCase__    )
    model.eval()
    __magic_name__				     =       model(UpperCamelCase__			,							attention_mask=UpperCamelCase__			,							labels=UpperCamelCase__    )
    self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.seq_length, self.vocab_size)    )
  def 							_lowercase							( self			: Any			,							UpperCamelCase__			: Dict			,							UpperCamelCase__			: Any			,							UpperCamelCase__			: str			,							UpperCamelCase__			: List[str]    )			->	Optional[Any]:







    """simple docstring"""






    __magic_name__				     =       self.num_labels
    __magic_name__				     =       GPTNeoXForQuestionAnswering(UpperCamelCase__    )
    model.to(UpperCamelCase__    )
    model.eval()
    __magic_name__				     =       model(UpperCamelCase__			,							attention_mask=UpperCamelCase__    )
    self.parent.assertEqual(result.start_logits.shape			,							(self.batch_size, self.seq_length)    )
    self.parent.assertEqual(result.end_logits.shape			,							(self.batch_size, self.seq_length)    )
  def 							_lowercase							( self			: Tuple			,							UpperCamelCase__			: Dict			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: Dict			,							UpperCamelCase__			: Tuple    )			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       self.num_labels
    __magic_name__				     =       GPTNeoXForSequenceClassification(UpperCamelCase__    )
    model.to(UpperCamelCase__    )
    model.eval()
    __magic_name__				     =       ids_tensor([self.batch_size]			,							self.type_sequence_label_size    )
    __magic_name__				     =       model(UpperCamelCase__			,							attention_mask=UpperCamelCase__			,							labels=UpperCamelCase__    )
    self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.num_labels)    )
  def 							_lowercase							( self			: List[str]			,							UpperCamelCase__			: List[Any]			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: List[Any]			,							UpperCamelCase__			: int    )			->	Tuple:







    """simple docstring"""






    __magic_name__				     =       self.num_labels
    __magic_name__				     =       GPTNeoXForTokenClassification(UpperCamelCase__    )
    model.to(UpperCamelCase__    )
    model.eval()
    __magic_name__				     =       model(UpperCamelCase__			,							attention_mask=UpperCamelCase__			,							labels=UpperCamelCase__    )
    self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.seq_length, self.num_labels)    )
  def 							_lowercase							( self			: Union[str, Any]			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: List[str]			,							UpperCamelCase__			: Dict    )			->	Dict:







    """simple docstring"""






    __magic_name__				     =       True
    __magic_name__				     =       GPTNeoXForCausalLM(config=UpperCamelCase__    )
    model.to(UpperCamelCase__    )
    model.eval()

    # first forward pass
    __magic_name__				     =       model(UpperCamelCase__			,							attention_mask=UpperCamelCase__			,							use_cache=UpperCamelCase__    )
    __magic_name__				     =       outputs.past_key_values

    # create hypothetical multiple next token and extent to next_input_ids
    __magic_name__				     =       ids_tensor((self.batch_size, 3)			,							config.vocab_size    )
    __magic_name__				     =       ids_tensor((self.batch_size, 3)			,							vocab_size=2    )

    # append to next input_ids and
    __magic_name__				     =       torch.cat([input_ids, next_tokens]			,							dim=-1    )
    __magic_name__				     =       torch.cat([input_mask, next_mask]			,							dim=-1    )

    __magic_name__				     =       model(UpperCamelCase__			,							attention_mask=UpperCamelCase__			,							output_hidden_states=UpperCamelCase__    )
    __magic_name__				     =       output_from_no_past["""hidden_states"""][0]
    __magic_name__				     =       model(
        UpperCamelCase__			,							attention_mask=UpperCamelCase__			,							past_key_values=UpperCamelCase__			,							output_hidden_states=UpperCamelCase__			,							)["""hidden_states"""][0]

    # select random slice
    __magic_name__				     =       ids_tensor((1,)			,							output_from_past.shape[-1]    ).item()
    __magic_name__				     =       output_from_no_past[:, -3:, random_slice_idx].detach()
    __magic_name__				     =       output_from_past[:, :, random_slice_idx].detach()

    self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]    )

    # test that outputs are equal for slice
    self.parent.assertTrue(torch.allclose(UpperCamelCase__			,							UpperCamelCase__			,							atol=1E-3    )    )





  def 							_lowercase							( self			: int    )			->	Optional[int]:







    """simple docstring"""






    __magic_name__				     =       self.prepare_config_and_inputs()
    __magic_name__     ,    __magic_name__     ,    __magic_name__     ,    __magic_name__				     =       config_and_inputs
    __magic_name__				     =       {"""input_ids""": input_ids, """attention_mask""": input_mask}
    return config, inputs_dict







@require_torch
class        UpperCAmelCase_			( _A  ,       _A  ,       _A  ,       unittest.TestCase     ):


  '''simple docstring'''







  a__           =		(
      (
          GPTNeoXModel,
          GPTNeoXForCausalLM,
          GPTNeoXForQuestionAnswering,
          GPTNeoXForSequenceClassification,
          GPTNeoXForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  a__           =		(GPTNeoXForCausalLM,) if is_torch_available() else ()
  a__           =		(
      {
          """feature-extraction""": GPTNeoXModel,
          """question-answering""": GPTNeoXForQuestionAnswering,
          """text-classification""": GPTNeoXForSequenceClassification,
          """text-generation""": GPTNeoXForCausalLM,
          """token-classification""": GPTNeoXForTokenClassification,
          """zero-shot""": GPTNeoXForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  a__           =		False
  a__           =		False
  a__           =		False
  a__           =		False
  def 							_lowercase							( self			: Optional[int]    )			->	Union[str, Any]:







    """simple docstring"""






    __magic_name__				     =       GPTNeoXModelTester(self    )
    __magic_name__				     =       ConfigTester(self			,							config_class=UpperCamelCase__			,							hidden_size=64			,							num_attention_heads=8    )
  def 							_lowercase							( self			: List[str]    )			->	List[str]:







    """simple docstring"""






    self.config_tester.run_common_tests()
  def 							_lowercase							( self			: str    )			->	List[str]:







    """simple docstring"""






    __magic_name__     ,    __magic_name__     ,    __magic_name__     ,    __magic_name__				     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(UpperCamelCase__			,							UpperCamelCase__			,							UpperCamelCase__    )
  def 							_lowercase							( self			: str    )			->	str:







    """simple docstring"""






    __magic_name__     ,    __magic_name__     ,    __magic_name__     ,    __magic_name__				     =       self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_model_as_decoder(UpperCamelCase__			,							UpperCamelCase__			,							UpperCamelCase__    )
  def 							_lowercase							( self			: Tuple    )			->	List[str]:







    """simple docstring"""






    __magic_name__     ,    __magic_name__     ,    __magic_name__     ,    __magic_name__				     =       self.model_tester.prepare_config_and_inputs_for_decoder()

    __magic_name__				     =       None

    self.model_tester.create_and_check_model_as_decoder(UpperCamelCase__			,							UpperCamelCase__			,							UpperCamelCase__    )
  def 							_lowercase							( self			: List[str]    )			->	Any:







    """simple docstring"""






    __magic_name__     ,    __magic_name__     ,    __magic_name__     ,    __magic_name__				     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_decoder_model_past_large_inputs(UpperCamelCase__			,							UpperCamelCase__			,							UpperCamelCase__    )
  def 							_lowercase							( self			: int    )			->	Any:







    """simple docstring"""






    __magic_name__				     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_causal_lm(*UpperCamelCase__    )
  def 							_lowercase							( self			: Optional[Any]    )			->	Dict:







    """simple docstring"""






    __magic_name__				     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_question_answering(*UpperCamelCase__    )
  def 							_lowercase							( self			: List[Any]    )			->	Union[str, Any]:







    """simple docstring"""






    __magic_name__				     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase__    )
  def 							_lowercase							( self			: List[str]    )			->	List[Any]:







    """simple docstring"""






    __magic_name__				     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__    )
  @unittest.skip(reason="""Feed forward chunking is not implemented"""    )
  def 							_lowercase							( self			: List[Any]    )			->	List[str]:







    """simple docstring"""






    pass





  @parameterized.expand([("""linear""",), ("""dynamic""",)]    )
  def 							_lowercase							( self			: List[str]			,							UpperCamelCase__			: Tuple    )			->	str:







    """simple docstring"""






    __magic_name__     ,    __magic_name__				     =       self.model_tester.prepare_config_and_inputs_for_common()
    __magic_name__				     =       ids_tensor([1, 10]			,							config.vocab_size    )
    __magic_name__				     =       ids_tensor([1, int(config.max_position_embeddings * 1.5    )]			,							config.vocab_size    )

    set_seed(42    )  # Fixed seed at init time so the two models get the same random weights
    __magic_name__				     =       GPTNeoXModel(UpperCamelCase__    )
    original_model.to(UpperCamelCase__    )
    original_model.eval()
    __magic_name__				     =       original_model(UpperCamelCase__    ).last_hidden_state
    __magic_name__				     =       original_model(UpperCamelCase__    ).last_hidden_state

    set_seed(42    )  # Fixed seed at init time so the two models get the same random weights
    __magic_name__				     =       {"""type""": scaling_type, """factor""": 10.0}
    __magic_name__				     =       GPTNeoXModel(UpperCamelCase__    )
    scaled_model.to(UpperCamelCase__    )
    scaled_model.eval()
    __magic_name__				     =       scaled_model(UpperCamelCase__    ).last_hidden_state
    __magic_name__				     =       scaled_model(UpperCamelCase__    ).last_hidden_state

    # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
    # maximum sequence length, so the outputs for the short input should match.
    if scaling_type == "dynamic":
      self.assertTrue(torch.allclose(UpperCamelCase__			,							UpperCamelCase__			,							atol=1E-5    )    )
    else:
      self.assertFalse(torch.allclose(UpperCamelCase__			,							UpperCamelCase__			,							atol=1E-5    )    )

    # The output should be different for long inputs
    self.assertFalse(torch.allclose(UpperCamelCase__			,							UpperCamelCase__			,							atol=1E-5    )    )







@require_torch
class        UpperCAmelCase_			( unittest.TestCase     ):


  '''simple docstring'''







  @slow
  def 							_lowercase							( self			: List[str]    )			->	Optional[int]:







    """simple docstring"""






    __magic_name__				     =       AutoTokenizer.from_pretrained("""EleutherAI/pythia-410m-deduped"""    )
    for checkpointing in [True, False]:
      __magic_name__				     =       GPTNeoXForCausalLM.from_pretrained("""EleutherAI/pythia-410m-deduped"""    )

      if checkpointing:
        model.gradient_checkpointing_enable()
      else:
        model.gradient_checkpointing_disable()
      model.to(UpperCamelCase__    )

      __magic_name__				     =       tokenizer("""My favorite food is"""			,							return_tensors="""pt"""    ).to(UpperCamelCase__    )
      # The hub repo. is updated on 2023-04-04, resulting in poor outputs.
      # See: https://github.com/huggingface/transformers/pull/24193
      __magic_name__				     =       """My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI'm not sure"""

      __magic_name__				     =       model.generate(**UpperCamelCase__			,							do_sample=UpperCamelCase__			,							max_new_tokens=20    )
      __magic_name__				     =       tokenizer.batch_decode(UpperCamelCase__    )[0]

      self.assertEqual(UpperCamelCase__			,							UpperCamelCase__    )



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			:     Optional[int]           =   ['\nclass', '\ndef', '\n#', '\n@', '\nprint', '\nif']




class        UpperCAmelCase_			( _A     ):


  '''simple docstring'''







  def __init__( self			: Union[str, Any]			,							UpperCamelCase__			: List[Any]			,							UpperCamelCase__			: Optional[int]			,							UpperCamelCase__			: Any=None			,							UpperCamelCase__			: Union[str, Any]=1    )			->	str:







    """simple docstring"""






    __magic_name__				     =       tokenizer
    __magic_name__				     =       dataset
    __magic_name__				     =       len(UpperCamelCase__    ) if n_tasks is None else n_tasks
    __magic_name__				     =       n_copies





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







    """simple docstring"""






    __magic_name__				     =       []
    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()    )
    __magic_name__				     =       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     ):


  '''simple docstring'''







  def __init__( self			: Union[str, Any]			,							UpperCamelCase__			: Dict			,							UpperCamelCase__			: str			,							UpperCamelCase__			: str    )			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       start_length
    __magic_name__				     =       eof_strings
    __magic_name__				     =       tokenizer





  def __call__( self			: List[Any]			,							UpperCamelCase__			: Tuple			,							UpperCamelCase__			: Optional[Any]			,							**UpperCamelCase__			: Optional[int]    )			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       self.tokenizer.batch_decode(input_ids[:, self.start_length :]    )
    __magic_name__				     =       []
    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 a__    (       A_					):






  '''simple docstring'''





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

def a__    (       A_,     A_,     A_,     A_,     A_,     A_=20,     **A_					):






  '''simple docstring'''





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

      __magic_name__     ,    __magic_name__				     =       accelerator.gather((generated_tokens, generated_tasks)					)
      __magic_name__				     =       generated_tokens.cpu().numpy()
      __magic_name__				     =       generated_tasks.cpu().numpy()

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

  __magic_name__				     =       [[] for _ in range(A_					)]
  for task, generated_tokens in gen_token_dict.items():
    for s in generated_tokens:
      __magic_name__				     =       tokenizer.decode(A_,     skip_special_tokens=A_,     clean_up_tokenization_spaces=A_					)
      code_gens[task].append(remove_last_block(A_					)					)
  return code_gens

def a__    (       ):






  '''simple docstring'''





  __magic_name__				     =       HfArgumentParser(A_					)
  __magic_name__				     =       parser.parse_args()

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

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

  # Use dataset load to feed to accelerate
  __magic_name__				     =       Accelerator()
  set_seed(args.seed,     device_specific=A_					)

  # Load model and tokenizer
  __magic_name__				     =       AutoTokenizer.from_pretrained(args.model_ckpt					)
  __magic_name__				     =       tokenizer.eos_token
  __magic_name__				     =       AutoModelForCausalLM.from_pretrained(args.model_ckpt					)

  # Generation settings
  __magic_name__				     =       {
      """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,     A_,     A_					)]					),
  }

  # Load evaluation dataset and metric
  __magic_name__				     =       load_dataset("""openai_humaneval"""					)
  __magic_name__				     =       load_metric("""code_eval"""					)

  __magic_name__				     =       args.num_tasks if args.num_tasks is not None else len(human_eval["""test"""]					)
  __magic_name__				     =       args.n_samples // args.batch_size

  __magic_name__				     =       TokenizedDataset(A_,     human_eval["""test"""],     n_copies=A_,     n_tasks=A_					)
  # do not confuse args.batch_size, which is actually the num_return_sequences
  __magic_name__				     =       DataLoader(A_,     batch_size=1					)

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

  __magic_name__     ,    __magic_name__				     =       accelerator.prepare(A_,     A_					)

  __magic_name__				     =       complete_code(
      A_,     A_,     A_,     A_,     n_tasks=A_,     batch_size=args.batch_size,     **A_,     )

  if accelerator.is_main_process:
    __magic_name__				     =       []

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

    # Evaluate completions with "code_eval" metric
    __magic_name__     ,    __magic_name__				     =       code_eval_metric.compute(
        references=A_,     predictions=A_,     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(A_,     A_					)


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

