


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							: int      =   logging.get_logger(__name__)






def 							__lowercase     (				__lowerCAmelCase							:       bool      ,					__lowerCAmelCase							:       bool     ):
							def run_func(__lowerCAmelCase							:       Any     ):
														@wraps(__lowerCAmelCase     )
														def run_in_eager_mode(*__lowerCAmelCase							:       Optional[Any]      ,					**__lowerCAmelCase							:       List[Any]     ):
																					return func(*__lowerCAmelCase      ,					**__lowerCAmelCase     )

														@wraps(__lowerCAmelCase     )
														@tf.function(experimental_compile=__lowerCAmelCase     )
														def run_in_graph_mode(*__lowerCAmelCase							:       Union[str, Any]      ,					**__lowerCAmelCase							:       int     ):
																					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 							__lowercase     (				__lowerCAmelCase							:       int      ,					__lowerCAmelCase							:       int      ,					__lowerCAmelCase							:       int     ):
							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  snake_case_     (lowerCamelCase_   ):
					UpperCAmelCase__      :					TensorFlowBenchmarkArguments
					UpperCAmelCase__      :					PretrainedConfig
					UpperCAmelCase__      :					str					   =    "TensorFlow"
					@property
					def 				lowerCamelCase__(   self     :Any   )      ->					Tuple:
												return tf.__version__
					def 				lowerCamelCase__(   self     :str    ,__snake_case     :str    ,__snake_case     :int    ,__snake_case     :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(__snake_case    ,__snake_case    ,__snake_case   )
												return self._measure_speed(_inference   )
					def 				lowerCamelCase__(   self     :Optional[Any]    ,__snake_case     :str    ,__snake_case     :int    ,__snake_case     :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(__snake_case    ,__snake_case    ,__snake_case   )
												return self._measure_speed(_train   )
					def 				lowerCamelCase__(   self     :Union[str, Any]    ,__snake_case     :str    ,__snake_case     :int    ,__snake_case     :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]    ,__snake_case   )
												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(__snake_case    ,__snake_case    ,__snake_case   )
												return self._measure_memory(_inference   )
					def 				lowerCamelCase__(   self     :Tuple    ,__snake_case     :str    ,__snake_case     :int    ,__snake_case     :int   )      ->					[Memory, Optional[MemorySummary]]:
												if self.args.is_gpu:
																			tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]    ,__snake_case   )
												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(__snake_case    ,__snake_case    ,__snake_case   )
												return self._measure_memory(_train   )
					def 				lowerCamelCase__(   self     :int    ,__snake_case     :str    ,__snake_case     :int    ,__snake_case     :int   )      ->					Callable[[], None]:
												a__    =  self.config_dict[model_name]

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

												a__    =  (
												    hasattr(__snake_case    ,'architectures'   )
												    and isinstance(config.architectures    ,__snake_case   )
												    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(__snake_case    ,__snake_case   )
																										a__    =  model_cls(__snake_case   )
																			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__](__snake_case   )

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

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

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

												a__    =  encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

												return _inference
					def 				lowerCamelCase__(   self     :Optional[int]    ,__snake_case     :str    ,__snake_case     :int    ,__snake_case     :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(__snake_case    ,'architectures'   )
												    and isinstance(config.architectures    ,__snake_case   )
												    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(__snake_case    ,__snake_case   )
																										a__    =  model_cls(__snake_case   )
																			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__](__snake_case   )

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

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

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

												a__    =  encoder_decoder_train if config.is_encoder_decoder else encoder_train

												return _train
					def 				lowerCamelCase__(   self     :Union[str, Any]    ,__snake_case     :Optional[int]   )      ->					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(__snake_case    ,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(
																										    __snake_case    ,repeat=self.args.repeat    ,number=10    ,)

																										return min(__snake_case   ) / 10.0
																			except ResourceExhaustedError as e:
																										self.print_fn(F'Doesn\'t fit on GPU. {e}'   )



					def 				lowerCamelCase__(   self     :str    ,__snake_case     :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(__snake_case   )
																																								a__    =  meminfo.used
																																								a__    =  Memory(__snake_case   )
																																								# 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(__snake_case   )
																																								a__    =  Memory(__snake_case   ) if isinstance(__snake_case    ,__snake_case   ) else memory_bytes
																										if self.args.trace_memory_line_by_line:
																																	a__    =  stop_memory_tracing(__snake_case   )
																																	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








# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..utils import cached_file


# docstyle-ignore
snake_case							: Dict      =   '''
Human: <<task>>

Assistant: '''


snake_case							: Optional[int]      =   '''huggingface-tools/default-prompts'''
snake_case							: Tuple      =   {'''chat''': '''chat_prompt_template.txt''', '''run''': '''run_prompt_template.txt'''}






def 							__lowercase     (				__lowerCAmelCase							:       Any      ,					__lowerCAmelCase							:       Union[str, Any]      ,					__lowerCAmelCase							:       Any="run"     ):
							if prompt_or_repo_id is None:
														a__    =  DEFAULT_PROMPTS_REPO

							# prompt is considered a repo ID when it does not contain any kind of space
							if re.search('\\s'      ,					__lowerCAmelCase     ) is not None:
														return prompt_or_repo_id

							a__    =  cached_file(
							    __lowerCAmelCase      ,					PROMPT_FILES[mode]      ,					repo_type='dataset'      ,					user_agent={'agent': agent_name}     )
							with open(__lowerCAmelCase      ,					'r'      ,					encoding='utf-8'     ) as f:
														return f.read()




