



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


a_       =       ['''\nclass''', '''\ndef''', '''\n#''', '''\n@''', '''\nprint''', '''\nif''']







class  __lowercase  (       _UpperCAmelCase):







						"""simple docstring"""







						def __init__(self ,		lowercase__ ,		lowercase__ ,		lowercase__=None ,		lowercase__=1							):
											snake_case_       :		str			       =       tokenizer
											snake_case_       :		Tuple			       =       dataset
											snake_case_       :		List[Any]			       =       len(SCREAMING_SNAKE_CASE_							) if n_tasks is None else n_tasks
											snake_case_       :		List[str]			       =       n_copies







						def __iter__(self							):
											snake_case_       :		Tuple			       =       []
											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_       :		Dict			       =       self.tokenizer(SCREAMING_SNAKE_CASE_ ,		padding=SCREAMING_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  __lowercase  (       _UpperCAmelCase):







						"""simple docstring"""







						def __init__(self ,		lowercase__ ,		lowercase__ ,		lowercase__							):
											snake_case_       :		Any			       =       start_length
											snake_case_       :		int			       =       eof_strings
											snake_case_       :		int			       =       tokenizer







						def __call__(self ,		lowercase__ ,		lowercase__ ,		**lowercase__							):
											snake_case_       :		List[Any]			       =       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(SCREAMING_SNAKE_CASE_							)







def 				SCREAMING_SNAKE_CASE__   (		SCREAMING_SNAKE_CASE__			:	Dict      ):



					"""simple docstring"""





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







def 				SCREAMING_SNAKE_CASE__   (		SCREAMING_SNAKE_CASE__			:	Union[str, Any]					,     SCREAMING_SNAKE_CASE__			:	List[str]					,     SCREAMING_SNAKE_CASE__			:	Optional[int]					,     SCREAMING_SNAKE_CASE__			:	Optional[Any]					,     SCREAMING_SNAKE_CASE__			:	Optional[Any]					,     SCREAMING_SNAKE_CASE__			:	List[Any]=2_0					,     **SCREAMING_SNAKE_CASE__			:	str      ):



					"""simple docstring"""





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

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

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

					snake_case_       :		Tuple			       =       [[] for _ in range(snake_case__      )]
					for task, generated_tokens in gen_token_dict.items():
										for s in generated_tokens:
															snake_case_       :		str			       =       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 				SCREAMING_SNAKE_CASE__   (		):



					"""simple docstring"""





					snake_case_       :		Optional[Any]			       =       HfArgumentParser(snake_case__      )
					snake_case_       :		Optional[Any]			       =       parser.parse_args()

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

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

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

					# Load model and tokenizer
					snake_case_       :		Tuple			       =       AutoTokenizer.from_pretrained(args.model_ckpt      )
					snake_case_       :		Tuple			       =       tokenizer.eos_token
					snake_case_       :		str			       =       AutoModelForCausalLM.from_pretrained(args.model_ckpt      )

					# Generation settings
					snake_case_       :		Optional[int]			       =       {
					    """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
					snake_case_       :		int			       =       load_dataset("""openai_humaneval"""      )
					snake_case_       :		List[Any]			       =       load_metric("""code_eval"""      )

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

					snake_case_       :		Tuple			       =       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
					snake_case_       :		Union[str, Any]			       =       DataLoader(snake_case__					,     batch_size=1      )

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

					snake_case_       :		Union[str, Any]			       =       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:
										snake_case_       :		Optional[Any]			       =       []

										for task in tqdm(range(snake_case__      )      ):
															snake_case_       :		str			       =       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_       :		Any			       =       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()






import argparse
import logging
import os
import re

import tensorflow as tf

from transformers import (
    AutoConfig,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
    PushToHubCallback,
    TFAutoModelForMaskedLM,
    create_optimizer,
)


__a         =   logging.getLogger(__name__)

__a         =   tf.data.AUTOTUNE




def 							a		(				):


				'''simple docstring'''
				lowercase_	    =   argparse.ArgumentParser(description='''Train a masked language model on TPU.'''   )
				parser.add_argument(
				    '''--pretrained_model_config'''		,  type=snake_case__		,  default='''roberta-base'''		,  help='''The model config to use. Note that we don\'t copy the model\'s weights, only the config!'''		,  )
				parser.add_argument(
				    '''--tokenizer'''		,  type=snake_case__		,  default='''unigram-tokenizer-wikitext'''		,  help='''The name of the tokenizer to load. We use the pretrained tokenizer to initialize the model\'s vocab size.'''		,  )

				parser.add_argument(
				    '''--per_replica_batch_size'''		,  type=snake_case__		,  default=8		,  help='''Batch size per TPU core.'''		,  )

				parser.add_argument(
				    '''--no_tpu'''		,  action='''store_true'''		,  help='''If set, run on CPU and don\'t try to initialize a TPU. Useful for debugging on non-TPU instances.'''		,  )

				parser.add_argument(
				    '''--tpu_name'''		,  type=snake_case__		,  help='''Name of TPU resource to initialize. Should be blank on Colab, and \'local\' on TPU VMs.'''		,  default='''local'''		,  )

				parser.add_argument(
				    '''--tpu_zone'''		,  type=snake_case__		,  help='''Google cloud zone that TPU resource is located in. Only used for non-Colab TPU nodes.'''		,  )

				parser.add_argument(
				    '''--gcp_project'''		,  type=snake_case__		,  help='''Google cloud project name. Only used for non-Colab TPU nodes.'''   )

				parser.add_argument(
				    '''--bfloat16'''		,  action='''store_true'''		,  help='''Use mixed-precision bfloat16 for training. This is the recommended lower-precision format for TPU.'''		,  )

				parser.add_argument(
				    '''--train_dataset'''		,  type=snake_case__		,  help='''Path to training dataset to load. If the path begins with `gs://`'''
				    ''' then the dataset will be loaded from a Google Cloud Storage bucket.'''		,  )

				parser.add_argument(
				    '''--shuffle_buffer_size'''		,  type=snake_case__		,  default=2**18		,  help='''Size of the shuffle buffer (in samples)'''		,  )

				parser.add_argument(
				    '''--eval_dataset'''		,  type=snake_case__		,  help='''Path to evaluation dataset to load. If the path begins with `gs://`'''
				    ''' then the dataset will be loaded from a Google Cloud Storage bucket.'''		,  )

				parser.add_argument(
				    '''--num_epochs'''		,  type=snake_case__		,  default=1		,  help='''Number of epochs to train for.'''		,  )

				parser.add_argument(
				    '''--learning_rate'''		,  type=snake_case__		,  default=1e-4		,  help='''Learning rate to use for training.'''		,  )

				parser.add_argument(
				    '''--weight_decay_rate'''		,  type=snake_case__		,  default=1e-3		,  help='''Weight decay rate to use for training.'''		,  )

				parser.add_argument(
				    '''--max_length'''		,  type=snake_case__		,  default=512		,  help='''Maximum length of tokenized sequences. Should match the setting used in prepare_tfrecord_shards.py'''		,  )

				parser.add_argument(
				    '''--mlm_probability'''		,  type=snake_case__		,  default=0.1_5		,  help='''Fraction of tokens to mask during training.'''		,  )

				parser.add_argument('''--output_dir'''		,  type=snake_case__		,  required=snake_case__		,  help='''Path to save model checkpoints to.'''   )
				parser.add_argument('''--hub_model_id'''		,  type=snake_case__		,  help='''Model ID to upload to on the Hugging Face Hub.'''   )

				lowercase_	    =   parser.parse_args()
				return args




def 							a		(				snake_case__:   int   ):


				'''simple docstring'''
				try:
								if args.tpu_name:
												lowercase_	    =   tf.distribute.cluster_resolver.TPUClusterResolver(
												    args.tpu_name		,  zone=args.tpu_zone		,  project=args.gcp_project   )
								else:
												lowercase_	    =   tf.distribute.cluster_resolver.TPUClusterResolver()
				except ValueError:
								raise RuntimeError(
								    '''Couldn\'t connect to TPU! Most likely you need to specify --tpu_name, --tpu_zone, or '''
								    '''--gcp_project. When running on a TPU VM, use --tpu_name local.'''   )

				tf.config.experimental_connect_to_cluster(snake_case__   )
				tf.tpu.experimental.initialize_tpu_system(snake_case__   )

				return tpu




def 							a		(				snake_case__:   Optional[int]   ):


				'''simple docstring'''
				lowercase_	    =   0
				for file in file_list:
								lowercase_	    =   file.split('''/'''   )[-1]
								lowercase_	    =   re.search(r'''-\d+-(\d+)\.tfrecord'''		,  snake_case__   ).group(1   )
								lowercase_	    =   int(snake_case__   )
								num_samples += sample_count

				return num_samples




def 							a		(				snake_case__:   Dict		,  snake_case__:   Dict		,  snake_case__:   Union[str, Any]		,  snake_case__:   Optional[int]		,  snake_case__:   Optional[int]		,  snake_case__:   Dict=None   ):


				'''simple docstring'''
				lowercase_	    =   count_samples(snake_case__   )
				lowercase_	    =   tf.data.Dataset.from_tensor_slices(snake_case__   )
				if shuffle:
								lowercase_	    =   dataset.shuffle(len(snake_case__   )   )
				lowercase_	    =   tf.data.TFRecordDataset(snake_case__		,  num_parallel_reads=snake_case__   )
				# TF can't infer the total sample count because it doesn't read all the records yet, so we assert it here
				lowercase_	    =   dataset.apply(tf.data.experimental.assert_cardinality(snake_case__   )   )
				lowercase_	    =   dataset.map(snake_case__		,  num_parallel_calls=snake_case__   )
				if shuffle:
								assert shuffle_buffer_size is not None
								lowercase_	    =   dataset.shuffle(args.shuffle_buffer_size   )
				lowercase_	    =   dataset.batch(snake_case__		,  drop_remainder=snake_case__   )
				lowercase_	    =   dataset.map(snake_case__		,  num_parallel_calls=snake_case__   )
				lowercase_	    =   dataset.prefetch(snake_case__   )
				return dataset




def 							a		(				snake_case__:   Tuple   ):


				'''simple docstring'''
				if not args.no_tpu:
								lowercase_	    =   initialize_tpu(snake_case__   )
								lowercase_	    =   tf.distribute.TPUStrategy(snake_case__   )
				else:
								lowercase_	    =   tf.distribute.OneDeviceStrategy(device='''/gpu:0'''   )

				if args.bfloataa:
								tf.keras.mixed_precision.set_global_policy('''mixed_bfloat16'''   )

				lowercase_	    =   AutoTokenizer.from_pretrained(args.tokenizer   )
				lowercase_	    =   AutoConfig.from_pretrained(args.pretrained_model_config   )
				lowercase_	    =   tokenizer.vocab_size

				lowercase_	    =   tf.io.gfile.glob(os.path.join(args.train_dataset		,  '''*.tfrecord'''   )   )
				if not training_records:
								raise ValueError(F'''No .tfrecord files found in {args.train_dataset}.'''   )
				lowercase_	    =   tf.io.gfile.glob(os.path.join(args.eval_dataset		,  '''*.tfrecord'''   )   )
				if not eval_records:
								raise ValueError(F'''No .tfrecord files found in {args.eval_dataset}.'''   )

				lowercase_	    =   count_samples(snake_case__   )

				lowercase_	    =   num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
				lowercase_	    =   steps_per_epoch * args.num_epochs

				with strategy.scope():
								lowercase_	    =   TFAutoModelForMaskedLM.from_config(snake_case__   )
								model(model.dummy_inputs   )  # Pass some dummy inputs through the model to ensure all the weights are built
								lowercase_    ,      lowercase_	    =   create_optimizer(
								    num_train_steps=snake_case__		,  num_warmup_steps=total_train_steps // 20		,  init_lr=args.learning_rate		,  weight_decay_rate=args.weight_decay_rate		,  )

								# Transformers models compute the right loss for their task by default when labels are passed, and will
								# use this for training unless you specify your own loss function in compile().
								model.compile(optimizer=snake_case__		,  metrics=['''accuracy''']   )

				def decode_fn(snake_case__:   Dict   ):
								lowercase_	    =   {
								    '''input_ids''': tf.io.FixedLenFeature(dtype=tf.intaa		,  shape=(args.max_length,)   ),
								    '''attention_mask''': tf.io.FixedLenFeature(dtype=tf.intaa		,  shape=(args.max_length,)   ),
								}
								return tf.io.parse_single_example(snake_case__		,  snake_case__   )

				# Many of the data collators in Transformers are TF-compilable when return_tensors == "tf", so we can
				# use their methods in our data pipeline.
				lowercase_	    =   DataCollatorForLanguageModeling(
				    tokenizer=snake_case__		,  mlm_probability=args.mlm_probability		,  mlm=snake_case__		,  return_tensors='''tf'''   )

				def mask_with_collator(snake_case__:   str   ):
								# TF really needs an isin() function
								lowercase_	    =   (
								    ~tf.cast(batch['''attention_mask''']		,  tf.bool   )
								    | (batch['''input_ids'''] == tokenizer.cls_token_id)
								    | (batch['''input_ids'''] == tokenizer.sep_token_id)
								)
								lowercase_    ,      lowercase_	    =   data_collator.tf_mask_tokens(
								    batch['''input_ids''']		,  vocab_size=len(snake_case__   )		,  mask_token_id=tokenizer.mask_token_id		,  special_tokens_mask=snake_case__		,  )
								return batch

				lowercase_	    =   args.per_replica_batch_size * strategy.num_replicas_in_sync

				lowercase_	    =   prepare_dataset(
				    snake_case__		,  decode_fn=snake_case__		,  mask_fn=snake_case__		,  batch_size=snake_case__		,  shuffle=snake_case__		,  shuffle_buffer_size=args.shuffle_buffer_size		,  )

				lowercase_	    =   prepare_dataset(
				    snake_case__		,  decode_fn=snake_case__		,  mask_fn=snake_case__		,  batch_size=snake_case__		,  shuffle=snake_case__		,  )

				lowercase_	    =   []
				if args.hub_model_id:
								callbacks.append(
								    PushToHubCallback(output_dir=args.output_dir		,  hub_model_id=args.hub_model_id		,  tokenizer=snake_case__   )   )

				model.fit(
				    snake_case__		,  validation_data=snake_case__		,  epochs=args.num_epochs		,  callbacks=snake_case__		,  )

				model.save_pretrained(args.output_dir   )


if __name__ == "__main__":
						__a         =   parse_args()
						main(args)




