






import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
	import tensorflow as tf
	from tensorflow.python.framework.errors_impl import ResourceExhaustedError

	from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
	import pyanvml.pyanvml as nvml

__lowerCamelCase :      List[Any]										=				logging.get_logger(__name__)







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							bool					,  snake_case_ :							bool       ):
	def run_func(snake_case_ :							str       ):
		@wraps(snake_case_       )
		def run_in_eager_mode(*snake_case_ :							str					,  **snake_case_ :							Union[str, Any]       ):
			return func(*snake_case_					,  **snake_case_       )

		@wraps(snake_case_       )
		@tf.function(experimental_compile=snake_case_       )
		def run_in_graph_mode(*snake_case_ :							List[Any]					,  **snake_case_ :							List[Any]       ):
			return func(*snake_case_					,  **snake_case_       )

		if do_eager_mode is True:
			if use_xla is not False:
				raise ValueError(
				    "Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`."       )
			return run_in_eager_mode
		else:
			return run_in_graph_mode

	return run_func







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							int					,  snake_case_ :							int					,  snake_case_ :							int       ):
	snake_case__			:				Dict       =  random.Random()
	snake_case__			:				List[Any]       =  [rng.randint(0					,  vocab_size - 1       ) for i in range(batch_size * sequence_length       )]
	return tf.constant(snake_case_					,  shape=(batch_size, sequence_length)					,  dtype=tf.intaa       )





