


import argparse
import json
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


snake_case__				:	str					       =					16
snake_case__				:	List[Any]					       =					32


def    __lowerCamelCase   (      A__						: Accelerator						,						A__						: int = 16						,						A__						: str = "bert-base-cased"					)     ->    int:
   lowerCamelCase_						:   List[Any]		= AutoTokenizer.from_pretrained(A__					)
   lowerCamelCase_						:   Any		= load_dataset("""glue"""						,						"""mrpc"""					)

   def tokenize_function(A__						: Dict					):
      # max_length=None => use the model max length (it's actually the default)
      lowerCamelCase_						:   str		= tokenizer(examples["""sentence1"""]						,						examples["""sentence2"""]						,						truncation=A__						,						max_length=A__					)
      return outputs

   # Apply the method we just defined to all the examples in all the splits of the dataset
   lowerCamelCase_						:   List[str]		= datasets.map(
       A__						,						batched=A__						,						remove_columns=["""idx""", """sentence1""", """sentence2"""]						,						load_from_cache_file=A__					)

   # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
   # transformers library
   lowerCamelCase_						:   List[Any]		= tokenized_datasets.rename_column("""label"""						,						"""labels"""					)

   def collate_fn(A__						: List[str]					):
      # On TPU it's best to pad everything to the same length or training will be very slow.
      if accelerator.distributed_type == DistributedType.TPU:
         return tokenizer.pad(A__						,						padding="""max_length"""						,						max_length=128						,						return_tensors="""pt"""					)
      return tokenizer.pad(A__						,						padding="""longest"""						,						return_tensors="""pt"""					)

   # Instantiate dataloaders.
   lowerCamelCase_						:   Dict		= DataLoader(
       tokenized_datasets["""train"""]						,						shuffle=A__						,						collate_fn=A__						,						batch_size=A__					)
   lowerCamelCase_						:   List[Any]		= DataLoader(
       tokenized_datasets["""validation"""]						,						shuffle=A__						,						collate_fn=A__						,						batch_size=A__					)

   return train_dataloader, eval_dataloader


