



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

__snake_case     :List[str]			       =logging.get_logger(__name__)






def  lowerCamelCase_	(      lowerCAmelCase__	:	bool			,	lowerCAmelCase__	:	bool				)       ->							str:

							'''simple docstring'''






							def run_func(lowerCAmelCase__	:	Tuple				):
														@wraps(lowerCAmelCase__				)
														def run_in_eager_mode(*lowerCAmelCase__	:	int			,	**lowerCAmelCase__	:	Optional[Any]				):
																					return func(*lowerCAmelCase__			,	**lowerCAmelCase__				)

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

														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  lowerCamelCase_	(      lowerCAmelCase__	:	int			,	lowerCAmelCase__	:	int			,	lowerCAmelCase__	:	int				)       ->							["tf.Tensor"]:

							'''simple docstring'''






							A      						=    random.Random()
							A      						=    [rng.randint(0			,	vocab_size - 1				) for i in range(batch_size * sequence_length				)]
							return tf.constant(lowerCAmelCase__			,	shape=(batch_size, sequence_length)			,	dtype=tf.intaa				)



class        lowerCAmelCase__					(  _lowerCamelCase       ):
				A_   :     TensorFlowBenchmarkArguments
				A_   :     PretrainedConfig
				A_   :     str								=     "TensorFlow"


				@property
				def __UpperCamelCase						(							self      :  Optional[int]       )		->     Optional[Any]:
											return tf.__version__


				def __UpperCamelCase						(							self      :  int						, __UpperCamelCase      :  str						, __UpperCamelCase      :  int						, __UpperCamelCase      :  int       )		->     float:
											# initialize GPU on separate process
											A      						=    self.args.strategy
											if strategy is None:
																		raise ValueError('A device strategy has to be initialized before using TensorFlow.'       )
											A      						=    self._prepare_inference_func(__UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase       )
											return self._measure_speed(_inference       )


				def __UpperCamelCase						(							self      :  Dict						, __UpperCamelCase      :  str						, __UpperCamelCase      :  int						, __UpperCamelCase      :  int       )		->     float:
											A      						=    self.args.strategy
											if strategy is None:
																		raise ValueError('A device strategy has to be initialized before using TensorFlow.'       )
											A      						=    self._prepare_train_func(__UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase       )
											return self._measure_speed(_train       )


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


				def __UpperCamelCase						(							self      :  Optional[int]						, __UpperCamelCase      :  str						, __UpperCamelCase      :  int						, __UpperCamelCase      :  int       )		->     [Memory, Optional[MemorySummary]]:
											if self.args.is_gpu:
																		tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]						, __UpperCamelCase       )
											A      						=    self.args.strategy
											if strategy is None:
																		raise ValueError('A device strategy has to be initialized before using TensorFlow.'       )

											A      						=    self._prepare_train_func(__UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase       )
											return self._measure_memory(_train       )


				def __UpperCamelCase						(							self      :  List[str]						, __UpperCamelCase      :  str						, __UpperCamelCase      :  int						, __UpperCamelCase      :  int       )		->     Callable[[], None]:
											A      						=    self.config_dict[model_name]

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

											A      						=    (
											    hasattr(__UpperCamelCase						, 'architectures'       )
											    and isinstance(config.architectures						, __UpperCamelCase       )
											    and len(config.architectures       ) > 0
											)
											if not self.args.only_pretrain_model and has_model_class_in_config:
																		try:
																									A      						=    'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
																									A      						=    __import__('transformers'						, fromlist=[model_class]       )
																									A      						=    getattr(__UpperCamelCase						, __UpperCamelCase       )
																									A      						=    model_cls(__UpperCamelCase       )
																		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:
																		A      						=    TF_MODEL_MAPPING[config.__class__](__UpperCamelCase       )

											# encoder-decoder has vocab size saved differently
											A      						=    config.vocab_size if hasattr(__UpperCamelCase						, 'vocab_size'       ) else config.encoder.vocab_size
											A      						=    random_input_ids(__UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase       )

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

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

											A      						=    encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

											return _inference


				def __UpperCamelCase						(							self      :  Tuple						, __UpperCamelCase      :  str						, __UpperCamelCase      :  int						, __UpperCamelCase      :  int       )		->     Callable[[], None]:
											A      						=    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.'       )

											A      						=    (
											    hasattr(__UpperCamelCase						, 'architectures'       )
											    and isinstance(config.architectures						, __UpperCamelCase       )
											    and len(config.architectures       ) > 0
											)
											if not self.args.only_pretrain_model and has_model_class_in_config:
																		try:
																									A      						=    'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
																									A      						=    __import__('transformers'						, fromlist=[model_class]       )
																									A      						=    getattr(__UpperCamelCase						, __UpperCamelCase       )
																									A      						=    model_cls(__UpperCamelCase       )
																		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:
																		A      						=    TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__UpperCamelCase       )

											# encoder-decoder has vocab size saved differently
											A      						=    config.vocab_size if hasattr(__UpperCamelCase						, 'vocab_size'       ) else config.encoder.vocab_size
											A      						=    random_input_ids(__UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase       )

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

											@run_with_tf_optimizations(self.args.eager_mode						, self.args.use_xla       )
											def encoder_train():
																		A      						=    model(__UpperCamelCase						, labels=__UpperCamelCase						, training=__UpperCamelCase       )[0]
																		A      						=    tf.gradients(__UpperCamelCase						, model.trainable_variables       )
																		return gradients

											A      						=    encoder_decoder_train if config.is_encoder_decoder else encoder_train

											return _train


				def __UpperCamelCase						(							self      :  Optional[Any]						, __UpperCamelCase      :  List[str]       )		->     float:
											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(__UpperCamelCase						, 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
																									A      						=    timeit.repeat(
																									    __UpperCamelCase						, repeat=self.args.repeat						, number=10						, )

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


				def __UpperCamelCase						(							self      :  Dict						, __UpperCamelCase      :  Callable[[], None]       )		->     [Memory, MemorySummary]:
											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.'       )
																																A      						=    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.'       )
																																							A      						=    '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()
																																							A      						=    nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx       )
																																							A      						=    nvml.nvmlDeviceGetMemoryInfo(__UpperCamelCase       )
																																							A      						=    meminfo.used
																																							A      						=    Memory(__UpperCamelCase       )
																																							# 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.'       )
																																							A      						=    None
																																else:
																																							A      						=    measure_peak_memory_cpu(__UpperCamelCase       )
																																							A      						=    Memory(__UpperCamelCase       ) if isinstance(__UpperCamelCase						, __UpperCamelCase       ) else memory_bytes
																									if self.args.trace_memory_line_by_line:
																																A      						=    stop_memory_tracing(__UpperCamelCase       )
																																if memory is None:
																																							A      						=    summary.total
																									else:
																																A      						=    None

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





