




"""simple docstring"""


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 	lowercase_				(			)						->    Tuple:
      lowerCAmelCase__      :			List[Any]   =	argparse.ArgumentParser(description="""Train a masked language model on TPU."""				)
      parser.add_argument(
          """--pretrained_model_config"""    ,		type=__UpperCAmelCase    ,		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=__UpperCAmelCase    ,		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=__UpperCAmelCase    ,		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=__UpperCAmelCase    ,		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=__UpperCAmelCase    ,		help="""Google cloud zone that TPU resource is located in. Only used for non-Colab TPU nodes."""    ,		)

      parser.add_argument(
          """--gcp_project"""    ,		type=__UpperCAmelCase    ,		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=__UpperCAmelCase    ,		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=__UpperCAmelCase    ,		default=2**18    ,		help="""Size of the shuffle buffer (in samples)"""    ,		)

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

      parser.add_argument(
          """--learning_rate"""    ,		type=__UpperCAmelCase    ,		default=1E-4    ,		help="""Learning rate to use for training."""    ,		)

      parser.add_argument(
          """--weight_decay_rate"""    ,		type=__UpperCAmelCase    ,		default=1E-3    ,		help="""Weight decay rate to use for training."""    ,		)

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

      parser.add_argument(
          """--mlm_probability"""    ,		type=__UpperCAmelCase    ,		default=0.15    ,		help="""Fraction of tokens to mask during training."""    ,		)

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

      lowerCAmelCase__      :			Tuple   =	parser.parse_args()
      return args



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

      return tpu



def 	lowercase_				(			__UpperCAmelCase				)						->    Optional[int]:
      lowerCAmelCase__      :			Optional[Any]   =	0
      for file in file_list:
            lowerCAmelCase__      :			List[str]   =	file.split("""/"""				)[-1]
            lowerCAmelCase__      :			Optional[int]   =	re.search(R"""-\d+-(\d+)\.tfrecord"""    ,		__UpperCAmelCase				).group(1				)
            lowerCAmelCase__      :			Dict   =	int(__UpperCAmelCase				)
            num_samples += sample_count

      return num_samples



def 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase=None				)						->    Any:
      lowerCAmelCase__      :			List[str]   =	count_samples(__UpperCAmelCase				)
      lowerCAmelCase__      :			str   =	tf.data.Dataset.from_tensor_slices(__UpperCAmelCase				)
      if shuffle:
            lowerCAmelCase__      :			List[Any]   =	dataset.shuffle(len(__UpperCAmelCase				)				)
      lowerCAmelCase__      :			Optional[int]   =	tf.data.TFRecordDataset(__UpperCAmelCase    ,		num_parallel_reads=__UpperCAmelCase				)
      # TF can't infer the total sample count because it doesn't read all the records yet, so we assert it here
      lowerCAmelCase__      :			int   =	dataset.apply(tf.data.experimental.assert_cardinality(__UpperCAmelCase				)				)
      lowerCAmelCase__      :			List[Any]   =	dataset.map(__UpperCAmelCase    ,		num_parallel_calls=__UpperCAmelCase				)
      if shuffle:
            assert shuffle_buffer_size is not None
            lowerCAmelCase__      :			List[Any]   =	dataset.shuffle(args.shuffle_buffer_size				)
      lowerCAmelCase__      :			int   =	dataset.batch(__UpperCAmelCase    ,		drop_remainder=__UpperCAmelCase				)
      lowerCAmelCase__      :			Any   =	dataset.map(__UpperCAmelCase    ,		num_parallel_calls=__UpperCAmelCase				)
      lowerCAmelCase__      :			List[Any]   =	dataset.prefetch(__UpperCAmelCase				)
      return dataset