def    __lowerCamelCase   (      A__						: Union[str, Any]						,						A__						: Dict					)     ->    int:
   # Initialize accelerator
   lowerCamelCase_						:   Tuple		= Accelerator()

   # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
   lowerCamelCase_						:   List[str]		= config["""lr"""]
   lowerCamelCase_						:   Tuple		= int(config["""num_epochs"""]					)
   lowerCamelCase_						:   Optional[int]		= int(config["""seed"""]					)
   lowerCamelCase_						:   Any		= int(config["""batch_size"""]					)
   lowerCamelCase_						:   List[str]		= args.model_name_or_path

   set_seed(A__					)
   lowerCamelCase_,	lowerCamelCase_						:   Optional[Any]		= get_dataloaders(A__						,						A__						,						A__					)

   # Instantiate the model (we build the model here so that the seed also control new weights initialization)
   lowerCamelCase_						:   int		= AutoModelForSequenceClassification.from_pretrained(A__						,						return_dict=A__					)

   # Instantiate optimizer
   lowerCamelCase_						:   Tuple		= (
       AdamW
       if accelerator.state.deepspeed_plugin is None
       or """optimizer""" not in accelerator.state.deepspeed_plugin.deepspeed_config
       else DummyOptim
   )
   lowerCamelCase_						:   Tuple		= optimizer_cls(params=model.parameters()						,						lr=A__					)

   if accelerator.state.deepspeed_plugin is not None:
      lowerCamelCase_						:   int		= accelerator.state.deepspeed_plugin.deepspeed_config[
          """gradient_accumulation_steps"""
      ]
   else:
      lowerCamelCase_						:   List[Any]		= 1
   lowerCamelCase_						:   int		= (len(A__					) * num_epochs) // gradient_accumulation_steps

   # Instantiate scheduler
   if (
       accelerator.state.deepspeed_plugin is None
       or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
   ):
      lowerCamelCase_						:   Any		= get_linear_schedule_with_warmup(
          optimizer=A__						,						num_warmup_steps=0						,						num_training_steps=A__						,						)
   else:
      lowerCamelCase_						:   List[str]		= DummyScheduler(A__						,						total_num_steps=A__						,						warmup_num_steps=0					)

   # Prepare everything
   # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
   # prepare method.
   lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_						:   Any		= accelerator.prepare(
       A__						,						A__						,						A__						,						A__						,						A__					)

   # We need to keep track of how many total steps we have iterated over
   lowerCamelCase_						:   Tuple		= 0
   # We also need to keep track of the stating epoch so files are named properly
   lowerCamelCase_						:   Dict		= 0

   # Now we train the model
   lowerCamelCase_						:   Any		= evaluate.load("""glue"""						,						"""mrpc"""					)
   lowerCamelCase_						:   Dict		= 0
   lowerCamelCase_						:   Optional[Any]		= {}
   for epoch in range(A__						,						A__					):
      model.train()
      for step, batch in enumerate(A__					):
         lowerCamelCase_						:   str		= model(**A__					)
         lowerCamelCase_						:   List[Any]		= outputs.loss
         lowerCamelCase_						:   Optional[int]		= loss / gradient_accumulation_steps
         accelerator.backward(A__					)
         if step % gradient_accumulation_steps == 0:
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

         overall_step += 1

      model.eval()
      lowerCamelCase_						:   Tuple		= 0
      for step, batch in enumerate(A__					):
         # We could avoid this line since we set the accelerator with `device_placement=True`.
         batch.to(accelerator.device					)
         with torch.no_grad():
            lowerCamelCase_						:   str		= model(**A__					)
         lowerCamelCase_						:   str		= outputs.logits.argmax(dim=-1					)
         # It is slightly faster to call this once, than multiple times
         lowerCamelCase_,	lowerCamelCase_						:   Union[str, Any]		= accelerator.gather(
             (predictions, batch["""labels"""])					)  # If we are in a multiprocess environment, the last batch has duplicates
         if accelerator.use_distributed:
            if step == len(A__					) - 1:
               lowerCamelCase_						:   List[str]		= predictions[: len(eval_dataloader.dataset					) - samples_seen]
               lowerCamelCase_						:   Any		= references[: len(eval_dataloader.dataset					) - samples_seen]
            else:
               samples_seen += references.shape[0]
         metric.add_batch(
             predictions=A__						,						references=A__						,						)

      lowerCamelCase_						:   Optional[Any]		= metric.compute()
      # Use accelerator.print to print only on the main process.
      accelerator.print(f'''epoch {epoch}:'''						,						A__					)
      lowerCamelCase_						:   List[Any]		= eval_metric["""accuracy"""]

      if best_performance < eval_metric["accuracy"]:
         lowerCamelCase_						:   str		= eval_metric["""accuracy"""]

   if args.performance_lower_bound is not None:
      assert (
          args.performance_lower_bound <= best_performance
      ), f'''Best performance metric {best_performance} is lower than the lower bound {args.performance_lower_bound}'''

   accelerator.wait_for_everyone()
   if accelerator.is_main_process:
      with open(os.path.join(args.output_dir						,						"""all_results.json"""					)						,						"""w"""					) as f:
         json.dump(A__						,						A__					)