import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
   from .tokenization_fnet import FNetTokenizer
else:
   lowerCamelCase__					:  Optional[Any]      							=				None

lowerCamelCase__					:  List[str]      							=				logging.get_logger(__name__)
lowerCamelCase__					:  List[str]      							=				{'vocab_file': 'spiece.model', 'tokenizer_file': 'tokenizer.json'}

lowerCamelCase__					:  List[str]      							=				{
    'vocab_file': {
        'google/fnet-base': 'https://huggingface.co/google/fnet-base/resolve/main/spiece.model',
        'google/fnet-large': 'https://huggingface.co/google/fnet-large/resolve/main/spiece.model',
    },
    'tokenizer_file': {
        'google/fnet-base': 'https://huggingface.co/google/fnet-base/resolve/main/tokenizer.json',
        'google/fnet-large': 'https://huggingface.co/google/fnet-large/resolve/main/tokenizer.json',
    },
}

lowerCamelCase__					:  Optional[Any]      							=				{
    'google/fnet-base': 512,
    'google/fnet-large': 512,
}

lowerCamelCase__					:  List[Any]      							=				'▁'





class 			lowerCamelCase_				( _SCREAMING_SNAKE_CASE   ):


 '''simple docstring'''







 lowercase_							       =       VOCAB_FILES_NAMES
 lowercase_							       =       PRETRAINED_VOCAB_FILES_MAP
 lowercase_							       =       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 lowercase_							       =       ["input_ids", "token_type_ids"]
 lowercase_							       =       FNetTokenizer



 def __init__(				self    :						List[Any]							,							_lowerCAmelCase    :						Dict=None							,							_lowerCAmelCase    :						Dict=None							,							_lowerCAmelCase    :						List[str]=False							,							_lowerCAmelCase    :						Optional[Any]=True							,							_lowerCAmelCase    :						Tuple=True							,							_lowerCAmelCase    :						List[Any]="<unk>"							,							_lowerCAmelCase    :						Optional[Any]="[SEP]"							,							_lowerCAmelCase    :						Optional[Any]="<pad>"							,							_lowerCAmelCase    :						Optional[int]="[CLS]"							,							_lowerCAmelCase    :						Optional[Any]="[MASK]"							,							**_lowerCAmelCase    :						Any							,							):
       # Mask token behave like a normal word, i.e. include the space before it and
       # is included in the raw text, there should be a match in a non-normalized sentence.
       SCREAMING_SNAKE_CASE_      		=     (
           AddedToken(_lowerCAmelCase							,							lstrip=_lowerCAmelCase							,							rstrip=_lowerCAmelCase							,							normalized=_lowerCAmelCase		)
           if isinstance(_lowerCAmelCase							,							_lowerCAmelCase		)
           else mask_token
       )

       super().__init__(
           _lowerCAmelCase							,							tokenizer_file=_lowerCAmelCase							,							do_lower_case=_lowerCAmelCase							,							remove_space=_lowerCAmelCase							,							keep_accents=_lowerCAmelCase							,							unk_token=_lowerCAmelCase							,							sep_token=_lowerCAmelCase							,							pad_token=_lowerCAmelCase							,							cls_token=_lowerCAmelCase							,							mask_token=_lowerCAmelCase							,							**_lowerCAmelCase							,							)

       SCREAMING_SNAKE_CASE_      		=     do_lower_case
       SCREAMING_SNAKE_CASE_      		=     remove_space
       SCREAMING_SNAKE_CASE_      		=     keep_accents
       SCREAMING_SNAKE_CASE_      		=     vocab_file
       SCREAMING_SNAKE_CASE_      		=     False if not self.vocab_file else True



 def  lowerCAmelCase_    (				self    :						Any							,							_lowerCAmelCase    :						List[int]							,							_lowerCAmelCase    :						Optional[List[int]] = None		):
       SCREAMING_SNAKE_CASE_      		=     [self.sep_token_id]
       SCREAMING_SNAKE_CASE_      		=     [self.cls_token_id]
       if token_ids_a is None:
             return cls + token_ids_a + sep
       return cls + token_ids_a + sep + token_ids_a + sep



 def  lowerCAmelCase_    (				self    :						Optional[int]							,							_lowerCAmelCase    :						List[int]							,							_lowerCAmelCase    :						Optional[List[int]] = None		):
       SCREAMING_SNAKE_CASE_      		=     [self.sep_token_id]
       SCREAMING_SNAKE_CASE_      		=     [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		) * [0] + len(token_ids_a + sep		) * [1]




 def  lowerCAmelCase_    (				self    :						Optional[Any]							,							_lowerCAmelCase    :						str							,							_lowerCAmelCase    :						Optional[str] = None		):
       if not os.path.isdir(_lowerCAmelCase		):
             logger.error(F"Vocabulary path ({save_directory}) should be a directory"		)
             return
       SCREAMING_SNAKE_CASE_      		=     os.path.join(
           _lowerCAmelCase							,							(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']		)

       if os.path.abspath(self.vocab_file		) != os.path.abspath(_lowerCAmelCase		):
             copyfile(self.vocab_file							,							_lowerCAmelCase		)

       return (out_vocab_file,)