def 	lowercase_				(			__UpperCAmelCase				)						->    int:
      if not args.no_tpu:
            lowerCAmelCase__      :			Optional[int]   =	initialize_tpu(__UpperCAmelCase				)
            lowerCAmelCase__      :			Tuple   =	tf.distribute.TPUStrategy(__UpperCAmelCase				)
      else:
            lowerCAmelCase__      :			Union[str, Any]   =	tf.distribute.OneDeviceStrategy(device="""/gpu:0"""				)

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

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

      lowerCAmelCase__      :			Any   =	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__      :			Union[str, 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__      :			str   =	count_samples(__UpperCAmelCase				)

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

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

      def decode_fn(__UpperCAmelCase				):
            lowerCAmelCase__      :			str   =	{
                """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(__UpperCAmelCase    ,		__UpperCAmelCase				)

      # 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__      :			int   =	DataCollatorForLanguageModeling(
          tokenizer=__UpperCAmelCase    ,		mlm_probability=args.mlm_probability    ,		mlm=__UpperCAmelCase    ,		return_tensors="""tf"""				)

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

      lowerCAmelCase__      :			int   =	args.per_replica_batch_size * strategy.num_replicas_in_sync

      lowerCAmelCase__      :			Optional[int]   =	prepare_dataset(
          __UpperCAmelCase    ,		decode_fn=__UpperCAmelCase    ,		mask_fn=__UpperCAmelCase    ,		batch_size=__UpperCAmelCase    ,		shuffle=__UpperCAmelCase    ,		shuffle_buffer_size=args.shuffle_buffer_size    ,		)

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

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

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

      model.save_pretrained(args.output_dir				)


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











"""simple docstring"""


import logging
import os
from typing import List, TextIO, Union

from conllu import parse_incr
from utils_ner import InputExample, Split, TokenClassificationTask


_A													=      logging.getLogger(__name__)


class _lowerCamelCase    (					a_						):






   def __init__(		self       :		List[Any]   ,     UpperCamelCase       :		Dict=-1   )			->						List[Any]:

         """simple docstring"""
         # in NER datasets, the last column is usually reserved for NER label
         lowerCAmelCase__      :			Optional[int]   =	label_idx






   def 	_lowerCAmelCase   (		self       :		Optional[int]   ,     UpperCamelCase       :		Optional[int]   ,     UpperCamelCase       :		Union[Split, str]   )			->						List[InputExample]:

         """simple docstring"""
         if isinstance(UpperCamelCase   ,     UpperCamelCase   ):
               lowerCAmelCase__      :			int   =	mode.value
         lowerCAmelCase__      :			Dict   =	os.path.join(UpperCamelCase   ,     f"""{mode}.txt"""   )
         lowerCAmelCase__      :			str   =	1
         lowerCAmelCase__      :			List[Any]   =	[]
         with open(UpperCamelCase   ,     encoding="""utf-8"""   ) as f:
               lowerCAmelCase__      :			str   =	[]
               lowerCAmelCase__      :			Any   =	[]
               for line in f:
                     if line.startswith("""-DOCSTART-"""   ) or line == "" or line == "\n":
                           if words:
                                 examples.append(InputExample(guid=f"""{mode}-{guid_index}"""   ,     words=UpperCamelCase   ,     labels=UpperCamelCase   )   )
                                 guid_index += 1
                                 lowerCAmelCase__      :			Optional[Any]   =	[]
                                 lowerCAmelCase__      :			Tuple   =	[]
                     else:
                           lowerCAmelCase__      :			Optional[int]   =	line.split(""" """   )
                           words.append(splits[0]   )
                           if len(UpperCamelCase   ) > 1:
                                 labels.append(splits[self.label_idx].replace("""\n"""   ,     """"""   )   )
                           else:
                                 # Examples could have no label for mode = "test"
                                 labels.append("""O"""   )
               if words:
                     examples.append(InputExample(guid=f"""{mode}-{guid_index}"""   ,     words=UpperCamelCase   ,     labels=UpperCamelCase   )   )
         return examples






   def 	_lowerCAmelCase   (		self       :		Any   ,     UpperCamelCase       :		TextIO   ,     UpperCamelCase       :		TextIO   ,     UpperCamelCase       :		List   )			->						List[Any]:

         """simple docstring"""
         lowerCAmelCase__      :			Dict   =	0
         for line in test_input_reader:
               if line.startswith("""-DOCSTART-"""   ) or line == "" or line == "\n":
                     writer.write(UpperCamelCase   )
                     if not preds_list[example_id]:
                           example_id += 1
               elif preds_list[example_id]:
                     lowerCAmelCase__      :			Union[str, Any]   =	line.split()[0] + """ """ + preds_list[example_id].pop(0   ) + """\n"""
                     writer.write(UpperCamelCase   )
               else:
                     logger.warning("""Maximum sequence length exceeded: No prediction for '%s'."""   ,     line.split()[0]   )






   def 	_lowerCAmelCase   (		self       :		str   ,     UpperCamelCase       :		str   )			->						List[str]:

         """simple docstring"""
         if path:
               with open(UpperCamelCase   ,     """r"""   ) as f:
                     lowerCAmelCase__      :			Any   =	f.read().splitlines()
               if "O" not in labels:
                     lowerCAmelCase__      :			List[str]   =	["""O"""] + labels
               return labels
         else:
               return ["O", "B-MISC", "I-MISC", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]






class _lowerCamelCase    (					a_						):






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

         """simple docstring"""
         # in CONLL2003 dataset chunk column is second-to-last
         super().__init__(label_idx=-2   )






   def 	_lowerCAmelCase   (		self       :		Optional[Any]   ,     UpperCamelCase       :		str   )			->						List[str]:

         """simple docstring"""
         if path:
               with open(UpperCamelCase   ,     """r"""   ) as f:
                     lowerCAmelCase__      :			Any   =	f.read().splitlines()
               if "O" not in labels:
                     lowerCAmelCase__      :			str   =	["""O"""] + labels
               return labels
         else:
               return [
                   "O",
                   "B-ADVP",
                   "B-INTJ",
                   "B-LST",
                   "B-PRT",
                   "B-NP",
                   "B-SBAR",
                   "B-VP",
                   "B-ADJP",
                   "B-CONJP",
                   "B-PP",
                   "I-ADVP",
                   "I-INTJ",
                   "I-LST",
                   "I-PRT",
                   "I-NP",
                   "I-SBAR",
                   "I-VP",
                   "I-ADJP",
                   "I-CONJP",
                   "I-PP",
               ]






class _lowerCamelCase    (					a_						):






   def 	_lowerCAmelCase   (		self       :		Union[str, Any]   ,     UpperCamelCase       :		Any   ,     UpperCamelCase       :		Union[Split, str]   )			->						List[InputExample]:

         """simple docstring"""
         if isinstance(UpperCamelCase   ,     UpperCamelCase   ):
               lowerCAmelCase__      :			Optional[Any]   =	mode.value
         lowerCAmelCase__      :			int   =	os.path.join(UpperCamelCase   ,     f"""{mode}.txt"""   )
         lowerCAmelCase__      :			Optional[int]   =	1
         lowerCAmelCase__      :			List[Any]   =	[]

         with open(UpperCamelCase   ,     encoding="""utf-8"""   ) as f:
               for sentence in parse_incr(UpperCamelCase   ):
                     lowerCAmelCase__      :			Union[str, Any]   =	[]
                     lowerCAmelCase__      :			List[Any]   =	[]
                     for token in sentence:
                           words.append(token["""form"""]   )
                           labels.append(token["""upos"""]   )
                     assert len(UpperCamelCase   ) == len(UpperCamelCase   )
                     if words:
                           examples.append(InputExample(guid=f"""{mode}-{guid_index}"""   ,     words=UpperCamelCase   ,     labels=UpperCamelCase   )   )
                           guid_index += 1
         return examples






   def 	_lowerCAmelCase   (		self       :		List[str]   ,     UpperCamelCase       :		TextIO   ,     UpperCamelCase       :		TextIO   ,     UpperCamelCase       :		List   )			->						List[Any]:

         """simple docstring"""
         lowerCAmelCase__      :			Any   =	0
         for sentence in parse_incr(UpperCamelCase   ):
               lowerCAmelCase__      :			Union[str, Any]   =	preds_list[example_id]
               lowerCAmelCase__      :			List[Any]   =	""""""
               for token in sentence:
                     out += f"""{token['form']} ({token['upos']}|{s_p.pop(0   )}) """
               out += "\n"
               writer.write(UpperCamelCase   )
               example_id += 1






   def 	_lowerCAmelCase   (		self       :		Dict   ,     UpperCamelCase       :		str   )			->						List[str]:

         """simple docstring"""
         if path:
               with open(UpperCamelCase   ,     """r"""   ) as f:
                     return f.read().splitlines()
         else:
               return [
                   "ADJ",
                   "ADP",
                   "ADV",
                   "AUX",
                   "CCONJ",
                   "DET",
                   "INTJ",
                   "NOUN",
                   "NUM",
                   "PART",
                   "PRON",
                   "PROPN",
                   "PUNCT",
                   "SCONJ",
                   "SYM",
                   "VERB",
                   "X",
               ]





