
import argparse
import logging
import os
import re

import tensorflow as tf

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


lowerCamelCase__          =		logging.getLogger(__name__)

lowerCamelCase__          =		tf.data.AUTOTUNE
def  lowerCAmelCase__			(  )				->		str:
			lowerCAmelCase__      :      Tuple         =		argparse.ArgumentParser(description='Train a masked language model on TPU.'     )
			parser.add_argument(
			    '--pretrained_model_config'       ,							type=SCREAMING_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=SCREAMING_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=SCREAMING_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=SCREAMING_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=SCREAMING_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=SCREAMING_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=SCREAMING_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=SCREAMING_SNAKE_CASE_       ,							default=2**18       ,							help='Size of the shuffle buffer (in samples)'       ,							)

			parser.add_argument(
			    '--eval_dataset'       ,							type=SCREAMING_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=SCREAMING_SNAKE_CASE_       ,							default=1       ,							help='Number of epochs to train for.'       ,							)

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

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

			parser.add_argument(
			    '--max_length'       ,							type=SCREAMING_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=SCREAMING_SNAKE_CASE_       ,							default=0.15       ,							help='Fraction of tokens to mask during training.'       ,							)

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

			lowerCAmelCase__      :      Union[str, Any]         =		parser.parse_args()
			return args
def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_     )				->		Tuple:
			try:
						if args.tpu_name:
									lowerCAmelCase__      :      Optional[Any]         =		tf.distribute.cluster_resolver.TPUClusterResolver(
									    args.tpu_name       ,							zone=args.tpu_zone       ,							project=args.gcp_project     )
						else:
									lowerCAmelCase__      :      List[str]         =		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(SCREAMING_SNAKE_CASE_     )
			tf.tpu.experimental.initialize_tpu_system(SCREAMING_SNAKE_CASE_     )

			return tpu
def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_     )				->		Dict:
			lowerCAmelCase__      :      Any         =		0
			for file in file_list:
						lowerCAmelCase__      :      List[str]         =		file.split('/'     )[-1]
						lowerCAmelCase__      :      Optional[Any]         =		re.search(r'-\d+-(\d+)\.tfrecord'       ,							SCREAMING_SNAKE_CASE_     ).group(1     )
						lowerCAmelCase__      :      Optional[Any]         =		int(SCREAMING_SNAKE_CASE_     )
						num_samples += sample_count

			return num_samples
def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_=None     )				->		Any:
			lowerCAmelCase__      :      Any         =		count_samples(SCREAMING_SNAKE_CASE_     )
			lowerCAmelCase__      :      List[Any]         =		tf.data.Dataset.from_tensor_slices(SCREAMING_SNAKE_CASE_     )
			if shuffle:
						lowerCAmelCase__      :      Optional[Any]         =		dataset.shuffle(len(SCREAMING_SNAKE_CASE_     )     )
			lowerCAmelCase__      :      int         =		tf.data.TFRecordDataset(SCREAMING_SNAKE_CASE_       ,							num_parallel_reads=SCREAMING_SNAKE_CASE_     )
			# TF can't infer the total sample count because it doesn't read all the records yet, so we assert it here
			lowerCAmelCase__      :      List[Any]         =		dataset.apply(tf.data.experimental.assert_cardinality(SCREAMING_SNAKE_CASE_     )     )
			lowerCAmelCase__      :      Tuple         =		dataset.map(SCREAMING_SNAKE_CASE_       ,							num_parallel_calls=SCREAMING_SNAKE_CASE_     )
			if shuffle:
						assert shuffle_buffer_size is not None
						lowerCAmelCase__      :      Optional[Any]         =		dataset.shuffle(args.shuffle_buffer_size     )
			lowerCAmelCase__      :      Any         =		dataset.batch(SCREAMING_SNAKE_CASE_       ,							drop_remainder=SCREAMING_SNAKE_CASE_     )
			lowerCAmelCase__      :      Union[str, Any]         =		dataset.map(SCREAMING_SNAKE_CASE_       ,							num_parallel_calls=SCREAMING_SNAKE_CASE_     )
			lowerCAmelCase__      :      Optional[Any]         =		dataset.prefetch(SCREAMING_SNAKE_CASE_     )
			return dataset
