





'''simple docstring'''


import argparse
import logging
import os
import re

import tensorflow as tf

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


_SCREAMING_SNAKE_CASE	 =							logging.getLogger(__name__)

_SCREAMING_SNAKE_CASE	 =							tf.data.AUTOTUNE


def 	__a():



				'''simple docstring'''





				_lowerCAmelCase				=      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				=      parser.parse_args()
				return args


def 	__a(SCREAMING_SNAKE_CASE_  :					List[Any]    ):



				'''simple docstring'''





				try:
								if args.tpu_name:
												_lowerCAmelCase				=      tf.distribute.cluster_resolver.TPUClusterResolver(
												    args.tpu_name		,       zone=args.tpu_zone		,       project=args.gcp_project    )
								else:
												_lowerCAmelCase				=      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 	__a(SCREAMING_SNAKE_CASE_  :					int    ):



				'''simple docstring'''





				_lowerCAmelCase				=      0
				for file in file_list:
								_lowerCAmelCase				=      file.split("/"    )[-1]
								_lowerCAmelCase				=      re.search(R"-\d+-(\d+)\.tfrecord"		,       SCREAMING_SNAKE_CASE_    ).group(1    )
								_lowerCAmelCase				=      int(SCREAMING_SNAKE_CASE_    )
								num_samples += sample_count

				return num_samples


def 	__a(SCREAMING_SNAKE_CASE_  :					Tuple		,       SCREAMING_SNAKE_CASE_  :					str		,       SCREAMING_SNAKE_CASE_  :					Any		,       SCREAMING_SNAKE_CASE_  :					List[str]		,       SCREAMING_SNAKE_CASE_  :					List[Any]		,       SCREAMING_SNAKE_CASE_  :					Tuple=None    ):



				'''simple docstring'''





				_lowerCAmelCase				=      count_samples(SCREAMING_SNAKE_CASE_    )
				_lowerCAmelCase				=      tf.data.Dataset.from_tensor_slices(SCREAMING_SNAKE_CASE_    )
				if shuffle:
								_lowerCAmelCase				=      dataset.shuffle(len(SCREAMING_SNAKE_CASE_    )    )
				_lowerCAmelCase				=      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				=      dataset.apply(tf.data.experimental.assert_cardinality(SCREAMING_SNAKE_CASE_    )    )
				_lowerCAmelCase				=      dataset.map(SCREAMING_SNAKE_CASE_		,       num_parallel_calls=SCREAMING_SNAKE_CASE_    )
				if shuffle:
								assert shuffle_buffer_size is not None
								_lowerCAmelCase				=      dataset.shuffle(args.shuffle_buffer_size    )
				_lowerCAmelCase				=      dataset.batch(SCREAMING_SNAKE_CASE_		,       drop_remainder=SCREAMING_SNAKE_CASE_    )
				_lowerCAmelCase				=      dataset.map(SCREAMING_SNAKE_CASE_		,       num_parallel_calls=SCREAMING_SNAKE_CASE_    )
				_lowerCAmelCase				=      dataset.prefetch(SCREAMING_SNAKE_CASE_    )
				return dataset


def 	__a(SCREAMING_SNAKE_CASE_  :					Optional[Any]    ):



				'''simple docstring'''





				if not args.no_tpu:
								_lowerCAmelCase				=      initialize_tpu(SCREAMING_SNAKE_CASE_    )
								_lowerCAmelCase				=      tf.distribute.TPUStrategy(SCREAMING_SNAKE_CASE_    )
				else:
								_lowerCAmelCase				=      tf.distribute.OneDeviceStrategy(device="/gpu:0"    )

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

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

				_lowerCAmelCase				=      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				=      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				=      count_samples(SCREAMING_SNAKE_CASE_    )

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

				with strategy.scope():
								_lowerCAmelCase				=      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				=      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_  :					Optional[int]    ):
								_lowerCAmelCase				=      {
								    "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				=      DataCollatorForLanguageModeling(
				    tokenizer=SCREAMING_SNAKE_CASE_		,       mlm_probability=args.mlm_probability		,       mlm=SCREAMING_SNAKE_CASE_		,       return_tensors="tf"    )

				def mask_with_collator(SCREAMING_SNAKE_CASE_  :					str    ):
								# TF really needs an isin() function
								_lowerCAmelCase				=      (
								    ~tf.cast(batch["attention_mask"]		,       tf.bool    )
								    | (batch["input_ids"] == tokenizer.cls_token_id)
								    | (batch["input_ids"] == tokenizer.sep_token_id)
								)
								_lowerCAmelCase			,			_lowerCAmelCase				=      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				=      args.per_replica_batch_size * strategy.num_replicas_in_sync

				_lowerCAmelCase				=      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				=      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				=      []
				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__":
							_SCREAMING_SNAKE_CASE	 =							parse_args()
							main(args)









