





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


UpperCamelCase__           =			16
UpperCamelCase__           =			32





def      _UpperCamelCase  (a__ :List[str]							):






							"""simple docstring"""





							return int(x / 2**20							)

class 							__SCREAMING_SNAKE_CASE       :






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







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






def      _UpperCamelCase  (a__ :Accelerator ,      a__ :int = 16 ,      a__ :str = "bert-base-cased" ,      a__ :int = 320 ,      a__ :int = 160 ,      ):






							"""simple docstring"""





							UpperCamelCase__					  =      AutoTokenizer.from_pretrained(a__							)
							UpperCamelCase__					  =      load_dataset(
							    """glue""" ,      """mrpc""" ,      split={"""train""": f"""train[:{n_train}]""", """validation""": f"""validation[:{n_val}]"""}							)

							def tokenize_function(a__ :List[str]							):
														# max_length=None => use the model max length (it's actually the default)
														UpperCamelCase__					  =      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
							UpperCamelCase__					  =      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
							UpperCamelCase__					  =      tokenized_datasets.rename_column("""label""" ,      """labels"""							)

							def collate_fn(a__ :Optional[Any]							):
														# 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.
							UpperCamelCase__					  =      DataLoader(
							    tokenized_datasets["""train"""] ,      shuffle=a__ ,      collate_fn=a__ ,      batch_size=a__							)
							UpperCamelCase__					  =      DataLoader(
							    tokenized_datasets["""validation"""] ,      shuffle=a__ ,      collate_fn=a__ ,      batch_size=a__							)

							return train_dataloader, eval_dataloader





def      _UpperCamelCase  (a__ :List[Any] ,      a__ :int							):






							"""simple docstring"""





							UpperCamelCase__					  =      Accelerator()

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

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

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

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

							if accelerator.state.deepspeed_plugin is not None:
														UpperCamelCase__					  =      accelerator.state.deepspeed_plugin.deepspeed_config[
														    """gradient_accumulation_steps"""
														]
							else:
														UpperCamelCase__					  =      1
							UpperCamelCase__					  =      (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
							):
														UpperCamelCase__					  =      get_linear_schedule_with_warmup(
														    optimizer=a__ ,      num_warmup_steps=0 ,      num_training_steps=a__ ,      )
							else:
														UpperCamelCase__					  =      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.
							UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__					  =      accelerator.prepare(
							    a__ ,      a__ ,      a__ ,      a__ ,      a__							)

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

							# Now we train the model
							UpperCamelCase__					  =      {}
							for epoch in range(a__ ,      a__							):
														with TorchTracemalloc() as tracemalloc:
																					model.train()
																					for step, batch in enumerate(a__							):
																												UpperCamelCase__					  =      model(**a__							)
																												UpperCamelCase__					  =      outputs.loss
																												UpperCamelCase__					  =      loss / gradient_accumulation_steps
																												accelerator.backward(a__							)
																												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							)							)							)
														UpperCamelCase__					  =      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(a__ ,      a__							)





def      _UpperCamelCase  ():






							"""simple docstring"""





							UpperCamelCase__					  =      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(
							    """--peak_memory_upper_bound""" ,      type=a__ ,      default=a__ ,      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=a__ ,      default=320 ,      help="""Number of training examples to use.""" ,      )
							parser.add_argument(
							    """--n_val""" ,      type=a__ ,      default=160 ,      help="""Number of validation examples to use.""" ,      )
							parser.add_argument(
							    """--num_epochs""" ,      type=a__ ,      default=1 ,      help="""Number of train epochs.""" ,      )
							UpperCamelCase__					  =      parser.parse_args()
							UpperCamelCase__					  =      {"""lr""": 2e-5, """num_epochs""": args.num_epochs, """seed""": 42, """batch_size""": 16}
							training_function(a__ ,      a__							)


if __name__ == "__main__":
					main()











import json
import os
import re
import unittest

from transformers import CodeGenTokenizer, CodeGenTokenizerFast
from transformers.models.codegen.tokenization_codegen import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin

@require_tokenizers
class 							__SCREAMING_SNAKE_CASE       (  _a ,		unittest.TestCase							):
							snake_case  :   Dict      =							CodeGenTokenizer
							snake_case  :   Dict      =							CodeGenTokenizerFast
							snake_case  :   Tuple      =							True
							snake_case  :   Optional[int]      =							{"""add_prefix_space""": True}
							snake_case  :   int      =							False






							def 					_lowerCamelCase		(    self	):
														super().setUp()

														# Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
														UpperCamelCase__					  =      [
														    """l""",
														    """o""",
														    """w""",
														    """e""",
														    """r""",
														    """s""",
														    """t""",
														    """i""",
														    """d""",
														    """n""",
														    """\u0120""",
														    """\u0120l""",
														    """\u0120n""",
														    """\u0120lo""",
														    """\u0120low""",
														    """er""",
														    """\u0120lowest""",
														    """\u0120newer""",
														    """\u0120wider""",
														    """<unk>""",
														    """<|endoftext|>""",
														]
														UpperCamelCase__					  =      dict(zip(__lowerCAmelCase	,			range(len(__lowerCAmelCase	)	)	)	)
														UpperCamelCase__					  =      ["""#version: 0.2""", """\u0120 l""", """\u0120l o""", """\u0120lo w""", """e r""", """"""]
														UpperCamelCase__					  =      {"""unk_token""": """<unk>"""}

														UpperCamelCase__					  =      os.path.join(self.tmpdirname	,			VOCAB_FILES_NAMES["""vocab_file"""]	)
														UpperCamelCase__					  =      os.path.join(self.tmpdirname	,			VOCAB_FILES_NAMES["""merges_file"""]	)
														with open(self.vocab_file	,			"""w"""	,			encoding="""utf-8"""	) as fp:
																					fp.write(json.dumps(__lowerCAmelCase	) + """\n"""	)
														with open(self.merges_file	,			"""w"""	,			encoding="""utf-8"""	) as fp:
																					fp.write("""\n""".join(__lowerCAmelCase	)	)






							def 					_lowerCamelCase		(    self	,			**__lowerCAmelCase	):
														kwargs.update(self.special_tokens_map	)
														return CodeGenTokenizer.from_pretrained(self.tmpdirname	,			**__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			**__lowerCAmelCase	):
														kwargs.update(self.special_tokens_map	)
														return CodeGenTokenizerFast.from_pretrained(self.tmpdirname	,			**__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      """lower newer"""
														UpperCamelCase__					  =      """lower newer"""
														return input_text, output_text






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      CodeGenTokenizer(self.vocab_file	,			self.merges_file	,			**self.special_tokens_map	)
														UpperCamelCase__					  =      """lower newer"""
														UpperCamelCase__					  =      ["""\u0120low""", """er""", """\u0120""", """n""", """e""", """w""", """er"""]
														UpperCamelCase__					  =      tokenizer.tokenize(__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														UpperCamelCase__					  =      tokens + [tokenizer.unk_token]
														UpperCamelCase__					  =      [14, 15, 10, 9, 3, 2, 15, 19]
														self.assertListEqual(tokenizer.convert_tokens_to_ids(__lowerCAmelCase	)	,			__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	):
														if not self.test_rust_tokenizer:
																					return

														UpperCamelCase__					  =      self.get_tokenizer()
														UpperCamelCase__					  =      self.get_rust_tokenizer(add_prefix_space=__lowerCAmelCase	)

														UpperCamelCase__					  =      """lower newer"""

														# Testing tokenization
														UpperCamelCase__					  =      tokenizer.tokenize(__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      rust_tokenizer.tokenize(__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														# Testing conversion to ids without special tokens
														UpperCamelCase__					  =      tokenizer.encode(__lowerCAmelCase	,			add_special_tokens=__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      rust_tokenizer.encode(__lowerCAmelCase	,			add_special_tokens=__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														# Testing conversion to ids with special tokens
														UpperCamelCase__					  =      self.get_rust_tokenizer(add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      tokenizer.encode(__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      rust_tokenizer.encode(__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														# Testing the unknown token
														UpperCamelCase__					  =      tokens + [rust_tokenizer.unk_token]
														UpperCamelCase__					  =      [14, 15, 10, 9, 3, 2, 15, 19]
														self.assertListEqual(rust_tokenizer.convert_tokens_to_ids(__lowerCAmelCase	)	,			__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			*__lowerCAmelCase	,			**__lowerCAmelCase	):
														# It's very difficult to mix/test pretokenization with byte-level
														# And get both CodeGen and Roberta to work at the same time (mostly an issue of adding a space before the string)
														pass






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase=15	):
														for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
																					with self.subTest(f"""{tokenizer.__class__.__name__} ({pretrained_name})"""	):
																												UpperCamelCase__					  =      self.rust_tokenizer_class.from_pretrained(__lowerCAmelCase	,			**__lowerCAmelCase	)

																												# Simple input
																												UpperCamelCase__					  =      """This is a simple input"""
																												UpperCamelCase__					  =      ["""This is a simple input 1""", """This is a simple input 2"""]
																												UpperCamelCase__					  =      ("""This is a simple input""", """This is a pair""")
																												UpperCamelCase__					  =      [
																												    ("""This is a simple input 1""", """This is a simple input 2"""),
																												    ("""This is a simple pair 1""", """This is a simple pair 2"""),
																												]

																												# Simple input tests
																												self.assertRaises(__lowerCAmelCase	,			tokenizer_r.encode	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	)

																												# Simple input
																												self.assertRaises(__lowerCAmelCase	,			tokenizer_r.encode_plus	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	)

																												# Simple input
																												self.assertRaises(
																												    __lowerCAmelCase	,			tokenizer_r.batch_encode_plus	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	,			)

																												# Pair input
																												self.assertRaises(__lowerCAmelCase	,			tokenizer_r.encode	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	)

																												# Pair input
																												self.assertRaises(__lowerCAmelCase	,			tokenizer_r.encode_plus	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	)

																												# Pair input
																												self.assertRaises(
																												    __lowerCAmelCase	,			tokenizer_r.batch_encode_plus	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	,			)






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      CodeGenTokenizer.from_pretrained(self.tmpdirname	,			pad_token="""<pad>"""	)

														# Simple input
														UpperCamelCase__					  =      """This is a simple input"""
														UpperCamelCase__					  =      ["""This is a simple input looooooooong""", """This is a simple input"""]
														UpperCamelCase__					  =      ("""This is a simple input""", """This is a pair""")
														UpperCamelCase__					  =      [
														    ("""This is a simple input loooooong""", """This is a simple input"""),
														    ("""This is a simple pair loooooong""", """This is a simple pair"""),
														]

														UpperCamelCase__					  =      tokenizer.pad_token_id

														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	,			padding="""max_length"""	,			max_length=30	,			return_tensors="""np"""	)
														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	,			padding=__lowerCAmelCase	,			truncate=__lowerCAmelCase	,			return_tensors="""np"""	)
														UpperCamelCase__					  =      tokenizer(*__lowerCAmelCase	,			padding="""max_length"""	,			max_length=60	,			return_tensors="""np"""	)
														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	,			padding=__lowerCAmelCase	,			truncate=__lowerCAmelCase	,			return_tensors="""np"""	)

														# s
														# test single string max_length padding
														self.assertEqual(out_s["""input_ids"""].shape[-1]	,			30	)
														self.assertTrue(pad_token_id in out_s["""input_ids"""]	)
														self.assertTrue(0 in out_s["""attention_mask"""]	)

														# s2
														# test automatic padding
														self.assertEqual(out_sa["""input_ids"""].shape[-1]	,			33	)
														# long slice doesn't have padding
														self.assertFalse(pad_token_id in out_sa["""input_ids"""][0]	)
														self.assertFalse(0 in out_sa["""attention_mask"""][0]	)
														# short slice does have padding
														self.assertTrue(pad_token_id in out_sa["""input_ids"""][1]	)
														self.assertTrue(0 in out_sa["""attention_mask"""][1]	)

														# p
														# test single pair max_length padding
														self.assertEqual(out_p["""input_ids"""].shape[-1]	,			60	)
														self.assertTrue(pad_token_id in out_p["""input_ids"""]	)
														self.assertTrue(0 in out_p["""attention_mask"""]	)

														# p2
														# test automatic padding pair
														self.assertEqual(out_pa["""input_ids"""].shape[-1]	,			52	)
														# long slice pair doesn't have padding
														self.assertFalse(pad_token_id in out_pa["""input_ids"""][0]	)
														self.assertFalse(0 in out_pa["""attention_mask"""][0]	)
														# short slice pair does have padding
														self.assertTrue(pad_token_id in out_pa["""input_ids"""][1]	)
														self.assertTrue(0 in out_pa["""attention_mask"""][1]	)






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      """$$$"""
														UpperCamelCase__					  =      CodeGenTokenizer.from_pretrained(self.tmpdirname	,			bos_token=__lowerCAmelCase	,			add_bos_token=__lowerCAmelCase	)

														UpperCamelCase__					  =      """This is a simple input"""
														UpperCamelCase__					  =      ["""This is a simple input 1""", """This is a simple input 2"""]

														UpperCamelCase__					  =      tokenizer.bos_token_id

														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	)
														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	)

														self.assertEqual(out_s.input_ids[0]	,			__lowerCAmelCase	)
														self.assertTrue(all(o[0] == bos_token_id for o in out_sa.input_ids	)	)

														UpperCamelCase__					  =      tokenizer.decode(out_s.input_ids	)
														UpperCamelCase__					  =      tokenizer.batch_decode(out_sa.input_ids	)

														self.assertEqual(decode_s.split()[0]	,			__lowerCAmelCase	)
														self.assertTrue(all(d.split()[0] == bos_token for d in decode_sa	)	)






							@slow
							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      CodeGenTokenizer.from_pretrained("""Salesforce/codegen-350M-mono"""	)

														UpperCamelCase__					  =      """\nif len_a > len_b:\n    result = a\nelse:\n    result = b\n\n\n\n#"""
														UpperCamelCase__					  =      """\nif len_a > len_b:      result = a\nelse:      result = b"""

														UpperCamelCase__					  =      tokenizer.encode(__lowerCAmelCase	)
														UpperCamelCase__					  =      ["""^#""", re.escape("""<|endoftext|>"""	), """^'''""", """^\"\"\"""", """\n\n\n"""]
														UpperCamelCase__					  =      tokenizer.decode(__lowerCAmelCase	,			truncate_before_pattern=__lowerCAmelCase	)
														self.assertEqual(__lowerCAmelCase	,			__lowerCAmelCase	)







							def 					_lowerCamelCase		(    self	):
														pass




