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

class    a_   (  __UpperCamelCase      ):




					def __init__(						self    :  List[Any]	,    snake_case__    :  str	,    snake_case__    :  List[Any]	,    snake_case__    :  List[Any]=None	,    snake_case__    :  str=1  ):
							lowerCAmelCase__   = tokenizer
							lowerCAmelCase__   = dataset
							lowerCAmelCase__   = len(snake_case__  ) if n_tasks is None else n_tasks
							lowerCAmelCase__   = n_copies




					def __iter__(						self    :  Any  ):
							lowerCAmelCase__   = []
							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()  )
							lowerCAmelCase__   = 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    a_   (  __UpperCamelCase      ):




					def __init__(						self    :  List[str]	,    snake_case__    :  List[str]	,    snake_case__    :  Any	,    snake_case__    :  str  ):
							lowerCAmelCase__   = start_length
							lowerCAmelCase__   = eof_strings
							lowerCAmelCase__   = tokenizer




					def __call__(						self    :  List[Any]	,    snake_case__    :  Dict	,    snake_case__    :  List[Any]	,    **snake_case__    :  Tuple  ):
							lowerCAmelCase__   = self.tokenizer.batch_decode(input_ids[:, self.start_length :]  )
							lowerCAmelCase__   = []
							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 				_UpperCAmelCase		(							lowerCamelCase__  ):





		"""simple docstring"""




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







def 				_UpperCAmelCase		(							lowerCamelCase__      ,	lowerCamelCase__      ,	lowerCamelCase__      ,	lowerCamelCase__      ,	lowerCamelCase__      ,	lowerCamelCase__=20      ,	**lowerCamelCase__  ):





		"""simple docstring"""




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

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

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

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







