





'''simple docstring'''





import argparse
import gc
import json
import os

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


_lowerCamelCase   				=			16
_lowerCamelCase   				=			32






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							str    ):
		return int(x / 2**20    )




class   UpperCamelCase_       :


							def __enter__(							self    :int		)   ->   Tuple:



									"""simple docstring"""



									gc.collect()
									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
									SCREAMING_SNAKE_CASE__			=      torch.cuda.memory_allocated()
									return self


							def __exit__(							self    :int							,				*__A    :Any		)   ->   List[str]:



									"""simple docstring"""



									gc.collect()
									torch.cuda.empty_cache()
									SCREAMING_SNAKE_CASE__			=      torch.cuda.memory_allocated()
									SCREAMING_SNAKE_CASE__			=      torch.cuda.max_memory_allocated()
									SCREAMING_SNAKE_CASE__			=      bamb(self.end - self.begin		)
									SCREAMING_SNAKE_CASE__			=      bamb(self.peak - self.begin		)
									# print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							Accelerator					,   UpperCamelCase__:							int = 16					,   UpperCamelCase__:							str = "bert-base-cased"					,   UpperCamelCase__:							int = 320					,   UpperCamelCase__:							int = 160					,   ):
		SCREAMING_SNAKE_CASE__			=      AutoTokenizer.from_pretrained(UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      load_dataset(
		    """glue"""					,   """mrpc"""					,   split={"""train""": f'''train[:{n_train}]''', """validation""": f'''validation[:{n_val}]'''}    )

		def tokenize_function(UpperCamelCase__:							Optional[Any]    ):
				# max_length=None => use the model max length (it's actually the default)
				SCREAMING_SNAKE_CASE__			=      tokenizer(examples["""sentence1"""]					,   examples["""sentence2"""]					,   truncation=UpperCamelCase__					,   max_length=UpperCamelCase__    )
				return outputs

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

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

		def collate_fn(UpperCamelCase__:							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(UpperCamelCase__					,   padding="""max_length"""					,   max_length=128					,   return_tensors="""pt"""    )
				return tokenizer.pad(UpperCamelCase__					,   padding="""longest"""					,   return_tensors="""pt"""    )

		# Instantiate dataloaders.
		SCREAMING_SNAKE_CASE__			=      DataLoader(
		    tokenized_datasets["""train"""]					,   shuffle=UpperCamelCase__					,   collate_fn=UpperCamelCase__					,   batch_size=UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      DataLoader(
		    tokenized_datasets["""validation"""]					,   shuffle=UpperCamelCase__					,   collate_fn=UpperCamelCase__					,   batch_size=UpperCamelCase__    )

		return train_dataloader, eval_dataloader






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							Tuple					,   UpperCamelCase__:							Tuple    ):
		# Initialize accelerator
		SCREAMING_SNAKE_CASE__			=      Accelerator()

		# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
		SCREAMING_SNAKE_CASE__			=      config["""lr"""]
		SCREAMING_SNAKE_CASE__			=      int(config["""num_epochs"""]    )
		SCREAMING_SNAKE_CASE__			=      int(config["""seed"""]    )
		SCREAMING_SNAKE_CASE__			=      int(config["""batch_size"""]    )
		SCREAMING_SNAKE_CASE__			=      args.model_name_or_path

		set_seed(UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      get_dataloaders(UpperCamelCase__					,   UpperCamelCase__					,   UpperCamelCase__					,   args.n_train					,   args.n_val    )

		# Instantiate the model (we build the model here so that the seed also control new weights initialization)
		SCREAMING_SNAKE_CASE__			=      AutoModelForSequenceClassification.from_pretrained(UpperCamelCase__					,   return_dict=UpperCamelCase__    )

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

		if accelerator.state.deepspeed_plugin is not None:
				SCREAMING_SNAKE_CASE__			=      accelerator.state.deepspeed_plugin.deepspeed_config[
				    """gradient_accumulation_steps"""
				]
		else:
				SCREAMING_SNAKE_CASE__			=      1
		SCREAMING_SNAKE_CASE__			=      (len(UpperCamelCase__    ) * num_epochs) // gradient_accumulation_steps

		# Instantiate scheduler
		if (
		    accelerator.state.deepspeed_plugin is None
		    or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
		):
				SCREAMING_SNAKE_CASE__			=      get_linear_schedule_with_warmup(
				    optimizer=UpperCamelCase__					,   num_warmup_steps=0					,   num_training_steps=UpperCamelCase__					,   )
		else:
				SCREAMING_SNAKE_CASE__			=      DummyScheduler(UpperCamelCase__					,   total_num_steps=UpperCamelCase__					,   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.
		SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      accelerator.prepare(
		    UpperCamelCase__					,   UpperCamelCase__					,   UpperCamelCase__					,   UpperCamelCase__					,   UpperCamelCase__    )

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

		# Now we train the model
		SCREAMING_SNAKE_CASE__			=      {}
		for epoch in range(UpperCamelCase__					,   UpperCamelCase__    ):
				with TorchTracemalloc() as tracemalloc:
						model.train()
						for step, batch in enumerate(UpperCamelCase__    ):
								SCREAMING_SNAKE_CASE__			=      model(**UpperCamelCase__    )
								SCREAMING_SNAKE_CASE__			=      outputs.loss
								SCREAMING_SNAKE_CASE__			=      loss / gradient_accumulation_steps
								accelerator.backward(UpperCamelCase__    )
								if step % gradient_accumulation_steps == 0:
										optimizer.step()
										lr_scheduler.step()
										optimizer.zero_grad()

								overall_step += 1

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
				accelerator.print("""Memory before entering the train : {}""".format(bamb(tracemalloc.begin    )    )    )
				accelerator.print("""Memory consumed at the end of the train (end-begin): {}""".format(tracemalloc.used    )    )
				accelerator.print("""Peak Memory consumed during the train (max-begin): {}""".format(tracemalloc.peaked    )    )
				accelerator.print(
				    """Total Peak Memory consumed during the train (max): {}""".format(
				        tracemalloc.peaked + bamb(tracemalloc.begin    )    )    )
				SCREAMING_SNAKE_CASE__			=      tracemalloc.peaked + bamb(tracemalloc.begin    )
				if args.peak_memory_upper_bound is not None:
						assert (
						    train_total_peak_memory[f'''epoch-{epoch}'''] <= args.peak_memory_upper_bound
						), "Peak memory usage exceeded the upper bound"

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






def 				SCREAMING_SNAKE_CASE__							(       ):
		SCREAMING_SNAKE_CASE__			=      argparse.ArgumentParser(description="""Simple example of training script tracking peak GPU memory usage."""    )
		parser.add_argument(
		    """--model_name_or_path"""					,   type=UpperCamelCase__					,   default="""bert-base-cased"""					,   help="""Path to pretrained model or model identifier from huggingface.co/models."""					,   required=UpperCamelCase__					,   )
		parser.add_argument(
		    """--output_dir"""					,   type=UpperCamelCase__					,   default="""."""					,   help="""Optional save directory where all checkpoint folders will be stored. Default is the current working directory."""					,   )
		parser.add_argument(
		    """--peak_memory_upper_bound"""					,   type=UpperCamelCase__					,   default=UpperCamelCase__					,   help="""The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value."""					,   )
		parser.add_argument(
		    """--n_train"""					,   type=UpperCamelCase__					,   default=320					,   help="""Number of training examples to use."""					,   )
		parser.add_argument(
		    """--n_val"""					,   type=UpperCamelCase__					,   default=160					,   help="""Number of validation examples to use."""					,   )
		parser.add_argument(
		    """--num_epochs"""					,   type=UpperCamelCase__					,   default=1					,   help="""Number of train epochs."""					,   )
		SCREAMING_SNAKE_CASE__			=      parser.parse_args()
		SCREAMING_SNAKE_CASE__			=      {"""lr""": 2e-5, """num_epochs""": args.num_epochs, """seed""": 42, """batch_size""": 16}
		training_function(UpperCamelCase__					,   UpperCamelCase__    )


if __name__ == "__main__":
	main()






import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


_lowerCamelCase   				=			logging.getLogger(__name__)






def 				SCREAMING_SNAKE_CASE__							(       ):
		SCREAMING_SNAKE_CASE__			=      argparse.ArgumentParser(
		    description="""Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset."""    )
		parser.add_argument(
		    """--dataset_name"""					,   type=UpperCamelCase__					,   default="""wikitext"""					,   help="""Name of the training. Explore datasets at: hf.co/datasets."""					,   )
		parser.add_argument(
		    """--dataset_config"""					,   type=UpperCamelCase__					,   default="""wikitext-103-raw-v1"""					,   help="""Configuration name of the dataset."""    )
		parser.add_argument(
		    """--tokenizer_name_or_path"""					,   type=UpperCamelCase__					,   default="""sayakpaul/unigram-tokenizer-wikitext"""					,   help="""Tokenizer identifier. Can be a local filepath or a Hub identifier."""					,   )
		parser.add_argument(
		    """--shard_size"""					,   type=UpperCamelCase__					,   default=1_000					,   help="""Number of entries to go in a single shard."""					,   )
		parser.add_argument("""--split"""					,   type=UpperCamelCase__					,   default="""train"""					,   choices=["""train""", """test""", """validation"""]    )
		parser.add_argument(
		    """--limit"""					,   default=UpperCamelCase__					,   type=UpperCamelCase__					,   help="""Limit the number of shards (used for debugging)."""					,   )
		parser.add_argument(
		    """--max_length"""					,   type=UpperCamelCase__					,   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=UpperCamelCase__					,   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."""					,   )

		SCREAMING_SNAKE_CASE__			=      parser.parse_args()
		return args






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							List[Any]    ):
		def fn(UpperCamelCase__:							Any    ):
				return tokenizer(examples["""text"""]    )

		return fn






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							Any    ):
		SCREAMING_SNAKE_CASE__			=      []
		for i in range(len(tokenized_data["""input_ids"""]    )    ):
				SCREAMING_SNAKE_CASE__			=      {
				    """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]    )    ),
				}
				SCREAMING_SNAKE_CASE__			=      tf.train.Features(feature=UpperCamelCase__    )
				SCREAMING_SNAKE_CASE__			=      tf.train.Example(features=UpperCamelCase__    )
				SCREAMING_SNAKE_CASE__			=      example.SerializeToString()
				records.append(UpperCamelCase__    )
		return records






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							List[str]    ):
		SCREAMING_SNAKE_CASE__			=      datasets.load_dataset(args.dataset_name					,   args.dataset_config					,   split=args.split    )

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

		SCREAMING_SNAKE_CASE__			=      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    )
				SCREAMING_SNAKE_CASE__			=      os.path.join(args.output_dir					,   args.split    )
				if not os.path.exists(UpperCamelCase__    ):
						os.makedirs(UpperCamelCase__    )
		else:
				SCREAMING_SNAKE_CASE__			=      os.path.join(args.output_dir					,   args.split    )

		# Tokenize the whole dataset at once.
		SCREAMING_SNAKE_CASE__			=      tokenize_function(UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      dataset.map(UpperCamelCase__					,   batched=UpperCamelCase__					,   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(UpperCamelCase__:							int    ):
				# Concatenate all texts.
				SCREAMING_SNAKE_CASE__			=      {k: sum(examples[k]					,   []    ) for k in examples.keys()}
				SCREAMING_SNAKE_CASE__			=      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 🫀
				SCREAMING_SNAKE_CASE__			=      (total_length // args.max_length) * args.max_length
				# Split by chunks of max_len.
				SCREAMING_SNAKE_CASE__			=      {
				    k: [t[i : i + args.max_length] for i in range(0					,   UpperCamelCase__					,   args.max_length    )]
				    for k, t in concatenated_examples.items()
				}
				return result

		SCREAMING_SNAKE_CASE__			=      dataset_tokenized.map(UpperCamelCase__					,   batched=UpperCamelCase__					,   batch_size=1_000					,   num_proc=4    )

		SCREAMING_SNAKE_CASE__			=      0
		SCREAMING_SNAKE_CASE__			=      0
		for shard in range(0					,   len(UpperCamelCase__    )					,   args.shard_size    ):
				SCREAMING_SNAKE_CASE__			=      grouped_dataset[shard : shard + args.shard_size]
				SCREAMING_SNAKE_CASE__			=      len(dataset_snapshot["""input_ids"""]    )
				SCREAMING_SNAKE_CASE__			=      os.path.join(UpperCamelCase__					,   f'''dataset-{shard_count}-{records_containing}.tfrecord'''    )
				SCREAMING_SNAKE_CASE__			=      get_serialized_examples(UpperCamelCase__    )

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

				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=UpperCamelCase__    )


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