'''simple docstring'''


import argparse
import collections
import os
import re

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_table.py
_SCREAMING_SNAKE_CASE	 =							"src/transformers"
_SCREAMING_SNAKE_CASE	 =							"docs/source/en"
_SCREAMING_SNAKE_CASE	 =							"."


def 	__a(SCREAMING_SNAKE_CASE_  :					int		,       SCREAMING_SNAKE_CASE_  :					Tuple		,       SCREAMING_SNAKE_CASE_  :					Union[str, Any]    ):



    '''simple docstring'''





    with open(SCREAMING_SNAKE_CASE_		,       "r"		,       encoding="utf-8"		,       newline="\n"    ) as f:
        _lowerCAmelCase				=      f.readlines()
    # Find the start prompt.
    _lowerCAmelCase				=      0
    while not lines[start_index].startswith(SCREAMING_SNAKE_CASE_    ):
        start_index += 1
    start_index += 1

    _lowerCAmelCase				=      start_index
    while not lines[end_index].startswith(SCREAMING_SNAKE_CASE_    ):
        end_index += 1
    end_index -= 1

    while len(lines[start_index]    ) <= 1:
        start_index += 1
    while len(lines[end_index]    ) <= 1:
        end_index -= 1
    end_index += 1
    return "".join(lines[start_index:end_index]    ), start_index, end_index, lines