def    __lowerCamelCase   (      )     ->    str:
   lowerCamelCase_						:   Dict		= argparse.ArgumentParser(description="""Simple example of training script tracking peak GPU memory usage."""					)
   parser.add_argument(
       """--model_name_or_path"""						,						type=A__						,						default="""bert-base-cased"""						,						help="""Path to pretrained model or model identifier from huggingface.co/models."""						,						required=A__						,						)
   parser.add_argument(
       """--output_dir"""						,						type=A__						,						default="""."""						,						help="""Optional save directory where all checkpoint folders will be stored. Default is the current working directory."""						,						)
   parser.add_argument(
       """--performance_lower_bound"""						,						type=A__						,						default=A__						,						help="""Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value."""						,						)
   parser.add_argument(
       """--num_epochs"""						,						type=A__						,						default=3						,						help="""Number of train epochs."""						,						)
   lowerCamelCase_						:   str		= parser.parse_args()
   lowerCamelCase_						:   Optional[Any]		= {"""lr""": 2e-5, """num_epochs""": args.num_epochs, """seed""": 42, """batch_size""": 16}
   training_function(A__						,						A__					)


if __name__ == "__main__":
      main()









import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


snake_case__				:	Dict					       =					logging.getLogger(__name__)


def    __lowerCamelCase   (      )     ->    Any:
   lowerCamelCase_						:   str		= argparse.ArgumentParser(
       description="""Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset."""					)
   parser.add_argument(
       """--dataset_name"""						,						type=A__						,						default="""wikitext"""						,						help="""Name of the training. Explore datasets at: hf.co/datasets."""						,						)
   parser.add_argument(
       """--dataset_config"""						,						type=A__						,						default="""wikitext-103-raw-v1"""						,						help="""Configuration name of the dataset."""					)
   parser.add_argument(
       """--tokenizer_name_or_path"""						,						type=A__						,						default="""sayakpaul/unigram-tokenizer-wikitext"""						,						help="""Tokenizer identifier. Can be a local filepath or a Hub identifier."""						,						)
   parser.add_argument(
       """--shard_size"""						,						type=A__						,						default=1000						,						help="""Number of entries to go in a single shard."""						,						)
   parser.add_argument("""--split"""						,						type=A__						,						default="""train"""						,						choices=["""train""", """test""", """validation"""]					)
   parser.add_argument(
       """--limit"""						,						default=A__						,						type=A__						,						help="""Limit the number of shards (used for debugging)."""						,						)
   parser.add_argument(
       """--max_length"""						,						type=A__						,						default=512						,						help="""Maximum sequence length. For training on TPUs, it helps to have a maximum"""
       """ sequence length that is a multiple of 8."""						,						)
   parser.add_argument(
       """--output_dir"""						,						default="""tf-tpu"""						,						type=A__						,						help="""Output directory where the TFRecord shards will be saved. If the"""
       """ path is appended with `gs://` ('gs://tf-tpu', for example) then the TFRecord"""
       """ shards will be directly saved to a Google Cloud Storage bucket."""						,						)

   lowerCamelCase_						:   Optional[int]		= parser.parse_args()
   return args


def    __lowerCamelCase   (      A__						: Tuple					)     ->    str:
   def fn(A__						: Dict					):
      return tokenizer(examples["""text"""]					)

   return fn


def    __lowerCamelCase   (      A__						: Union[str, Any]					)     ->    Union[str, Any]:
   lowerCamelCase_						:   Dict		= []
   for i in range(len(tokenized_data["""input_ids"""]					)					):
      lowerCamelCase_						:   List[str]		= {
          """input_ids""": tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data["""input_ids"""][i]					)					),
          """attention_mask""": tf.train.Feature(
              intaa_list=tf.train.IntaaList(value=tokenized_data["""attention_mask"""][i]					)					),
      }
      lowerCamelCase_						:   Optional[int]		= tf.train.Features(feature=A__					)
      lowerCamelCase_						:   Tuple		= tf.train.Example(features=A__					)
      lowerCamelCase_						:   Optional[Any]		= example.SerializeToString()
      records.append(A__					)
   return records


