
'''simple docstring'''







import json
import os
from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple, Union

import numpy as np
import regex as re

from ...utils import is_tf_available, is_torch_available, logging


if TYPE_CHECKING:
							if is_torch_available():
														import torch
							if is_tf_available():
														import tensorflow as tf

from ...tokenization_utils import AddedToken, PreTrainedTokenizer


_lowercase				:						Any												= logging.get_logger(__name__)

_lowercase				:						int												= {
    "vocab_file": "vocab.json",
    "merges_file": "merges.txt",
}

_lowercase				:						List[Any]												= {
    "vocab_file": {
        "Salesforce/codegen-350M-mono": "https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/vocab.json",
    },
    "merges_file": {
        "Salesforce/codegen-350M-mono": "https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/merges.txt",
    },
}

_lowercase				:						Dict												= {
    "Salesforce/codegen-350M-mono": 2_0_4_8,
}
@lru_cache()
def       snake_case_	(       ):



					"""simple docstring"""





					lowercase_			:  str									= (
					    list(range(ord('''!'''							)   ,   ord('''~'''							) + 1							)							) + list(range(ord('''¡'''							)   ,   ord('''¬'''							) + 1							)							) + list(range(ord('''®'''							)   ,   ord('''ÿ'''							) + 1							)							)
					)
					lowercase_			:  Optional[int]									= bs[:]
					lowercase_			:  str									= 0
					for b in range(2**8							):
										if b not in bs:
															bs.append(__a							)
															cs.append(2**8 + n							)
															n += 1
					lowercase_			:  List[str]									= [chr(__a							) for n in cs]
					return dict(zip(__a   ,   __a							)							)
def       snake_case_	(       __SCREAMING_SNAKE_CASE      :				Union[str, Any]							):



					"""simple docstring"""





					lowercase_			:  int									= set()
					lowercase_			:  Optional[int]									= word[0]
					for char in word[1:]:
										pairs.add((prev_char, char)							)
										lowercase_			:  Optional[int]									= char
					return pairs