def 				_UpperCAmelCase		(							):





		"""simple docstring"""




		lowerCAmelCase__   = HfArgumentParser(lowerCamelCase__  )
		lowerCAmelCase__   = parser.parse_args()

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

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

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

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

		# Generation settings
		lowerCAmelCase__   = {
		    """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
		lowerCAmelCase__   = load_dataset("""openai_humaneval"""  )
		lowerCAmelCase__   = load_metric("""code_eval"""  )

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

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

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

		lowerCAmelCase__ ,       lowerCAmelCase__   = accelerator.prepare(lowerCamelCase__      ,	lowerCamelCase__  )

		lowerCAmelCase__   = complete_code(
		    lowerCamelCase__      ,	lowerCamelCase__      ,	lowerCamelCase__      ,	lowerCamelCase__      ,	n_tasks=lowerCamelCase__      ,	batch_size=args.batch_size      ,	**lowerCamelCase__      ,	)

		if accelerator.is_main_process:
				lowerCAmelCase__   = []

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

				# Evaluate completions with "code_eval" metric
				lowerCAmelCase__ ,       lowerCAmelCase__   = 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()





"""simple docstring"""




from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
			import numpy
			import tensorflow as tf

			from transformers import (
			    TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    BertConfig,
			    DPRConfig,
			    TFDPRContextEncoder,
			    TFDPRQuestionEncoder,
			    TFDPRReader,
			)

class    a_   :




					def __init__(						self    :  Union[str, Any]	,    snake_case__    :  Dict	,    snake_case__    :  int=13	,    snake_case__    :  List[str]=7	,    snake_case__    :  Any=True	,    snake_case__    :  Any=True	,    snake_case__    :  Dict=True	,    snake_case__    :  List[Any]=True	,    snake_case__    :  List[str]=99	,    snake_case__    :  Any=32	,    snake_case__    :  List[str]=2	,    snake_case__    :  Any=4	,    snake_case__    :  Dict=37	,    snake_case__    :  Optional[int]="gelu"	,    snake_case__    :  Optional[int]=0.1	,    snake_case__    :  Dict=0.1	,    snake_case__    :  Optional[int]=512	,    snake_case__    :  Union[str, Any]=16	,    snake_case__    :  str=2	,    snake_case__    :  Dict=0.02	,    snake_case__    :  Tuple=3	,    snake_case__    :  List[Any]=4	,    snake_case__    :  List[Any]=None	,    snake_case__    :  str=0	,    ):
							lowerCAmelCase__   = parent
							lowerCAmelCase__   = batch_size
							lowerCAmelCase__   = seq_length
							lowerCAmelCase__   = is_training
							lowerCAmelCase__   = use_input_mask
							lowerCAmelCase__   = use_token_type_ids
							lowerCAmelCase__   = use_labels
							lowerCAmelCase__   = vocab_size
							lowerCAmelCase__   = hidden_size
							lowerCAmelCase__   = num_hidden_layers
							lowerCAmelCase__   = num_attention_heads
							lowerCAmelCase__   = intermediate_size
							lowerCAmelCase__   = hidden_act
							lowerCAmelCase__   = hidden_dropout_prob
							lowerCAmelCase__   = attention_probs_dropout_prob
							lowerCAmelCase__   = max_position_embeddings
							lowerCAmelCase__   = type_vocab_size
							lowerCAmelCase__   = type_sequence_label_size
							lowerCAmelCase__   = initializer_range
							lowerCAmelCase__   = num_labels
							lowerCAmelCase__   = num_choices
							lowerCAmelCase__   = scope
							lowerCAmelCase__   = projection_dim




					def     _SCREAMING_SNAKE_CASE    (						self    :  Optional[int]  ):
							lowerCAmelCase__   = ids_tensor([self.batch_size, self.seq_length]	,    self.vocab_size  )

							lowerCAmelCase__   = None
							if self.use_input_mask:
									# follow test_modeling_tf_ctrl.py
									lowerCAmelCase__   = random_attention_mask([self.batch_size, self.seq_length]  )

							lowerCAmelCase__   = None
							if self.use_token_type_ids:
									lowerCAmelCase__   = ids_tensor([self.batch_size, self.seq_length]	,    self.type_vocab_size  )

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

							lowerCAmelCase__   = BertConfig(
							    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=snake_case__	,    initializer_range=self.initializer_range	,    )
							lowerCAmelCase__   = DPRConfig(projection_dim=self.projection_dim	,    **config.to_dict()  )

							return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




					def     _SCREAMING_SNAKE_CASE    (						self    :  Optional[Any]	,    snake_case__    :  int	,    snake_case__    :  Optional[int]	,    snake_case__    :  Optional[Any]	,    snake_case__    :  str	,    snake_case__    :  Tuple	,    snake_case__    :  Dict	,    snake_case__    :  List[Any]  ):
							lowerCAmelCase__   = TFDPRContextEncoder(config=snake_case__  )
							lowerCAmelCase__   = model(snake_case__	,    attention_mask=snake_case__	,    token_type_ids=snake_case__  )
							lowerCAmelCase__   = model(snake_case__	,    token_type_ids=snake_case__  )
							lowerCAmelCase__   = model(snake_case__  )
							self.parent.assertEqual(result.pooler_output.shape	,    (self.batch_size, self.projection_dim or self.hidden_size)  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  Any	,    snake_case__    :  str	,    snake_case__    :  Dict	,    snake_case__    :  str	,    snake_case__    :  Any	,    snake_case__    :  Union[str, Any]	,    snake_case__    :  List[str]	,    snake_case__    :  Dict  ):
							lowerCAmelCase__   = TFDPRQuestionEncoder(config=snake_case__  )
							lowerCAmelCase__   = model(snake_case__	,    attention_mask=snake_case__	,    token_type_ids=snake_case__  )
							lowerCAmelCase__   = model(snake_case__	,    token_type_ids=snake_case__  )
							lowerCAmelCase__   = model(snake_case__  )
							self.parent.assertEqual(result.pooler_output.shape	,    (self.batch_size, self.projection_dim or self.hidden_size)  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  Tuple	,    snake_case__    :  Tuple	,    snake_case__    :  Any	,    snake_case__    :  str	,    snake_case__    :  Tuple	,    snake_case__    :  List[str]	,    snake_case__    :  List[str]	,    snake_case__    :  Tuple  ):
							lowerCAmelCase__   = TFDPRReader(config=snake_case__  )
							lowerCAmelCase__   = model(snake_case__	,    attention_mask=snake_case__  )

							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)  )
							self.parent.assertEqual(result.relevance_logits.shape	,    (self.batch_size,)  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  Any  ):
							lowerCAmelCase__   = self.prepare_config_and_inputs()
							(
							    (
							    lowerCAmelCase__
							) ,       (
							    lowerCAmelCase__
							) ,       (
							    lowerCAmelCase__
							) ,       (
							    lowerCAmelCase__
							) ,       (
							    lowerCAmelCase__
							) ,       (
							    lowerCAmelCase__
							) ,       (
							    lowerCAmelCase__
							) ,       
							)   = config_and_inputs
							lowerCAmelCase__   = {"""input_ids""": input_ids}
							return config, inputs_dict

@require_tf
class    a_   (  __UpperCamelCase			,		__UpperCamelCase			,		unittest.TestCase      ):
					UpperCamelCase_			:						Optional[Any]      =     (
					    (
					        TFDPRContextEncoder,
					        TFDPRQuestionEncoder,
					        TFDPRReader,
					    )
					    if is_tf_available()
					    else ()
					)
					UpperCamelCase_			:						Any      =     {"feature-extraction": TFDPRQuestionEncoder} if is_tf_available() else {}

					UpperCamelCase_			:						Optional[Any]      =     False
					UpperCamelCase_			:						Tuple      =     False
					UpperCamelCase_			:						Tuple      =     False
					UpperCamelCase_			:						Optional[int]      =     False
					UpperCamelCase_			:						Optional[Any]      =     False




					def     _SCREAMING_SNAKE_CASE    (						self    :  Any  ):
							lowerCAmelCase__   = TFDPRModelTester(self  )
							lowerCAmelCase__   = ConfigTester(self	,    config_class=snake_case__	,    hidden_size=37  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  Optional[int]  ):
							self.config_tester.run_common_tests()




					def     _SCREAMING_SNAKE_CASE    (						self    :  Dict  ):
							lowerCAmelCase__   = self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_dpr_context_encoder(*snake_case__  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  Tuple  ):
							lowerCAmelCase__   = self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_dpr_question_encoder(*snake_case__  )




					def     _SCREAMING_SNAKE_CASE    (						self    :  int  ):
							lowerCAmelCase__   = self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_dpr_reader(*snake_case__  )




					@slow
					def     _SCREAMING_SNAKE_CASE    (						self    :  List[Any]  ):
							for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowerCAmelCase__   = TFDPRContextEncoder.from_pretrained(snake_case__  )
									self.assertIsNotNone(snake_case__  )

							for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowerCAmelCase__   = TFDPRContextEncoder.from_pretrained(snake_case__  )
									self.assertIsNotNone(snake_case__  )

							for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowerCAmelCase__   = TFDPRQuestionEncoder.from_pretrained(snake_case__  )
									self.assertIsNotNone(snake_case__  )

							for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowerCAmelCase__   = TFDPRReader.from_pretrained(snake_case__  )
									self.assertIsNotNone(snake_case__  )

@require_tf
class    a_   (  unittest.TestCase      ):




					@slow
					def     _SCREAMING_SNAKE_CASE    (						self    :  List[str]  ):
							lowerCAmelCase__   = TFDPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base"""  )

							lowerCAmelCase__   = tf.constant(
							    [[101, 7592, 1010, 2003, 2026, 3899, 10140, 1029, 102]]  )  # [CLS] hello, is my dog cute? [SEP]
							lowerCAmelCase__   = model(snake_case__  )[0]  # embedding shape = (1, 768)
							# compare the actual values for a slice.
							lowerCAmelCase__   = tf.constant(
							    [
							        [
							            0.0323_6253,
							            0.1275_3335,
							            0.1681_8509,
							            0.0027_9786,
							            0.389_6933,
							            0.2426_4945,
							            0.217_8971,
							            -0.0233_5227,
							            -0.0848_1959,
							            -0.1432_4117,
							        ]
							    ]  )
							self.assertTrue(numpy.allclose(output[:, :10].numpy()	,    expected_slice.numpy()	,    atol=1E-4  )  )




