

"""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


_lowercase   :	Dict												=						['\nclass', '\ndef', '\n#', '\n@', '\nprint', '\nif']
class      _UpperCAmelCase   (							_lowerCAmelCase       ):



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



	def __iter__(  self     :       int						):
		__UpperCAmelCase      		= []
		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      		= self.tokenizer(_lowercase			,    padding=_lowercase			,    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   (							_lowerCAmelCase       ):



	def __init__(  self     :       Union[str, Any]			,    _lowercase     :       Optional[int]			,    _lowercase     :       Optional[int]			,    _lowercase     :       int						):
		__UpperCAmelCase      		= start_length
		__UpperCAmelCase      		= eof_strings
		__UpperCAmelCase      		= tokenizer



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






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





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

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

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

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





def       lowercase__						(       ):
	# Setup configuration
	__UpperCAmelCase      		= HfArgumentParser(snake_case_    )
	__UpperCAmelCase      		= parser.parse_args()

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

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

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

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

	# Generation settings
	__UpperCAmelCase      		= {
	    '''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			,		snake_case_			,		snake_case_    )]    ),
	}

	# Load evaluation dataset and metric
	__UpperCAmelCase      		= load_dataset('''openai_humaneval'''    )
	__UpperCAmelCase      		= load_metric('''code_eval'''    )

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

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

	# Run a quick test to see if code evaluation is enabled
	try:
		__UpperCAmelCase      		= 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      		= accelerator.prepare(snake_case_			,		snake_case_    )

	__UpperCAmelCase      		= complete_code(
	    snake_case_			,		snake_case_			,		snake_case_			,		snake_case_			,		n_tasks=snake_case_			,		batch_size=args.batch_size			,		**snake_case_			,		)

	if accelerator.is_main_process:
		__UpperCAmelCase      		= []

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

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


# 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 typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


lowerCAmelCase_      :  Optional[Any]				     =  {
    'configuration_mobilevit': ['MOBILEVIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MobileViTConfig', 'MobileViTOnnxConfig'],
}

try:
 if not is_vision_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCAmelCase_      :  int				     =  ['MobileViTFeatureExtractor']
 lowerCAmelCase_      :  int				     =  ['MobileViTImageProcessor']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCAmelCase_      :  str				     =  [
     'MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST',
     'MobileViTForImageClassification',
     'MobileViTForSemanticSegmentation',
     'MobileViTModel',
     'MobileViTPreTrainedModel',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCAmelCase_      :  Optional[Any]				     =  [
     'TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST',
     'TFMobileViTForImageClassification',
     'TFMobileViTForSemanticSegmentation',
     'TFMobileViTModel',
     'TFMobileViTPreTrainedModel',
 ]

if TYPE_CHECKING:
 from .configuration_mobilevit import MOBILEVIT_PRETRAINED_CONFIG_ARCHIVE_MAP, MobileViTConfig, MobileViTOnnxConfig

 try:
  if not is_vision_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .feature_extraction_mobilevit import MobileViTFeatureExtractor
  from .image_processing_mobilevit import MobileViTImageProcessor

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_mobilevit import (
      MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
      MobileViTForImageClassification,
      MobileViTForSemanticSegmentation,
      MobileViTModel,
      MobileViTPreTrainedModel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_mobilevit import (
      TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFMobileViTForImageClassification,
      TFMobileViTForSemanticSegmentation,
      TFMobileViTModel,
      TFMobileViTPreTrainedModel,
  )


else:
 import sys

 lowerCAmelCase_      :  List[Any]				     =  _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