class 				lowerCAmelCase__     (						UpperCamelCase__  ):
					lowerCAmelCase_       =  VOCAB_FILES_NAMES
					lowerCAmelCase_       =  PRETRAINED_VOCAB_FILES_MAP
					lowerCAmelCase_       =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					lowerCAmelCase_       =  ['''input_ids''', '''attention_mask''']






					def __init__(     self       ,							__SCREAMING_SNAKE_CASE       ,							__SCREAMING_SNAKE_CASE       ,							__SCREAMING_SNAKE_CASE="replace"       ,							__SCREAMING_SNAKE_CASE="<|endoftext|>"       ,							__SCREAMING_SNAKE_CASE="<|endoftext|>"       ,							__SCREAMING_SNAKE_CASE="<|endoftext|>"       ,							__SCREAMING_SNAKE_CASE=None       ,							__SCREAMING_SNAKE_CASE=False       ,							__SCREAMING_SNAKE_CASE=False       ,							**__SCREAMING_SNAKE_CASE       ,							):


										"""simple docstring"""





										lowercase_			:  Dict									= AddedToken(snake_case_       ,							lstrip=snake_case_       ,							rstrip=snake_case_	) if isinstance(snake_case_       ,							snake_case_	) else bos_token
										lowercase_			:  Optional[int]									= AddedToken(snake_case_       ,							lstrip=snake_case_       ,							rstrip=snake_case_	) if isinstance(snake_case_       ,							snake_case_	) else eos_token
										lowercase_			:  List[Any]									= AddedToken(snake_case_       ,							lstrip=snake_case_       ,							rstrip=snake_case_	) if isinstance(snake_case_       ,							snake_case_	) else unk_token
										lowercase_			:  Dict									= AddedToken(snake_case_       ,							lstrip=snake_case_       ,							rstrip=snake_case_	) if isinstance(snake_case_       ,							snake_case_	) else pad_token
										super().__init__(
										    errors=snake_case_       ,							unk_token=snake_case_       ,							bos_token=snake_case_       ,							eos_token=snake_case_       ,							pad_token=snake_case_       ,							add_prefix_space=snake_case_       ,							add_bos_token=snake_case_       ,							**snake_case_       ,							)
										lowercase_			:  Optional[int]									= add_bos_token

										with open(snake_case_       ,							encoding='''utf-8'''	) as vocab_handle:
															lowercase_			:  List[str]									= json.load(snake_case_	)
										lowercase_			:  Optional[Any]									= {v: k for k, v in self.encoder.items()}
										lowercase_			:  str									= errors  # how to handle errors in decoding
										lowercase_			:  List[str]									= bytes_to_unicode()
										lowercase_			:  List[str]									= {v: k for k, v in self.byte_encoder.items()}
										with open(snake_case_       ,							encoding='''utf-8'''	) as merges_handle:
															lowercase_			:  Tuple									= merges_handle.read().split('''\n'''	)[1:-1]
										lowercase_			:  List[str]									= [tuple(merge.split()	) for merge in bpe_merges]
										lowercase_			:  Any									= dict(zip(snake_case_       ,							range(len(snake_case_	)	)	)	)
										lowercase_			:  List[Any]									= {}
										lowercase_			:  Optional[Any]									= add_prefix_space

										# Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
										lowercase_			:  str									= re.compile(R'''\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+'''	)






					@property
					def 			_snake_case					(     self	):


										"""simple docstring"""





										return len(self.encoder	)






					def 			_snake_case					(     self	):


										"""simple docstring"""





										return dict(self.encoder       ,							**self.added_tokens_encoder	)






					def 			_snake_case					(     self       ,							__SCREAMING_SNAKE_CASE	):


										"""simple docstring"""





										if token in self.cache:
															return self.cache[token]
										lowercase_			:  List[Any]									= tuple(snake_case_	)
										lowercase_			:  Dict									= get_pairs(snake_case_	)

										if not pairs:
															return token

										while True:
															lowercase_			:  Optional[Any]									= min(snake_case_       ,							key=lambda __SCREAMING_SNAKE_CASE	: self.bpe_ranks.get(snake_case_       ,							float('''inf'''	)	)	)
															if bigram not in self.bpe_ranks:
																				break
															lowercase_			:  List[Any]									= bigram
															lowercase_			:  List[Any]									= []
															lowercase_			:  Tuple									= 0
															while i < len(snake_case_	):
																				try:
																									lowercase_			:  int									= word.index(snake_case_       ,							snake_case_	)
																				except ValueError:
																									new_word.extend(word[i:]	)
																									break
																				else:
																									new_word.extend(word[i:j]	)
																									lowercase_			:  List[Any]									= j

																				if word[i] == first and i < len(snake_case_	) - 1 and word[i + 1] == second:
																									new_word.append(first + second	)
																									i += 2
																				else:
																									new_word.append(word[i]	)
																									i += 1
															lowercase_			:  Tuple									= tuple(snake_case_	)
															lowercase_			:  Dict									= new_word
															if len(snake_case_	) == 1:
																				break
															else:
																				lowercase_			:  Dict									= get_pairs(snake_case_	)
										lowercase_			:  Any									= ''' '''.join(snake_case_	)
										lowercase_			:  List[str]									= word
										return word






					def 			_snake_case					(     self       ,							__SCREAMING_SNAKE_CASE       ,							__SCREAMING_SNAKE_CASE=None	):


										"""simple docstring"""





										if self.add_bos_token:
															lowercase_			:  Dict									= [self.bos_token_id]
										else:
															lowercase_			:  Optional[Any]									= []

										lowercase_			:  Union[str, Any]									= bos_token_ids + token_ids_a

										if token_ids_a is None:
															return output

										return output + bos_token_ids + token_ids_a






					def 			_snake_case					(     self       ,							__SCREAMING_SNAKE_CASE	):


										"""simple docstring"""





										lowercase_			:  List[Any]									= []
										for token in re.findall(self.pat       ,							snake_case_	):
															lowercase_			:  Optional[int]									= ''''''.join(
															    self.byte_encoder[b] for b in token.encode('''utf-8'''	)	)  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
															bpe_tokens.extend(bpe_token for bpe_token in self.bpe(snake_case_	).split(''' '''	)	)
										return bpe_tokens






					def 			_snake_case					(     self       ,							__SCREAMING_SNAKE_CASE	):


										"""simple docstring"""





										return self.encoder.get(snake_case_       ,							self.encoder.get(self.unk_token	)	)






					def 			_snake_case					(     self       ,							__SCREAMING_SNAKE_CASE	):


										"""simple docstring"""





										return self.decoder.get(snake_case_	)






					def 			_snake_case					(     self       ,							__SCREAMING_SNAKE_CASE	):


										"""simple docstring"""





										lowercase_			:  Optional[int]									= ''''''.join(snake_case_	)
										lowercase_			:  Optional[Any]									= bytearray([self.byte_decoder[c] for c in text]	).decode('''utf-8'''       ,							errors=self.errors	)
										return text






					def 			_snake_case					(     self       ,							__SCREAMING_SNAKE_CASE       ,							__SCREAMING_SNAKE_CASE = None	):


										"""simple docstring"""





										if not os.path.isdir(snake_case_	):
															logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''	)
															return
										lowercase_			:  List[str]									= os.path.join(
										    snake_case_       ,							(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']	)
										lowercase_			:  List[str]									= os.path.join(
										    snake_case_       ,							(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file''']	)

										with open(snake_case_       ,							'''w'''       ,							encoding='''utf-8'''	) as f:
															f.write(json.dumps(self.encoder       ,							indent=2       ,							sort_keys=snake_case_       ,							ensure_ascii=snake_case_	) + '''\n'''	)

										lowercase_			:  Tuple									= 0
										with open(snake_case_       ,							'''w'''       ,							encoding='''utf-8'''	) as writer:
															writer.write('''#version: 0.2\n'''	)
															for bpe_tokens, token_index in sorted(self.bpe_ranks.items()       ,							key=lambda __SCREAMING_SNAKE_CASE	: kv[1]	):
																				if index != token_index:
																									logger.warning(
																									    F'''Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'''
																									    ''' Please check that the tokenizer is not corrupted!'''	)
																									lowercase_			:  int									= token_index
																				writer.write(''' '''.join(snake_case_	) + '''\n'''	)
																				index += 1

										return vocab_file, merge_file






					def 			_snake_case					(     self       ,							__SCREAMING_SNAKE_CASE       ,							__SCREAMING_SNAKE_CASE=False       ,							**__SCREAMING_SNAKE_CASE	):


										"""simple docstring"""





										lowercase_			:  str									= kwargs.pop('''add_prefix_space'''       ,							self.add_prefix_space	)
										if is_split_into_words or add_prefix_space:
															lowercase_			:  Optional[Any]									= ''' ''' + text
										return (text, kwargs)






					def 			_snake_case					(     self       ,							__SCREAMING_SNAKE_CASE       ,							__SCREAMING_SNAKE_CASE = False       ,							__SCREAMING_SNAKE_CASE = None       ,							__SCREAMING_SNAKE_CASE = None       ,							**__SCREAMING_SNAKE_CASE       ,							):


										"""simple docstring"""





										lowercase_			:  List[str]									= super()._decode(
										    token_ids=snake_case_       ,							skip_special_tokens=snake_case_       ,							clean_up_tokenization_spaces=snake_case_       ,							**snake_case_       ,							)

										if truncate_before_pattern is not None and len(snake_case_	) > 0:
															lowercase_			:  List[Any]									= self.truncate(snake_case_       ,							snake_case_	)

										return decoded_text






					def 			_snake_case					(     self       ,							__SCREAMING_SNAKE_CASE       ,							__SCREAMING_SNAKE_CASE	):


										"""simple docstring"""





										def find_re(__SCREAMING_SNAKE_CASE       ,							__SCREAMING_SNAKE_CASE       ,							__SCREAMING_SNAKE_CASE	):
															lowercase_			:  int									= pattern.search(snake_case_       ,							snake_case_	)
															return m.start() if m else -1

										lowercase_			:  Any									= [re.compile(snake_case_       ,							re.MULTILINE	) for pattern in truncate_before_pattern]

										lowercase_			:  List[str]									= list(re.finditer('''^print'''       ,							snake_case_       ,							re.MULTILINE	)	)

										if len(snake_case_	) > 1:
															lowercase_			:  Optional[int]									= completion[: prints[1].start()]

										lowercase_			:  Union[str, Any]									= list(re.finditer('''^def'''       ,							snake_case_       ,							re.MULTILINE	)	)

										if len(snake_case_	) > 1:
															lowercase_			:  List[Any]									= completion[: defs[1].start()]

										lowercase_			:  Any									= 0

										lowercase_			:  Any									= [
										    pos for pos in [find_re(snake_case_       ,							snake_case_       ,							snake_case_	) for terminal in terminals] if pos != -1
										]

										if len(snake_case_	) > 0:
															return completion[: min(snake_case_	)]
										else:
															return completion





'''simple docstring'''







import argparse
import os

# New Code #
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
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to combine both the gradient accumulation
# and automatic batch size finder utilities of Accelerate to perfrom
# automatic gradient accumulation
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################

_lowercase				:						Tuple												= 3_2
def       snake_case_	(       __SCREAMING_SNAKE_CASE      :				Accelerator   ,   __SCREAMING_SNAKE_CASE      :				int = 16							):



						"""simple docstring"""





						lowercase_			:  int									= AutoTokenizer.from_pretrained('''bert-base-cased'''							)
						lowercase_			:  int									= load_dataset('''glue'''   ,   '''mrpc'''							)

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

						# Apply the method we just defined to all the examples in all the splits of the dataset
						# starting with the main process first:
						with accelerator.main_process_first():
												lowercase_			:  Optional[Any]									= datasets.map(
												    __SCREAMING_SNAKE_CASE   ,   batched=__SCREAMING_SNAKE_CASE   ,   remove_columns=['''idx''', '''sentence1''', '''sentence2''']   ,   )

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

						def collate_fn(__SCREAMING_SNAKE_CASE      :				List[Any]							):
												# When using mixed precision we want round multiples of 8/16
												if accelerator.mixed_precision == "fp8":
																		lowercase_			:  List[str]									= 16
												elif accelerator.mixed_precision != "no":
																		lowercase_			:  List[Any]									= 8
												else:
																		lowercase_			:  List[Any]									= None

												return tokenizer.pad(
												    __SCREAMING_SNAKE_CASE   ,   padding='''longest'''   ,   pad_to_multiple_of=__SCREAMING_SNAKE_CASE   ,   return_tensors='''pt'''   ,   )

						# Instantiate dataloaders.
						lowercase_			:  List[str]									= DataLoader(
						    tokenized_datasets['''train''']   ,   shuffle=__SCREAMING_SNAKE_CASE   ,   collate_fn=__SCREAMING_SNAKE_CASE   ,   batch_size=__SCREAMING_SNAKE_CASE							)
						lowercase_			:  int									= DataLoader(
						    tokenized_datasets['''validation''']   ,   shuffle=__SCREAMING_SNAKE_CASE   ,   collate_fn=__SCREAMING_SNAKE_CASE   ,   batch_size=__SCREAMING_SNAKE_CASE							)

						return train_dataloader, eval_dataloader


# For testing only
if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
							from accelerate.test_utils.training import mocked_dataloaders

							_lowercase				:						Tuple												= mocked_dataloaders  # noqa: F811
def       snake_case_	(       __SCREAMING_SNAKE_CASE      :				str   ,   __SCREAMING_SNAKE_CASE      :				List[Any]							):



						"""simple docstring"""





						if os.environ.get('''TESTING_MOCKED_DATALOADERS'''   ,   __SCREAMING_SNAKE_CASE							) == "1":
												lowercase_			:  Dict									= 2
						# Initialize accelerator
						lowercase_			:  Any									= Accelerator(cpu=args.cpu   ,   mixed_precision=args.mixed_precision							)
						# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
						lowercase_			:  Optional[Any]									= config['''lr''']
						lowercase_			:  Optional[Any]									= int(config['''num_epochs''']							)
						lowercase_			:  str									= int(config['''seed''']							)
						lowercase_			:  Optional[Any]									= int(config['''batch_size''']							)

						lowercase_			:  int									= evaluate.load('''glue'''   ,   '''mrpc'''							)

						# New Code #
						# We use the `find_executable_batch_size` decorator, passing in the desired observed batch size
						# to train on. If a CUDA OOM error occurs, it will retry this loop cutting the batch size in
						# half each time. From this, we can calculate the number of gradient accumulation steps needed
						# and modify the Accelerator object as a result
						@find_executable_batch_size(starting_batch_size=int(__SCREAMING_SNAKE_CASE							)							)
						def inner_training_loop(__SCREAMING_SNAKE_CASE      :				Any							):
												# Since we need to modify the outside accelerator object, we need to bring it
												# to the local scope
												nonlocal accelerator

												# We can calculate the number of gradient accumulation steps based on the current
												# batch size vs the starting batch size
												lowercase_			:  Union[str, Any]									= observed_batch_size // batch_size

												# And then set it in the Accelerator directly:
												lowercase_			:  Tuple									= num_gradient_accumulation_steps

												# Next we need to free all of the stored model references in the Accelerator each time
												accelerator.free_memory()

												# And set the seed so our results are reproducable each reset
												set_seed(__SCREAMING_SNAKE_CASE							)

												# Instantiate the model (we build the model here so that the seed also control new weights initialization)
												lowercase_			:  Dict									= AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''   ,   return_dict=__SCREAMING_SNAKE_CASE							)

												# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
												# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
												# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
												lowercase_			:  Optional[int]									= model.to(accelerator.device							)

												# Instantiate optimizer
												lowercase_			:  int									= AdamW(params=model.parameters()   ,   lr=__SCREAMING_SNAKE_CASE							)
												lowercase_		,				lowercase_			:  List[str]									= get_dataloaders(__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE							)

												# Instantiate scheduler
												lowercase_			:  str									= get_linear_schedule_with_warmup(
												    optimizer=__SCREAMING_SNAKE_CASE   ,   num_warmup_steps=100   ,   num_training_steps=(len(__SCREAMING_SNAKE_CASE							) * num_epochs)   ,   )

												# 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.
												lowercase_		,				lowercase_		,				lowercase_		,				lowercase_		,				lowercase_			:  Any									= accelerator.prepare(
												    __SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE							)

												# Now we train the model
												for epoch in range(__SCREAMING_SNAKE_CASE							):
																		model.train()
																		for step, batch in enumerate(__SCREAMING_SNAKE_CASE							):
																								# And perform gradient accumulation
																								with accelerator.accumulate(__SCREAMING_SNAKE_CASE							):
																														# We could avoid this line since we set the accelerator with `device_placement=True`.
																														batch.to(accelerator.device							)
																														lowercase_			:  List[str]									= model(**__SCREAMING_SNAKE_CASE							)
																														lowercase_			:  Any									= outputs.loss
																														accelerator.backward(__SCREAMING_SNAKE_CASE							)
																														optimizer.step()
																														lr_scheduler.step()
																														optimizer.zero_grad()

																		model.eval()
																		for step, batch in enumerate(__SCREAMING_SNAKE_CASE							):
																								# We could avoid this line since we set the accelerator with `device_placement=True`.
																								batch.to(accelerator.device							)
																								with torch.no_grad():
																														lowercase_			:  Tuple									= model(**__SCREAMING_SNAKE_CASE							)
																								lowercase_			:  List[str]									= outputs.logits.argmax(dim=-1							)
																								lowercase_		,				lowercase_			:  int									= accelerator.gather_for_metrics((predictions, batch['''labels'''])							)
																								metric.add_batch(
																								    predictions=__SCREAMING_SNAKE_CASE   ,   references=__SCREAMING_SNAKE_CASE   ,   )

																		lowercase_			:  Union[str, Any]									= metric.compute()
																		# Use accelerator.print to print only on the main process.
																		accelerator.print(F'''epoch {epoch}:'''   ,   __SCREAMING_SNAKE_CASE							)

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
						inner_training_loop()
def       snake_case_	(       ):



						"""simple docstring"""





						lowercase_			:  int									= argparse.ArgumentParser(description='''Simple example of training script.'''							)
						parser.add_argument(
						    '''--mixed_precision'''   ,   type=__SCREAMING_SNAKE_CASE   ,   default=__SCREAMING_SNAKE_CASE   ,   choices=['''no''', '''fp16''', '''bf16''', '''fp8''']   ,   help='''Whether to use mixed precision. Choose'''
						    '''between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'''
						    '''and an Nvidia Ampere GPU.'''   ,   )
						parser.add_argument('''--cpu'''   ,   action='''store_true'''   ,   help='''If passed, will train on the CPU.'''							)
						lowercase_			:  int									= parser.parse_args()
						# New Code #
						# We modify the starting batch size to be an observed batch size of 256, to guarentee an initial CUDA OOM
						lowercase_			:  Union[str, Any]									= {'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 256}
						training_function(__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE							)


if __name__ == "__main__":
							main()