# Add here suffixes that are used to identify models, separated by |
_SCREAMING_SNAKE_CASE	 =							"Model|Encoder|Decoder|ForConditionalGeneration"
# Regexes that match TF/Flax/PT model names.
_SCREAMING_SNAKE_CASE	 =							re.compile(r"TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
_SCREAMING_SNAKE_CASE	 =							re.compile(r"Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
_SCREAMING_SNAKE_CASE	 =							re.compile(r"(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")


# This is to make sure the transformers module imported is the one in the repo.
_SCREAMING_SNAKE_CASE	 =							direct_transformers_import(TRANSFORMERS_PATH)


def 	__a(SCREAMING_SNAKE_CASE_  :					Union[str, Any]    ):



    '''simple docstring'''





    _lowerCAmelCase				=      re.finditer(".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)"		,       SCREAMING_SNAKE_CASE_    )
    return [m.group(0    ) for m in matches]


def 	__a(SCREAMING_SNAKE_CASE_  :					Dict		,       SCREAMING_SNAKE_CASE_  :					int    ):



    '''simple docstring'''





    _lowerCAmelCase				=      2 if text == "✅" or text == "❌" else len(SCREAMING_SNAKE_CASE_    )
    _lowerCAmelCase				=      (width - text_length) // 2
    _lowerCAmelCase				=      width - text_length - left_indent
    return " " * left_indent + text + " " * right_indent


def 	__a():



    '''simple docstring'''





    _lowerCAmelCase				=      transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
    _lowerCAmelCase				=      {
        name: config_maping_names[code]
        for code, name in transformers_module.MODEL_NAMES_MAPPING.items()
        if code in config_maping_names
    }
    _lowerCAmelCase				=      {name: config.replace("Config"		,       ""    ) for name, config in model_name_to_config.items()}

    # Dictionaries flagging if each model prefix has a slow/fast tokenizer, backend in PT/TF/Flax.
    _lowerCAmelCase				=      collections.defaultdict(SCREAMING_SNAKE_CASE_    )
    _lowerCAmelCase				=      collections.defaultdict(SCREAMING_SNAKE_CASE_    )
    _lowerCAmelCase				=      collections.defaultdict(SCREAMING_SNAKE_CASE_    )
    _lowerCAmelCase				=      collections.defaultdict(SCREAMING_SNAKE_CASE_    )
    _lowerCAmelCase				=      collections.defaultdict(SCREAMING_SNAKE_CASE_    )

    # Let's lookup through all transformers object (once).
    for attr_name in dir(SCREAMING_SNAKE_CASE_    ):
        _lowerCAmelCase				=      None
        if attr_name.endswith("Tokenizer"    ):
            _lowerCAmelCase				=      slow_tokenizers
            _lowerCAmelCase				=      attr_name[:-9]
        elif attr_name.endswith("TokenizerFast"    ):
            _lowerCAmelCase				=      fast_tokenizers
            _lowerCAmelCase				=      attr_name[:-13]
        elif _re_tf_models.match(SCREAMING_SNAKE_CASE_    ) is not None:
            _lowerCAmelCase				=      tf_models
            _lowerCAmelCase				=      _re_tf_models.match(SCREAMING_SNAKE_CASE_    ).groups()[0]
        elif _re_flax_models.match(SCREAMING_SNAKE_CASE_    ) is not None:
            _lowerCAmelCase				=      flax_models
            _lowerCAmelCase				=      _re_flax_models.match(SCREAMING_SNAKE_CASE_    ).groups()[0]
        elif _re_pt_models.match(SCREAMING_SNAKE_CASE_    ) is not None:
            _lowerCAmelCase				=      pt_models
            _lowerCAmelCase				=      _re_pt_models.match(SCREAMING_SNAKE_CASE_    ).groups()[0]

        if lookup_dict is not None:
            while len(SCREAMING_SNAKE_CASE_    ) > 0:
                if attr_name in model_name_to_prefix.values():
                    _lowerCAmelCase				=      True
                    break
                # Try again after removing the last word in the name
                _lowerCAmelCase				=      "".join(camel_case_split(SCREAMING_SNAKE_CASE_    )[:-1]    )

    # Let's build that table!
    _lowerCAmelCase				=      list(model_name_to_config.keys()    )
    model_names.sort(key=str.lower    )
    _lowerCAmelCase				=      ["Model", "Tokenizer slow", "Tokenizer fast", "PyTorch support", "TensorFlow support", "Flax Support"]
    # We'll need widths to properly display everything in the center (+2 is to leave one extra space on each side).
    _lowerCAmelCase				=      [len(SCREAMING_SNAKE_CASE_    ) + 2 for c in columns]
    _lowerCAmelCase				=      max([len(SCREAMING_SNAKE_CASE_    ) for name in model_names]    ) + 2

    # Build the table per se
    _lowerCAmelCase				=      "|" + "|".join([_center_text(SCREAMING_SNAKE_CASE_		,       SCREAMING_SNAKE_CASE_    ) for c, w in zip(SCREAMING_SNAKE_CASE_		,       SCREAMING_SNAKE_CASE_    )]    ) + "|\n"
    # Use ":-----:" format to center-aligned table cell texts
    table += "|" + "|".join([":" + "-" * (w - 2) + ":" for w in widths]    ) + "|\n"

    _lowerCAmelCase				=      {True: "✅", False: "❌"}
    for name in model_names:
        _lowerCAmelCase				=      model_name_to_prefix[name]
        _lowerCAmelCase				=      [
            name,
            check[slow_tokenizers[prefix]],
            check[fast_tokenizers[prefix]],
            check[pt_models[prefix]],
            check[tf_models[prefix]],
            check[flax_models[prefix]],
        ]
        table += "|" + "|".join([_center_text(SCREAMING_SNAKE_CASE_		,       SCREAMING_SNAKE_CASE_    ) for l, w in zip(SCREAMING_SNAKE_CASE_		,       SCREAMING_SNAKE_CASE_    )]    ) + "|\n"
    return table


def 	__a(SCREAMING_SNAKE_CASE_  :					Union[str, Any]=False    ):



    '''simple docstring'''





    _lowerCAmelCase			,			_lowerCAmelCase			,			_lowerCAmelCase			,			_lowerCAmelCase				=      _find_text_in_file(
        filename=os.path.join(SCREAMING_SNAKE_CASE_		,       "index.md"    )		,       start_prompt="<!--This table is updated automatically from the auto modules"		,       end_prompt="<!-- End table-->"		,       )
    _lowerCAmelCase				=      get_model_table_from_auto_modules()

    if current_table != new_table:
        if overwrite:
            with open(os.path.join(SCREAMING_SNAKE_CASE_		,       "index.md"    )		,       "w"		,       encoding="utf-8"		,       newline="\n"    ) as f:
                f.writelines(lines[:start_index] + [new_table] + lines[end_index:]    )
        else:
            raise ValueError(
                "The model table in the `index.md` has not been updated. Run `make fix-copies` to fix this."    )


if __name__ == "__main__":
       _SCREAMING_SNAKE_CASE	 =							argparse.ArgumentParser()
       parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
       _SCREAMING_SNAKE_CASE	 =							parser.parse_args()

       check_model_table(args.fix_and_overwrite)