def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_     )				->		Tuple:
			if not args.no_tpu:
						lowerCAmelCase__      :      Union[str, Any]         =		initialize_tpu(SCREAMING_SNAKE_CASE_     )
						lowerCAmelCase__      :      Optional[int]         =		tf.distribute.TPUStrategy(SCREAMING_SNAKE_CASE_     )
			else:
						lowerCAmelCase__      :      Union[str, Any]         =		tf.distribute.OneDeviceStrategy(device='/gpu:0'     )

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

			lowerCAmelCase__      :      Any         =		AutoTokenizer.from_pretrained(args.tokenizer     )
			lowerCAmelCase__      :      str         =		AutoConfig.from_pretrained(args.pretrained_model_config     )
			lowerCAmelCase__      :      str         =		tokenizer.vocab_size

			lowerCAmelCase__      :      Dict         =		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}.'''     )
			lowerCAmelCase__      :      List[Any]         =		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}.'''     )

			lowerCAmelCase__      :      Optional[Any]         =		count_samples(SCREAMING_SNAKE_CASE_     )

			lowerCAmelCase__      :      str         =		num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
			lowerCAmelCase__      :      List[Any]         =		steps_per_epoch * args.num_epochs

			with strategy.scope():
						lowerCAmelCase__      :      Dict         =		TFAutoModelForMaskedLM.from_config(SCREAMING_SNAKE_CASE_     )
						model(model.dummy_inputs     )  # Pass some dummy inputs through the model to ensure all the weights are built
						lowerCAmelCase__      ,				lowerCAmelCase__      :      int         =		create_optimizer(
						    num_train_steps=SCREAMING_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=SCREAMING_SNAKE_CASE_       ,							metrics=['accuracy']     )

			def decode_fn(SCREAMING_SNAKE_CASE_     ):
						lowerCAmelCase__      :      Tuple         =		{
						    '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(SCREAMING_SNAKE_CASE_       ,							SCREAMING_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.
			lowerCAmelCase__      :      List[Any]         =		DataCollatorForLanguageModeling(
			    tokenizer=SCREAMING_SNAKE_CASE_       ,							mlm_probability=args.mlm_probability       ,							mlm=SCREAMING_SNAKE_CASE_       ,							return_tensors='tf'     )

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

			lowerCAmelCase__      :      Optional[int]         =		args.per_replica_batch_size * strategy.num_replicas_in_sync

			lowerCAmelCase__      :      int         =		prepare_dataset(
			    SCREAMING_SNAKE_CASE_       ,							decode_fn=SCREAMING_SNAKE_CASE_       ,							mask_fn=SCREAMING_SNAKE_CASE_       ,							batch_size=SCREAMING_SNAKE_CASE_       ,							shuffle=SCREAMING_SNAKE_CASE_       ,							shuffle_buffer_size=args.shuffle_buffer_size       ,							)

			lowerCAmelCase__      :      Tuple         =		prepare_dataset(
			    SCREAMING_SNAKE_CASE_       ,							decode_fn=SCREAMING_SNAKE_CASE_       ,							mask_fn=SCREAMING_SNAKE_CASE_       ,							batch_size=SCREAMING_SNAKE_CASE_       ,							shuffle=SCREAMING_SNAKE_CASE_       ,							)

			lowerCAmelCase__      :      Dict         =		[]
			if args.hub_model_id:
						callbacks.append(
						    PushToHubCallback(output_dir=args.output_dir       ,							hub_model_id=args.hub_model_id       ,							tokenizer=SCREAMING_SNAKE_CASE_     )     )

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

			model.save_pretrained(args.output_dir     )


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

from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


lowerCamelCase__          =		logging.get_logger(__name__)

lowerCamelCase__          =		{
    """EleutherAI/gpt-j-6B""": """https://huggingface.co/EleutherAI/gpt-j-6B/resolve/main/config.json""",
    # See all GPT-J models at https://huggingface.co/models?filter=gpt_j
}







class      A__							(     __magic_name__	):
				lowercase			    =  'gptj'
				lowercase			    =  {
				    'max_position_embeddings': 'n_positions',
				    'hidden_size': 'n_embd',
				    'num_attention_heads': 'n_head',
				    'num_hidden_layers': 'n_layer',
				}






				def __init__(						self					:      Optional[Any]    ,     a					:      Dict=50_400    ,     a					:      Union[str, Any]=2_048    ,     a					:      List[str]=4_096    ,     a					:      Any=28    ,     a					:      Optional[Any]=16    ,     a					:      Optional[Any]=64    ,     a					:      int=None    ,     a					:      Any="gelu_new"    ,     a					:      Union[str, Any]=0.0    ,     a					:      List[Any]=0.0    ,     a					:      List[Any]=0.0    ,     a					:      Optional[Any]=1E-5    ,     a					:      Optional[int]=0.0_2    ,     a					:      int=True    ,     a					:      str=50_256    ,     a					:      str=50_256    ,     a					:      Any=False    ,     **a					:      Dict    ,     ):






							'''simple docstring'''







							lowerCAmelCase__      :      List[Any]         =		vocab_size
							lowerCAmelCase__      :      List[Any]         =		n_positions
							lowerCAmelCase__      :      str         =		n_embd
							lowerCAmelCase__      :      str         =		n_layer
							lowerCAmelCase__      :      str         =		n_head
							lowerCAmelCase__      :      Dict         =		n_inner
							lowerCAmelCase__      :      Union[str, Any]         =		rotary_dim
							lowerCAmelCase__      :      Optional[int]         =		activation_function
							lowerCAmelCase__      :      Any         =		resid_pdrop
							lowerCAmelCase__      :      int         =		embd_pdrop
							lowerCAmelCase__      :      int         =		attn_pdrop
							lowerCAmelCase__      :      List[Any]         =		layer_norm_epsilon
							lowerCAmelCase__      :      str         =		initializer_range
							lowerCAmelCase__      :      Dict         =		use_cache

							lowerCAmelCase__      :      str         =		bos_token_id
							lowerCAmelCase__      :      int         =		eos_token_id

							super().__init__(
							    bos_token_id=a    ,     eos_token_id=a    ,     tie_word_embeddings=a    ,     **a   )







class      A__							(     __magic_name__	):






				def __init__(						self					:      str    ,     a					:      PretrainedConfig    ,     a					:      str = "default"    ,     a					:      List[PatchingSpec] = None    ,     a					:      bool = False    ,     ):






							'''simple docstring'''







							super().__init__(a    ,     task=a    ,     patching_specs=a    ,     use_past=a   )
							if not getattr(self._config    ,     'pad_token_id'    ,     a   ):
										# TODO: how to do that better?
										lowerCAmelCase__      :      int         =		0






				@property
				def 							_lowerCamelCase		(						self					:      Tuple   ):






							'''simple docstring'''







							lowerCAmelCase__      :      Dict         =		OrderedDict({'input_ids': {0: 'batch', 1: 'sequence'}}   )
							if self.use_past:
										self.fill_with_past_key_values_(a    ,     direction='inputs'   )
										lowerCAmelCase__      :      Optional[Any]         =		{0: 'batch', 1: 'past_sequence + sequence'}
							else:
										lowerCAmelCase__      :      Tuple         =		{0: 'batch', 1: 'sequence'}

							return common_inputs






				@property
				def 							_lowerCamelCase		(						self					:      Union[str, Any]   ):






							'''simple docstring'''







							return self._config.n_layer






				@property
				def 							_lowerCamelCase		(						self					:      Optional[Any]   ):






							'''simple docstring'''







							return self._config.n_head






				def 							_lowerCamelCase		(						self					:      Tuple    ,     a					:      PreTrainedTokenizer    ,     a					:      int = -1    ,     a					:      int = -1    ,     a					:      bool = False    ,     a					:      Optional[TensorType] = None    ,     ):






							'''simple docstring'''







							lowerCAmelCase__      :      Tuple         =		super(a    ,     self   ).generate_dummy_inputs(
							    a    ,     batch_size=a    ,     seq_length=a    ,     is_pair=a    ,     framework=a   )

							# We need to order the input in the way they appears in the forward()
							lowerCAmelCase__      :      Optional[int]         =		OrderedDict({'input_ids': common_inputs['input_ids']}   )

							# Need to add the past_keys
							if self.use_past:
										if not is_torch_available():
													raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'   )
										else:
													import torch

													lowerCAmelCase__      ,				lowerCAmelCase__      :      int         =		common_inputs['input_ids'].shape
													# Not using the same length for past_key_values
													lowerCAmelCase__      :      Optional[int]         =		seqlen + 2
													lowerCAmelCase__      :      Tuple         =		(
													    batch,
													    self.num_attention_heads,
													    past_key_values_length,
													    self._config.hidden_size // self.num_attention_heads,
													)
													lowerCAmelCase__      :      Tuple         =		[
													    (torch.zeros(a   ), torch.zeros(a   )) for _ in range(self.num_layers   )
													]

							lowerCAmelCase__      :      Any         =		common_inputs['attention_mask']
							if self.use_past:
										lowerCAmelCase__      :      List[str]         =		ordered_inputs['attention_mask'].dtype
										lowerCAmelCase__      :      Optional[Any]         =		torch.cat(
										    [ordered_inputs['attention_mask'], torch.ones(a    ,     a    ,     dtype=a   )]    ,     dim=1   )

							return ordered_inputs






				@property
				def 							_lowerCamelCase		(						self					:      int   ):






							'''simple docstring'''







							return 13