





import json
import os
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import MgpstrTokenizer
from transformers.models.mgp_str.tokenization_mgp_str import VOCAB_FILES_NAMES
from transformers.testing_utils import require_torch, require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_torch_available, is_vision_available


if is_torch_available():
			import torch


if is_vision_available():
			from PIL import Image

			from transformers import MgpstrProcessor, ViTImageProcessor
@require_torch
@require_vision
class SCREAMING_SNAKE_CASE_  (   unittest.TestCase   ):
						__magic_name__:     Optional[Any]	 =    ViTImageProcessor if is_vision_available() else None






						@property
						def UpperCAmelCase_			( self : Tuple		)							->      Optional[Any]:







							"""simple docstring"""


							return self.image_processor_tester.prepare_image_processor_dict()






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







							"""simple docstring"""


							snake_case_						:     Any     =		(3, 32, 128)
							snake_case_						:     Tuple     =		tempfile.mkdtemp()

							# fmt: off
							snake_case_						:     Optional[int]     =		['[GO]', '[s]', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
							# fmt: on
							snake_case_						:     Optional[int]     =		dict(zip(__a						,				range(len(__a		)		)		)		)

							snake_case_						:     int     =		os.path.join(self.tmpdirname						,				VOCAB_FILES_NAMES['vocab_file']		)
							with open(self.vocab_file						,				'w'						,				encoding='utf-8'		) as fp:
								fp.write(json.dumps(__a		) + '\n'		)

							snake_case_						:     Union[str, Any]     =		{
							    'do_normalize': False,
							    'do_resize': True,
							    'image_processor_type': 'ViTImageProcessor',
							    'resample': 3,
							    'size': {'height': 32, 'width': 128},
							}
							snake_case_						:     Dict     =		os.path.join(self.tmpdirname						,				__a		)
							with open(self.image_processor_file						,				'w'						,				encoding='utf-8'		) as fp:
								json.dump(__a						,				__a		)






						def UpperCAmelCase_			( self : str						,				**_A : Any		)							->      str:







							"""simple docstring"""


							return MgpstrTokenizer.from_pretrained(self.tmpdirname						,				**__a		)






						def UpperCAmelCase_			( self : Optional[int]						,				**_A : Any		)							->      Union[str, Any]:







							"""simple docstring"""


							return ViTImageProcessor.from_pretrained(self.tmpdirname						,				**__a		)






						def UpperCAmelCase_			( self : str		)							->      Dict:







							"""simple docstring"""


							shutil.rmtree(self.tmpdirname		)






						def UpperCAmelCase_			( self : int		)							->      Dict:







							"""simple docstring"""


							snake_case_						:     List[Any]     =		np.random.randint(255						,				size=(3, 30, 400)						,				dtype=np.uinta		)

							snake_case_						:     Dict     =		Image.fromarray(np.moveaxis(__a						,				0						,				-1		)		)

							return image_input






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







							"""simple docstring"""


							snake_case_						:     Dict     =		self.get_tokenizer()
							snake_case_						:     List[str]     =		self.get_image_processor()

							snake_case_						:     Optional[int]     =		MgpstrProcessor(tokenizer=__a						,				image_processor=__a		)
							processor.save_pretrained(self.tmpdirname		)
							snake_case_						:     Union[str, Any]     =		MgpstrProcessor.from_pretrained(self.tmpdirname						,				use_fast=__a		)

							self.assertEqual(processor.char_tokenizer.get_vocab()						,				tokenizer.get_vocab()		)
							self.assertIsInstance(processor.char_tokenizer						,				__a		)

							self.assertEqual(processor.image_processor.to_json_string()						,				image_processor.to_json_string()		)
							self.assertIsInstance(processor.image_processor						,				__a		)






						def UpperCAmelCase_			( self : str		)							->      List[str]:







							"""simple docstring"""


							snake_case_						:     str     =		self.get_tokenizer()
							snake_case_						:     Optional[Any]     =		self.get_image_processor()

							snake_case_						:     Optional[int]     =		MgpstrProcessor(tokenizer=__a						,				image_processor=__a		)
							processor.save_pretrained(self.tmpdirname		)

							snake_case_						:     str     =		self.get_tokenizer(bos_token='(BOS)'						,				eos_token='(EOS)'		)
							snake_case_						:     str     =		self.get_image_processor(do_normalize=__a						,				padding_value=1.0		)

							snake_case_						:     Any     =		MgpstrProcessor.from_pretrained(
							    self.tmpdirname						,				bos_token='(BOS)'						,				eos_token='(EOS)'						,				do_normalize=__a						,				padding_value=1.0		)

							self.assertEqual(processor.char_tokenizer.get_vocab()						,				tokenizer_add_kwargs.get_vocab()		)
							self.assertIsInstance(processor.char_tokenizer						,				__a		)

							self.assertEqual(processor.image_processor.to_json_string()						,				image_processor_add_kwargs.to_json_string()		)
							self.assertIsInstance(processor.image_processor						,				__a		)






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







							"""simple docstring"""


							snake_case_						:     Optional[Any]     =		self.get_image_processor()
							snake_case_						:     Optional[Any]     =		self.get_tokenizer()

							snake_case_						:     Optional[Any]     =		MgpstrProcessor(tokenizer=__a						,				image_processor=__a		)

							snake_case_						:     Optional[Any]     =		self.prepare_image_inputs()

							snake_case_						:     Any     =		image_processor(__a						,				return_tensors='np'		)
							snake_case_						:     Optional[int]     =		processor(images=__a						,				return_tensors='np'		)

							for key in input_image_proc.keys():
								self.assertAlmostEqual(input_image_proc[key].sum()						,				input_processor[key].sum()						,				delta=1E-2		)






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







							"""simple docstring"""


							snake_case_						:     Dict     =		self.get_image_processor()
							snake_case_						:     List[Any]     =		self.get_tokenizer()

							snake_case_						:     int     =		MgpstrProcessor(tokenizer=__a						,				image_processor=__a		)

							snake_case_						:     Tuple     =		'test'

							snake_case_						:     Tuple     =		processor(text=__a		)

							snake_case_						:     Tuple     =		tokenizer(__a		)
							for key in encoded_tok.keys():
								self.assertListEqual(encoded_tok[key]						,				encoded_processor[key]		)






						def UpperCAmelCase_			( self : int		)							->      Dict:







							"""simple docstring"""


							snake_case_						:     Dict     =		self.get_image_processor()
							snake_case_						:     Optional[int]     =		self.get_tokenizer()

							snake_case_						:     Optional[Any]     =		MgpstrProcessor(tokenizer=__a						,				image_processor=__a		)

							snake_case_						:     Dict     =		'test'
							snake_case_						:     Tuple     =		self.prepare_image_inputs()

							snake_case_						:     Optional[Any]     =		processor(text=__a						,				images=__a		)

							self.assertListEqual(list(inputs.keys()		)						,				['pixel_values', 'labels']		)

							# test if it raises when no input is passed
							with pytest.raises(__a		):
								processor()






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







							"""simple docstring"""


							snake_case_						:     str     =		self.get_image_processor()
							snake_case_						:     Optional[Any]     =		self.get_tokenizer()

							snake_case_						:     Tuple     =		MgpstrProcessor(tokenizer=__a						,				image_processor=__a		)

							snake_case_						:     Any     =		[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9], [3, 4, 3, 1, 1, 8, 9]]

							snake_case_						:     List[Any]     =		processor.char_decode(__a		)
							snake_case_						:     Any     =		tokenizer.batch_decode(__a		)
							snake_case_						:     List[Any]     =		[seq.replace(' '						,				''		) for seq in decoded_tok]

							self.assertListEqual(__a						,				__a		)






						def UpperCAmelCase_			( self : Any		)							->      List[Any]:







							"""simple docstring"""


							snake_case_						:     List[str]     =		self.get_image_processor()
							snake_case_						:     Optional[Any]     =		self.get_tokenizer()

							snake_case_						:     Union[str, Any]     =		MgpstrProcessor(tokenizer=__a						,				image_processor=__a		)

							snake_case_						:     Optional[int]     =		None
							snake_case_						:     List[Any]     =		self.prepare_image_inputs()

							snake_case_						:     Optional[Any]     =		processor(text=__a						,				images=__a		)

							self.assertListEqual(list(inputs.keys()		)						,				processor.model_input_names		)






						def UpperCAmelCase_			( self : Any		)							->      Tuple:







							"""simple docstring"""


							snake_case_						:     Optional[Any]     =		self.get_image_processor()
							snake_case_						:     int     =		self.get_tokenizer()

							snake_case_						:     Optional[int]     =		MgpstrProcessor(tokenizer=__a						,				image_processor=__a		)

							snake_case_						:     str     =		torch.randn(1						,				27						,				38		)
							snake_case_						:     int     =		torch.randn(1						,				27						,				50257		)
							snake_case_						:     Optional[int]     =		torch.randn(1						,				27						,				30522		)

							snake_case_						:     Tuple     =		processor.batch_decode([char_input, bpe_input, wp_input]		)

							self.assertListEqual(list(results.keys()		)						,				['generated_text', 'scores', 'char_preds', 'bpe_preds', 'wp_preds']		)








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 		=						["""\nclass""", """\ndef""", """\n#""", """\n@""", """\nprint""", """\nif"""]
class SCREAMING_SNAKE_CASE_  (   snake_case_   ):






						def __init__( self : Tuple						,				_A : Any						,				_A : str						,				_A : int=None						,				_A : str=1		)							->      List[str]:







							"""simple docstring"""


							snake_case_						:     Union[str, Any]     =		tokenizer
							snake_case_						:     Optional[int]     =		dataset
							snake_case_						:     List[str]     =		len(_A		) if n_tasks is None else n_tasks
							snake_case_						:     List[Any]     =		n_copies






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







							"""simple docstring"""


							snake_case_						:     List[str]     =		[]
							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()		)
							snake_case_						:     Optional[int]     =		self.tokenizer(_A						,				padding=_A						,				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 SCREAMING_SNAKE_CASE_  (   snake_case_   ):






						def __init__( self : List[Any]						,				_A : Optional[int]						,				_A : str						,				_A : Dict		)							->      Any:







							"""simple docstring"""


							snake_case_						:     List[str]     =		start_length
							snake_case_						:     int     =		eof_strings
							snake_case_						:     Dict     =		tokenizer






						def __call__( self : Any						,				_A : Union[str, Any]						,				_A : Dict						,				**_A : List[Any]		)							->      List[str]:







							"""simple docstring"""


							snake_case_						:     Optional[int]     =		self.tokenizer.batch_decode(input_ids[:, self.start_length :]		)
							snake_case_						:     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(_A		)


def 							SCREAMING_SNAKE_CASE__  (      __a				):
	snake_case_						:     List[str]     =		re.split('(%s)' % '|'.join(__a				)				,	__a				)
	# last string should be ""
	return "".join(string_list[:-2]				)


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

			snake_case_  ,snake_case_						:     Union[str, Any]     =		accelerator.gather((generated_tokens, generated_tasks)				)
			snake_case_						:     Optional[Any]     =		generated_tokens.cpu().numpy()
			snake_case_						:     Dict     =		generated_tasks.cpu().numpy()

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

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


def 							SCREAMING_SNAKE_CASE__  (      ):
	# Setup configuration
	snake_case_						:     Optional[int]     =		HfArgumentParser(__a				)
	snake_case_						:     Tuple     =		parser.parse_args()

	transformers.logging.set_verbosity_error()
	# enables code execution in code_eval metric
	snake_case_						:     Optional[int]     =		args.HF_ALLOW_CODE_EVAL
	# make sure tokenizer plays nice with multiprocessing
	snake_case_						:     int     =		'false'

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

	# Use dataset load to feed to accelerate
	snake_case_						:     List[Any]     =		Accelerator()
	set_seed(args.seed				,	device_specific=__a				)

	# Load model and tokenizer
	snake_case_						:     Any     =		AutoTokenizer.from_pretrained(args.model_ckpt				)
	snake_case_						:     int     =		tokenizer.eos_token
	snake_case_						:     List[str]     =		AutoModelForCausalLM.from_pretrained(args.model_ckpt				)

	# Generation settings
	snake_case_						:     List[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				,	__a				,	__a				)]				),
	}

	# Load evaluation dataset and metric
	snake_case_						:     Dict     =		load_dataset('openai_humaneval'				)
	snake_case_						:     Optional[Any]     =		load_metric('code_eval'				)

	snake_case_						:     Union[str, Any]     =		args.num_tasks if args.num_tasks is not None else len(human_eval['test']				)
	snake_case_						:     Optional[int]     =		args.n_samples // args.batch_size

	snake_case_						:     Dict     =		TokenizedDataset(__a				,	human_eval['test']				,	n_copies=__a				,	n_tasks=__a				)
	# do not confuse args.batch_size, which is actually the num_return_sequences
	snake_case_						:     Optional[Any]     =		DataLoader(__a				,	batch_size=1				)

	# Run a quick test to see if code evaluation is enabled
	try:
		snake_case_						:     Union[str, Any]     =		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

	snake_case_  ,snake_case_						:     Union[str, Any]     =		accelerator.prepare(__a				,	__a				)

	snake_case_						:     str     =		complete_code(
	    __a				,	__a				,	__a				,	__a				,	n_tasks=__a				,	batch_size=args.batch_size				,	**__a				,	)

	if accelerator.is_main_process:
		snake_case_						:     Tuple     =		[]

		for task in tqdm(range(__a				)				):
			snake_case_						:     Union[str, Any]     =		human_eval['test'][task]['test']
			snake_case_						:     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
		snake_case_  ,snake_case_						:     int     =		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()

