
import unittest

from transformers.utils.backbone_utils import (
    BackboneMixin,
    get_aligned_output_features_output_indices,
    verify_out_features_out_indices,
)

class     _UpperCamelCase			(   unittest.TestCase					):






			'''simple docstring'''

			def  lowerCAmelCase__		(    self	:    Optional[int] ):
							UpperCamelCase_:							int		=  ["""a""", """b""", """c"""]

							# Defaults to last layer if both are None
							UpperCamelCase_,				UpperCamelCase_:							Union[str, Any]		=  get_aligned_output_features_output_indices(snake_case_  ,		snake_case_  ,		snake_case_ )
							self.assertEqual(snake_case_  ,		["""c"""] )
							self.assertEqual(snake_case_  ,		[2] )

							# Out indices set to match out features
							UpperCamelCase_,				UpperCamelCase_:							Union[str, Any]		=  get_aligned_output_features_output_indices(["""a""", """c"""]  ,		snake_case_  ,		snake_case_ )
							self.assertEqual(snake_case_  ,		["""a""", """c"""] )
							self.assertEqual(snake_case_  ,		[0, 2] )

							# Out features set to match out indices
							UpperCamelCase_,				UpperCamelCase_:							Optional[int]		=  get_aligned_output_features_output_indices(snake_case_  ,		[0, 2]  ,		snake_case_ )
							self.assertEqual(snake_case_  ,		["""a""", """c"""] )
							self.assertEqual(snake_case_  ,		[0, 2] )

							# Out features selected from negative indices
							UpperCamelCase_,				UpperCamelCase_:							Dict		=  get_aligned_output_features_output_indices(snake_case_  ,		[-3, -1]  ,		snake_case_ )
							self.assertEqual(snake_case_  ,		["""a""", """c"""] )
							self.assertEqual(snake_case_  ,		[-3, -1] )

			def  lowerCAmelCase__		(    self	:    int ):
							# Stage names must be set
							with self.assertRaises(snake_case_ ):
											verify_out_features_out_indices(["""a""", """b"""]  ,		(0, 1)  ,		snake_case_ )

							# Out features must be a list
							with self.assertRaises(snake_case_ ):
											verify_out_features_out_indices(("""a""", """b""")  ,		(0, 1)  ,		["""a""", """b"""] )

							# Out features must be a subset of stage names
							with self.assertRaises(snake_case_ ):
											verify_out_features_out_indices(["""a""", """b"""]  ,		(0, 1)  ,		["""a"""] )

							# Out indices must be a list or tuple
							with self.assertRaises(snake_case_ ):
											verify_out_features_out_indices(snake_case_  ,		0  ,		["""a""", """b"""] )

							# Out indices must be a subset of stage names
							with self.assertRaises(snake_case_ ):
											verify_out_features_out_indices(snake_case_  ,		(0, 1)  ,		["""a"""] )

							# Out features and out indices must be the same length
							with self.assertRaises(snake_case_ ):
											verify_out_features_out_indices(["""a""", """b"""]  ,		(0,)  ,		["""a""", """b""", """c"""] )

							# Out features should match out indices
							with self.assertRaises(snake_case_ ):
											verify_out_features_out_indices(["""a""", """b"""]  ,		(0, 2)  ,		["""a""", """b""", """c"""] )

							# Out features and out indices should be in order
							with self.assertRaises(snake_case_ ):
											verify_out_features_out_indices(["""b""", """a"""]  ,		(0, 1)  ,		["""a""", """b"""] )

							# Check passes with valid inputs
							verify_out_features_out_indices(["""a""", """b""", """d"""]  ,		(0, 1, -1)  ,		["""a""", """b""", """c""", """d"""] )




			def  lowerCAmelCase__		(    self	:    str ):
							UpperCamelCase_:							Any		=  BackboneMixin()

							UpperCamelCase_:							Optional[Any]		=  ["""a""", """b""", """c"""]
							UpperCamelCase_:							Dict		=  ["""a""", """c"""]
							UpperCamelCase_:							Optional[Any]		=  [0, 2]

							# Check that the output features and indices are set correctly
							self.assertEqual(backbone.out_features  ,		["""a""", """c"""] )
							self.assertEqual(backbone.out_indices  ,		[0, 2] )

							# Check out features and indices are updated correctly
							UpperCamelCase_:							Dict		=  ["""a""", """b"""]
							self.assertEqual(backbone.out_features  ,		["""a""", """b"""] )
							self.assertEqual(backbone.out_indices  ,		[0, 1] )

							UpperCamelCase_:							Dict		=  [-3, -1]
							self.assertEqual(backbone.out_features  ,		["""a""", """c"""] )
							self.assertEqual(backbone.out_indices  ,		[-3, -1] )








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_  :	str 				=					["""\nclass""", """\ndef""", """\n#""", """\n@""", """\nprint""", """\nif"""]