def    __lowerCamelCase   (      A__						: Tuple					)     ->    Union[str, Any]:
   lowerCamelCase_						:   int		= datasets.load_dataset(args.dataset_name						,						args.dataset_config						,						split=args.split					)

   if args.limit is not None:
      lowerCamelCase_						:   str		= min(len(A__					)						,						args.limit					)
      lowerCamelCase_						:   Dict		= dataset.select(range(A__					)					)
      print(f'''Limiting the dataset to {args.limit} entries.'''					)

   lowerCamelCase_						:   str		= AutoTokenizer.from_pretrained(args.tokenizer_name_or_path					)

   # Handle output directory creation.
   # For serializing into a Google Cloud Storage Bucket, one needs to first
   # create a bucket.
   if "gs" not in args.output_dir:
      if not os.path.exists(args.output_dir					):
         os.makedirs(args.output_dir					)
      lowerCamelCase_						:   Optional[int]		= os.path.join(args.output_dir						,						args.split					)
      if not os.path.exists(A__					):
         os.makedirs(A__					)
   else:
      lowerCamelCase_						:   Any		= os.path.join(args.output_dir						,						args.split					)

   # Tokenize the whole dataset at once.
   lowerCamelCase_						:   List[str]		= tokenize_function(A__					)
   lowerCamelCase_						:   List[Any]		= dataset.map(A__						,						batched=A__						,						num_proc=4						,						remove_columns=["""text"""]					)

   # We need to concatenate all our texts together, and then split the result
   # into chunks of a fixed size, which we will call block_size. To do this, we
   # will use the map method again, with the option batched=True. When we use batched=True,
   # the function we pass to map() will be passed multiple inputs at once, allowing us
   # to group them into more or fewer examples than we had in the input.
   # This allows us to create our new fixed-length samples. The advantage of this
   # method is that we don't lose a whole lot of content from the dataset compared to the
   # case where we simply tokenize with a pre-defined max_length.

   def group_texts(A__						: Union[str, Any]					):
      # Concatenate all texts.
      lowerCamelCase_						:   Optional[int]		= {k: sum(examples[k]						,						[]					) for k in examples.keys()}
      lowerCamelCase_						:   Union[str, Any]		= len(concatenated_examples[list(examples.keys()					)[0]]					)
      # We drop the small remainder, though you could add padding instead if the model supports it
      # In this, as in all things, we advise you to follow your heart 🫀
      lowerCamelCase_						:   Union[str, Any]		= (total_length // args.max_length) * args.max_length
      # Split by chunks of max_len.
      lowerCamelCase_						:   List[str]		= {
          k: [t[i : i + args.max_length] for i in range(0						,						A__						,						args.max_length					)]
          for k, t in concatenated_examples.items()
      }
      return result

   lowerCamelCase_						:   Optional[int]		= dataset_tokenized.map(A__						,						batched=A__						,						batch_size=1000						,						num_proc=4					)

   lowerCamelCase_						:   List[str]		= 0
   lowerCamelCase_						:   Tuple		= 0
   for shard in range(0						,						len(A__					)						,						args.shard_size					):
      lowerCamelCase_						:   Union[str, Any]		= grouped_dataset[shard : shard + args.shard_size]
      lowerCamelCase_						:   Optional[int]		= len(dataset_snapshot["""input_ids"""]					)
      lowerCamelCase_						:   List[str]		= os.path.join(A__						,						f'''dataset-{shard_count}-{records_containing}.tfrecord'''					)
      lowerCamelCase_						:   Optional[int]		= get_serialized_examples(A__					)

      with tf.io.TFRecordWriter(A__					) as out_file:
         for i in range(len(A__					)					):
            lowerCamelCase_						:   Dict		= serialized_examples[i]
            out_file.write(A__					)
         print("""Wrote file {} containing {} records""".format(A__						,						A__					)					)

      shard_count += 1
      total_records += records_containing

   with open(f'''split-{args.split}-records-count.txt'''						,						"""w"""					) as f:
      print(f'''Total {args.split} records: {total_records}'''						,						file=A__					)


if __name__ == "__main__":
      snake_case__				:	str					       =					parse_args()
      main(args)