class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					a_     =		42
					a_     =		42
					a_     =		"TensorFlow"



					@property
					def 	_lowercase					( self					:  List[str]	):
						return tf.__version__



					def 	_lowercase					( self					:  List[str]							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						# initialize GPU on separate process
						snake_case__			:				str       =  self.args.strategy
						if strategy is None:
							raise ValueError("A device strategy has to be initialized before using TensorFlow."	)
						snake_case__			:				Dict       =  self._prepare_inference_func(__A							,		__A							,		__A	)
						return self._measure_speed(_inference	)



					def 	_lowercase					( self					:  Tuple							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						snake_case__			:				Optional[int]       =  self.args.strategy
						if strategy is None:
							raise ValueError("A device strategy has to be initialized before using TensorFlow."	)
						snake_case__			:				Any       =  self._prepare_train_func(__A							,		__A							,		__A	)
						return self._measure_speed(_train	)



					def 	_lowercase					( self					:  List[Any]							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						# initialize GPU on separate process
						if self.args.is_gpu:
							tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]							,		__A	)
						snake_case__			:				List[str]       =  self.args.strategy
						if strategy is None:
							raise ValueError("A device strategy has to be initialized before using TensorFlow."	)
						snake_case__			:				Optional[Any]       =  self._prepare_inference_func(__A							,		__A							,		__A	)
						return self._measure_memory(_inference	)



					def 	_lowercase					( self					:  str							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						if self.args.is_gpu:
							tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]							,		__A	)
						snake_case__			:				List[Any]       =  self.args.strategy
						if strategy is None:
							raise ValueError("A device strategy has to be initialized before using TensorFlow."	)

						snake_case__			:				int       =  self._prepare_train_func(__A							,		__A							,		__A	)
						return self._measure_memory(_train	)



					def 	_lowercase					( self					:  Union[str, Any]							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						snake_case__			:				int       =  self.config_dict[model_name]

						if self.args.fpaa:
							raise NotImplementedError("Mixed precision is currently not supported."	)

						snake_case__			:				Tuple       =  (
						    hasattr(__A							,		"architectures"	)
						    and isinstance(config.architectures							,		__A	)
						    and len(config.architectures	) > 0
						)
						if not self.args.only_pretrain_model and has_model_class_in_config:
							try:
								snake_case__			:				Dict       =  "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
								snake_case__			:				Union[str, Any]       =  __import__("transformers"							,		fromlist=[model_class]	)
								snake_case__			:				Any       =  getattr(__A							,		__A	)
								snake_case__			:				Dict       =  model_cls(__A	)
							except ImportError:
								raise ImportError(
								    f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
								    " set `--only_pretrain_model` or `args.only_pretrain_model=True`."	)
						else:
							snake_case__			:				Dict       =  TF_MODEL_MAPPING[config.__class__](__A	)

						# encoder-decoder has vocab size saved differently
						snake_case__			:				Optional[int]       =  config.vocab_size if hasattr(__A							,		"vocab_size"	) else config.encoder.vocab_size
						snake_case__			:				List[Any]       =  random_input_ids(__A							,		__A							,		__A	)

						@run_with_tf_optimizations(self.args.eager_mode							,		self.args.use_xla	)
						def encoder_decoder_forward():
							return model(__A							,		decoder_input_ids=__A							,		training=__A	)

						@run_with_tf_optimizations(self.args.eager_mode							,		self.args.use_xla	)
						def encoder_forward():
							return model(__A							,		training=__A	)

						snake_case__			:				Optional[int]       =  encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

						return _inference



					def 	_lowercase					( self					:  List[str]							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						snake_case__			:				Optional[Any]       =  self.config_dict[model_name]

						if self.args.eager_mode is not False:
							raise ValueError("Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`."	)

						if self.args.fpaa:
							raise NotImplementedError("Mixed precision is currently not supported."	)

						snake_case__			:				Any       =  (
						    hasattr(__A							,		"architectures"	)
						    and isinstance(config.architectures							,		__A	)
						    and len(config.architectures	) > 0
						)
						if not self.args.only_pretrain_model and has_model_class_in_config:
							try:
								snake_case__			:				Dict       =  "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
								snake_case__			:				List[Any]       =  __import__("transformers"							,		fromlist=[model_class]	)
								snake_case__			:				Optional[int]       =  getattr(__A							,		__A	)
								snake_case__			:				str       =  model_cls(__A	)
							except ImportError:
								raise ImportError(
								    f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
								    " set `--only_pretrain_model` or `args.only_pretrain_model=True`."	)
						else:
							snake_case__			:				Union[str, Any]       =  TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__A	)

						# encoder-decoder has vocab size saved differently
						snake_case__			:				Union[str, Any]       =  config.vocab_size if hasattr(__A							,		"vocab_size"	) else config.encoder.vocab_size
						snake_case__			:				List[str]       =  random_input_ids(__A							,		__A							,		__A	)

						@run_with_tf_optimizations(self.args.eager_mode							,		self.args.use_xla	)
						def encoder_decoder_train():
							snake_case__			:				str       =  model(__A							,		decoder_input_ids=__A							,		labels=__A							,		training=__A	)[0]
							snake_case__			:				Dict       =  tf.gradients(__A							,		model.trainable_variables	)
							return gradients

						@run_with_tf_optimizations(self.args.eager_mode							,		self.args.use_xla	)
						def encoder_train():
							snake_case__			:				Optional[Any]       =  model(__A							,		labels=__A							,		training=__A	)[0]
							snake_case__			:				Dict       =  tf.gradients(__A							,		model.trainable_variables	)
							return gradients

						snake_case__			:				int       =  encoder_decoder_train if config.is_encoder_decoder else encoder_train

						return _train



					def 	_lowercase					( self					:  int							,		__A					:  List[Any]	):
						with self.args.strategy.scope():
							try:
								if self.args.is_tpu or self.args.use_xla:
									# run additional 10 times to stabilize compilation for tpu
									logger.info("Do inference on TPU. Running model 5 times to stabilize compilation"	)
									timeit.repeat(__A							,		repeat=1							,		number=5	)

								# as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
								snake_case__			:				Optional[Any]       =  timeit.repeat(
								    __A							,		repeat=self.args.repeat							,		number=1_0							,		)

								return min(__A	) / 1_0.0
							except ResourceExhaustedError as e:
								self.print_fn(f'''Doesn\'t fit on GPU. {e}'''	)






					def 	_lowercase					( self					:  str							,		__A					:  Callable[[], None]	):
						logger.info(
						    "Note that TensorFlow allocates more memory than "
						    "it might need to speed up computation. "
						    "The memory reported here corresponds to the memory "
						    "reported by `nvidia-smi`, which can vary depending "
						    "on total available memory on the GPU that is used."	)
						with self.args.strategy.scope():
							try:
								if self.args.trace_memory_line_by_line:
									if not self.args.eager_mode:
										raise ValueError(
										    "`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory"
										    " consumption line by line."	)
									snake_case__			:				Optional[int]       =  start_memory_tracing("transformers"	)

								if self.args.is_tpu:
									# tpu
									raise NotImplementedError(
									    "Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking"
									    " with `args.memory=False`"	)
								elif self.args.is_gpu:
									# gpu
									if not is_pyanvml_available():
										logger.warning(
										    "py3nvml not installed, we won't log GPU memory usage. "
										    "Install py3nvml (pip install py3nvml) to log information about GPU."	)
										snake_case__			:				List[str]       =  "N/A"
									else:
										logger.info(
										    "Measuring total GPU usage on GPU device. Make sure to not have additional processes"
										    " running on the same GPU."	)
										# init nvml
										nvml.nvmlInit()
										func()
										snake_case__			:				Any       =  nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
										snake_case__			:				Optional[Any]       =  nvml.nvmlDeviceGetMemoryInfo(__A	)
										snake_case__			:				Optional[int]       =  meminfo.used
										snake_case__			:				Any       =  Memory(__A	)
										# shutdown nvml
										nvml.nvmlShutdown()
								else:
									# cpu
									if self.args.trace_memory_line_by_line:
										logger.info(
										    "When enabling line by line tracing, the max peak memory for CPU is inaccurate in"
										    " TensorFlow."	)
										snake_case__			:				int       =  None
									else:
										snake_case__			:				Any       =  measure_peak_memory_cpu(__A	)
										snake_case__			:				Tuple       =  Memory(__A	) if isinstance(__A							,		__A	) else memory_bytes
								if self.args.trace_memory_line_by_line:
									snake_case__			:				Optional[int]       =  stop_memory_tracing(__A	)
									if memory is None:
										snake_case__			:				Dict       =  summary.total
								else:
									snake_case__			:				List[str]       =  None

								return memory, summary
							except ResourceExhaustedError as e:
								self.print_fn(f'''Doesn\'t fit on GPU. {e}'''	)
								return "N/A", None











import json
from typing import Dict, List, Optional, Tuple, Union

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding, EncodedInput
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import PaddingStrategy, logging
from .tokenization_led import LEDTokenizer


__lowerCamelCase :      int										=				logging.get_logger(__name__)


__lowerCamelCase :      List[str]										=				{"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt""", """tokenizer_file""": """tokenizer.json"""}

__lowerCamelCase :      Any										=				{
    """vocab_file""": {
        """allenai/led-base-16384""": """https://huggingface.co/allenai/led-base-16384/resolve/main/vocab.json""",
    },
    """merges_file""": {
        """allenai/led-base-16384""": """https://huggingface.co/allenai/led-base-16384/resolve/main/merges.txt""",
    },
    """tokenizer_file""": {
        """allenai/led-base-16384""": """https://huggingface.co/allenai/led-base-16384/resolve/main/tokenizer.json""",
    },
}

__lowerCamelCase :      Optional[int]										=				{
    """allenai/led-base-16384""": 1_6384,
}



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					a_     =		VOCAB_FILES_NAMES
					a_     =		PRETRAINED_VOCAB_FILES_MAP
					a_     =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					a_     =		LEDTokenizer
					a_     =		["input_ids", "attention_mask"]



					def __init__( self					:  int							,		__A					:  Optional[int]=None							,		__A					:  Any=None							,		__A					:  Any=None							,		__A					:  int="replace"							,		__A					:  Optional[Any]="<s>"							,		__A					:  List[Any]="</s>"							,		__A					:  Dict="</s>"							,		__A					:  Union[str, Any]="<s>"							,		__A					:  List[Any]="<unk>"							,		__A					:  Optional[Any]="<pad>"							,		__A					:  Optional[int]="<mask>"							,		__A					:  str=False							,		__A					:  Optional[Any]=True							,		**__A					:  List[Any]							,		):
						super().__init__(
						    __A							,		__A							,		tokenizer_file=__A							,		errors=__A							,		bos_token=__A							,		eos_token=__A							,		sep_token=__A							,		cls_token=__A							,		unk_token=__A							,		pad_token=__A							,		mask_token=__A							,		add_prefix_space=__A							,		trim_offsets=__A							,		**__A							,		)

						snake_case__			:				List[Any]       =  json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()	)
						if pre_tok_state.get("add_prefix_space"							,		__A	) != add_prefix_space:
							snake_case__			:				List[str]       =  getattr(__A							,		pre_tok_state.pop("type"	)	)
							snake_case__			:				Tuple       =  add_prefix_space
							snake_case__			:				Any       =  pre_tok_class(**__A	)

						snake_case__			:				List[Any]       =  add_prefix_space

						# the pre_tokenizer is already updated in the GPT2TokenizerFast `__init__`
						snake_case__			:				str       =  "post_processor"
						snake_case__			:				Union[str, Any]       =  getattr(self.backend_tokenizer							,		__A							,		__A	)
						if tokenizer_component_instance:
							snake_case__			:				Any       =  json.loads(tokenizer_component_instance.__getstate__()	)

							# The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
							if "sep" in state:
								snake_case__			:				int       =  tuple(state["sep"]	)
							if "cls" in state:
								snake_case__			:				Optional[int]       =  tuple(state["cls"]	)

							snake_case__			:				Union[str, Any]       =  False

							if state.get("add_prefix_space"							,		__A	) != add_prefix_space:
								snake_case__			:				Optional[Any]       =  add_prefix_space
								snake_case__			:				List[str]       =  True

							if state.get("trim_offsets"							,		__A	) != trim_offsets:
								snake_case__			:				str       =  trim_offsets
								snake_case__			:				Tuple       =  True

							if changes_to_apply:
								snake_case__			:				Any       =  getattr(__A							,		state.pop("type"	)	)
								snake_case__			:				Optional[Any]       =  component_class(**__A	)
								setattr(self.backend_tokenizer							,		__A							,		__A	)



					@property
					# Copied from transformers.models.bart.tokenization_bart_fast.BartTokenizerFast.mask_token with BART->LED
					def 	_lowercase					( self					:  Union[str, Any]	):
						if self._mask_token is None:
							if self.verbose:
								logger.error("Using mask_token, but it is not set yet."	)
							return None
						return str(self._mask_token	)



					@mask_token.setter
					def 	_lowercase					( self					:  str							,		__A					:  int	):
						snake_case__			:				str       =  AddedToken(__A							,		lstrip=__A							,		rstrip=__A	) if isinstance(__A							,		__A	) else value
						snake_case__			:				Optional[int]       =  value



					def 	_lowercase					( self					:  List[str]							,		*__A					:  Union[str, Any]							,		**__A					:  int	):
						snake_case__			:				Dict       =  kwargs.get("is_split_into_words"							,		__A	)

						if is_split_into_words and not self.add_prefix_space:
							raise ValueError(
							    f'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
							    "to use it with pretokenized inputs."	)

						return super()._batch_encode_plus(*__A							,		**__A	)



					def 	_lowercase					( self					:  Union[str, Any]							,		*__A					:  Tuple							,		**__A					:  Union[str, Any]	):
						snake_case__			:				List[str]       =  kwargs.get("is_split_into_words"							,		__A	)

						if is_split_into_words and not self.add_prefix_space:
							raise ValueError(
							    f'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
							    "to use it with pretokenized inputs."	)

						return super()._encode_plus(*__A							,		**__A	)



					def 	_lowercase					( self					:  Union[str, Any]							,		__A					:  str							,		__A					:  Optional[str] = None	):
						snake_case__			:				str       =  self._tokenizer.model.save(__A							,		name=__A	)
						return tuple(__A	)



					def 	_lowercase					( self					:  Union[str, Any]							,		__A					:  int							,		__A					:  str=None	):
						snake_case__			:				int       =  [self.bos_token_id] + token_ids_a + [self.eos_token_id]
						if token_ids_a is None:
							return output

						return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]



					def 	_lowercase					( self					:  Any							,		__A					:  List[int]							,		__A					:  Optional[List[int]] = None	):
						snake_case__			:				Optional[Any]       =  [self.sep_token_id]
						snake_case__			:				Any       =  [self.cls_token_id]

						if token_ids_a is None:
							return len(cls + token_ids_a + sep	) * [0]
						return len(cls + token_ids_a + sep + sep + token_ids_a + sep	) * [0]






					def 	_lowercase					( self					:  str							,		__A					:  Union[Dict[str, EncodedInput], BatchEncoding]							,		__A					:  Optional[int] = None							,		__A					:  PaddingStrategy = PaddingStrategy.DO_NOT_PAD							,		__A					:  Optional[int] = None							,		__A					:  Optional[bool] = None							,		):
						snake_case__			:				Dict       =  super()._pad(
						    encoded_inputs=__A							,		max_length=__A							,		padding_strategy=__A							,		pad_to_multiple_of=__A							,		return_attention_mask=__A							,		)

						# Load from model defaults
						if return_attention_mask is None:
							snake_case__			:				Optional[Any]       =  "attention_mask" in self.model_input_names

						if return_attention_mask and "global_attention_mask" in encoded_inputs:
							snake_case__			:				List[Any]       =  encoded_inputs[self.model_input_names[0]]
							# `global_attention_mask` need to have the same length as other (sequential) inputs.
							snake_case__			:				Union[str, Any]       =  len(encoded_inputs["global_attention_mask"]	) != len(__A	)

							if needs_to_be_padded:
								snake_case__			:				Dict       =  len(__A	) - len(encoded_inputs["global_attention_mask"]	)

								if self.padding_side == "right":
									# Use `-1` since `0` in `global_attention_mask` means `local attention` instead of `not to attend`
									snake_case__			:				Optional[Any]       =  (
									    encoded_inputs["global_attention_mask"] + [-1] * difference
									)
								elif self.padding_side == "left":
									snake_case__			:				Optional[Any]       =  [-1] * difference + encoded_inputs[
									    "global_attention_mask"
									]
								else:
									raise ValueError("Invalid padding strategy:" + str(self.padding_side	)	)

						return encoded_inputs