class     _UpperCamelCase			(   _A					):






			'''simple docstring'''

			def __init__(    self	:    Optional[int]  ,		snake_case_	:    List[str]  ,		snake_case_	:    Union[str, Any]  ,		snake_case_	:    Optional[Any]=None  ,		snake_case_	:    str=1 ):
							UpperCamelCase_:							List[str]		=  tokenizer
							UpperCamelCase_:							str		=  dataset
							UpperCamelCase_:							List[str]		=  len(snake_case_ ) if n_tasks is None else n_tasks
							UpperCamelCase_:							str		=  n_copies




			def __iter__(    self	:    Tuple ):
							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_:							List[str]		=  self.tokenizer(snake_case_  ,		padding=snake_case_  ,		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	:    Any  ,		snake_case_	:    List[str]  ,		snake_case_	:    Dict  ,		snake_case_	:    int ):
							UpperCamelCase_:							Union[str, Any]		=  start_length
							UpperCamelCase_:							Dict		=  eof_strings
							UpperCamelCase_:							List[str]		=  tokenizer




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







def     A__   (  lowerCamelCase  )					->		Optional[int]:
				UpperCamelCase_:							str		=  re.split("""(%s)""" % """|""".join(lowerCamelCase  )       ,      lowerCamelCase  )
				# last string should be ""
				return "".join(string_list[:-2]  )







def     A__   (  lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase=20       ,      **lowerCamelCase  )					->		int:
				UpperCamelCase_:							str		=  defaultdict(lowerCamelCase  )  # dict of list of generated tokens
				for step, batch in tqdm(enumerate(lowerCamelCase  )  ):
								with torch.no_grad():
												UpperCamelCase_:							Optional[int]		=  batch["""ids"""].shape[-1]
												UpperCamelCase_:							Dict		=  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_:							Optional[int]		=  batch["""task_id"""].repeat(lowerCamelCase  )
												UpperCamelCase_:							int		=  accelerator.pad_across_processes(
												    lowerCamelCase       ,      dim=1       ,      pad_index=tokenizer.pad_token_id  )

												UpperCamelCase_,				UpperCamelCase_:							Tuple		=  accelerator.gather((generated_tokens, generated_tasks)  )
												UpperCamelCase_:							Tuple		=  generated_tokens.cpu().numpy()
												UpperCamelCase_:							Any		=  generated_tasks.cpu().numpy()

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

				UpperCamelCase_:							Dict		=  [[] for _ in range(lowerCamelCase  )]
				for task, generated_tokens in gen_token_dict.items():
								for s in generated_tokens:
												UpperCamelCase_:							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     A__   (  )					->		Union[str, Any]:
				# Setup configuration
				UpperCamelCase_:							Optional[Any]		=  HfArgumentParser(lowerCamelCase  )
				UpperCamelCase_:							str		=  parser.parse_args()

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

				if args.num_workers is None:
								UpperCamelCase_:							Dict		=  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_:							Any		=  AutoTokenizer.from_pretrained(args.model_ckpt  )
				UpperCamelCase_:							Union[str, Any]		=  tokenizer.eos_token
				UpperCamelCase_:							Optional[Any]		=  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       ,      lowerCamelCase       ,      lowerCamelCase  )]  ),
				}

				# Load evaluation dataset and metric
				UpperCamelCase_:							Any		=  load_dataset("""openai_humaneval"""  )
				UpperCamelCase_:							Union[str, Any]		=  load_metric("""code_eval"""  )

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

				UpperCamelCase_:							Any		=  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_:							Optional[int]		=  DataLoader(lowerCamelCase       ,      batch_size=1  )

				# Run a quick test to see if code evaluation is enabled
				try:
								UpperCamelCase_:							List[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_:							Dict		=  accelerator.prepare(lowerCamelCase       ,      lowerCamelCase  )

				UpperCamelCase_:							Union[str, Any]		=  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_:							Optional[Any]		=  human_eval["""test"""][task]["""test"""]
												UpperCamelCase_:							Optional[int]		=  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()






