






'''simple docstring'''


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







def 	__magic_name__( _A						,      _A					):
      '''simple docstring'''





      def run_func(_A					):
            @wraps(lowerCAmelCase__					)
            def run_in_eager_mode(*_A						,      **_A					):
                  return func(*lowerCAmelCase__						,      **lowerCAmelCase__					)

            @wraps(lowerCAmelCase__					)
            @tf.function(experimental_compile=lowerCAmelCase__					)
            def run_in_graph_mode(*_A						,      **_A					):
                  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 	__magic_name__( _A						,      _A						,      _A					):
      '''simple docstring'''





      UpperCamelCase__         = random.Random()
      UpperCamelCase__         = [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    _SCREAMING_SNAKE_CASE							(		__A	):


   '''simple docstring'''





   __a							:							TensorFlowBenchmarkArguments
   __a							:							PretrainedConfig
   __a							:							str					     =					"TensorFlow"






   @property
   def 				A      (   self							:						Union[str, Any]			)							->    Dict:

         '''simple docstring'''







         return tf.__version__






   def 				A      (   self							:						Tuple ,      lowercase							:						List[str] ,      lowercase							:						Optional[int] ,      lowercase							:						Any			)							->    List[str]:

         '''simple docstring'''







         UpperCamelCase__         = self.args.strategy
         if strategy is None:
               raise ValueError("""A device strategy has to be initialized before using TensorFlow."""			)
         UpperCamelCase__         = self._prepare_inference_func(UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__			)
         return self._measure_speed(_inference			)






   def 				A      (   self							:						Optional[Any] ,      lowercase							:						int ,      lowercase							:						Optional[Any] ,      lowercase							:						int			)							->    Tuple:

         '''simple docstring'''







         UpperCamelCase__         = self.args.strategy
         if strategy is None:
               raise ValueError("""A device strategy has to be initialized before using TensorFlow."""			)
         UpperCamelCase__         = self._prepare_train_func(UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__			)
         return self._measure_speed(_train			)






   def 				A      (   self							:						Dict ,      lowercase							:						Tuple ,      lowercase							:						Any ,      lowercase							:						Tuple			)							->    str:

         '''simple docstring'''







         if self.args.is_gpu:
               tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx] ,      UpperCamelCase__			)
         UpperCamelCase__         = self.args.strategy
         if strategy is None:
               raise ValueError("""A device strategy has to be initialized before using TensorFlow."""			)
         UpperCamelCase__         = self._prepare_inference_func(UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__			)
         return self._measure_memory(_inference			)






   def 				A      (   self							:						Optional[int] ,      lowercase							:						str ,      lowercase							:						Dict ,      lowercase							:						Any			)							->    Union[str, Any]:

         '''simple docstring'''







         if self.args.is_gpu:
               tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx] ,      UpperCamelCase__			)
         UpperCamelCase__         = self.args.strategy
         if strategy is None:
               raise ValueError("""A device strategy has to be initialized before using TensorFlow."""			)

         UpperCamelCase__         = self._prepare_train_func(UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__			)
         return self._measure_memory(_train			)






   def 				A      (   self							:						List[str] ,      lowercase							:						Any ,      lowercase							:						Optional[Any] ,      lowercase							:						int			)							->    Dict:

         '''simple docstring'''







         UpperCamelCase__         = self.config_dict[model_name]

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

         UpperCamelCase__         = (
             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:
                     UpperCamelCase__         = 'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
                     UpperCamelCase__         = __import__("""transformers""" ,      fromlist=[model_class]			)
                     UpperCamelCase__         = getattr(UpperCamelCase__ ,      UpperCamelCase__			)
                     UpperCamelCase__         = 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:
               UpperCamelCase__         = TF_MODEL_MAPPING[config.__class__](UpperCamelCase__			)

         # encoder-decoder has vocab size saved differently
         UpperCamelCase__         = config.vocab_size if hasattr(UpperCamelCase__ ,      """vocab_size"""			) else config.encoder.vocab_size
         UpperCamelCase__         = 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__			)

         UpperCamelCase__         = encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

         return _inference






   def 				A      (   self							:						Tuple ,      lowercase							:						Tuple ,      lowercase							:						str ,      lowercase							:						List[str]			)							->    Union[str, Any]:

         '''simple docstring'''







         UpperCamelCase__         = 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."""			)

         UpperCamelCase__         = (
             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:
                     UpperCamelCase__         = 'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
                     UpperCamelCase__         = __import__("""transformers""" ,      fromlist=[model_class]			)
                     UpperCamelCase__         = getattr(UpperCamelCase__ ,      UpperCamelCase__			)
                     UpperCamelCase__         = 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:
               UpperCamelCase__         = TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](UpperCamelCase__			)

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

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

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

         UpperCamelCase__         = encoder_decoder_train if config.is_encoder_decoder else encoder_train

         return _train






   def 				A      (   self							:						Tuple ,      lowercase							:						Any			)							->    str:

         '''simple docstring'''







         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
                     UpperCamelCase__         = timeit.repeat(
                         UpperCamelCase__ ,      repeat=self.args.repeat ,      number=1_0 ,      )

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







   def 				A      (   self							:						Union[str, Any] ,      lowercase							:						Tuple			)							->    str:

         '''simple docstring'''







         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."""			)
                           UpperCamelCase__         = 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."""			)
                                 UpperCamelCase__         = '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()
                                 UpperCamelCase__         = nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx			)
                                 UpperCamelCase__         = nvml.nvmlDeviceGetMemoryInfo(UpperCamelCase__			)
                                 UpperCamelCase__         = meminfo.used
                                 UpperCamelCase__         = 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."""			)
                                 UpperCamelCase__         = None
                           else:
                                 UpperCamelCase__         = measure_peak_memory_cpu(UpperCamelCase__			)
                                 UpperCamelCase__         = Memory(UpperCamelCase__			) if isinstance(UpperCamelCase__ ,      UpperCamelCase__			) else memory_bytes
                     if self.args.trace_memory_line_by_line:
                           UpperCamelCase__         = stop_memory_tracing(UpperCamelCase__			)
                           if memory is None:
                                 UpperCamelCase__         = summary.total
                     else:
                           UpperCamelCase__         = None

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








'''simple docstring'''


from collections import Counter
from pathlib import Path
from typing import Optional, Tuple

import yaml
class    _SCREAMING_SNAKE_CASE							(		yaml.SafeLoader	):


			'''simple docstring'''





			def 				A      (   self							:						List[str] ,      lowercase							:						List[Any]			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = [self.constructed_objects[key_node] for key_node, _ in node.value]
									UpperCamelCase__         = [tuple(lowercase			) if isinstance(lowercase ,      lowercase			) else key for key in keys]
									UpperCamelCase__         = Counter(lowercase			)
									UpperCamelCase__         = [key for key in counter if counter[key] > 1]
									if duplicate_keys:
															raise TypeError(f"Got duplicate yaml keys: {duplicate_keys}"			)







			def 				A      (   self							:						List[str] ,      lowercase							:						int ,      lowercase							:						str=False			)							->    Any:

									'''simple docstring'''







									UpperCamelCase__         = super().construct_mapping(lowercase ,      deep=lowercase			)
									self._check_no_duplicates_on_constructed_node(lowercase			)
									return mapping







def 	__magic_name__( _A					):
						'''simple docstring'''





						UpperCamelCase__         = list(readme_content.splitlines()					)
						if full_content and full_content[0] == "---" and "---" in full_content[1:]:
												UpperCamelCase__         = full_content[1:].index("""---"""					) + 1
												UpperCamelCase__         = """\n""".join(full_content[1:sep_idx]					)
												return yamlblock, "\n".join(full_content[sep_idx + 1 :]					)

						return None, "\n".join(_A					)
class    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE	):


			'''simple docstring'''





			__a							:							Tuple					     =					{"train_eval_index"}  # train-eval-index in the YAML metadata






			@classmethod
			def 				A      (   cls							:						int ,      lowercase							:						Path			)							->    "DatasetMetadata":

									'''simple docstring'''







									with open(lowercase ,      encoding="""utf-8"""			) as readme_file:
															UpperCamelCase__    ,   UpperCamelCase__         = _split_yaml_from_readme(readme_file.read()			)
									if yaml_string is not None:
															return cls.from_yaml_string(lowercase			)
									else:
															return cls()






			def 				A      (   self							:						int ,      lowercase							:						Path			)							->    Dict:

									'''simple docstring'''







									if path.exists():
															with open(lowercase ,      encoding="""utf-8"""			) as readme_file:
																					UpperCamelCase__         = readme_file.read()
									else:
															UpperCamelCase__         = None
									UpperCamelCase__         = self._to_readme(lowercase			)
									with open(lowercase ,      """w""" ,      encoding="""utf-8"""			) as readme_file:
															readme_file.write(lowercase			)






			def 				A      (   self							:						Any ,      lowercase							:						Optional[str] = None			)							->    str:

									'''simple docstring'''







									if readme_content is not None:
															UpperCamelCase__    ,   UpperCamelCase__         = _split_yaml_from_readme(lowercase			)
															UpperCamelCase__         = """---\n""" + self.to_yaml_string() + """---\n""" + content
									else:
															UpperCamelCase__         = """---\n""" + self.to_yaml_string() + """---\n"""
									return full_content






			@classmethod
			def 				A      (   cls							:						Tuple ,      lowercase							:						str			)							->    "DatasetMetadata":

									'''simple docstring'''







									UpperCamelCase__         = yaml.load(lowercase ,      Loader=_NoDuplicateSafeLoader			) or {}

									# Convert the YAML keys to DatasetMetadata fields
									UpperCamelCase__         = {
									    (key.replace("""-""" ,      """_"""			) if key.replace("""-""" ,      """_"""			) in cls._FIELDS_WITH_DASHES else key): value
									    for key, value in metadata_dict.items()
									}
									return cls(**lowercase			)







			def 				A      (   self							:						Dict			)							->    str:

									'''simple docstring'''







									return yaml.safe_dump(
									    {
									        (key.replace("""_""" ,      """-"""			) if key in self._FIELDS_WITH_DASHES else key): value
									        for key, value in self.items()
									    } ,      sort_keys=lowercase ,      allow_unicode=lowercase ,      encoding="""utf-8""" ,      ).decode("""utf-8"""			)
lowerCamelCase_				: str    	=							{
    '''image-classification''': [],
    '''translation''': [],
    '''image-segmentation''': [],
    '''fill-mask''': [],
    '''automatic-speech-recognition''': [],
    '''token-classification''': [],
    '''sentence-similarity''': [],
    '''audio-classification''': [],
    '''question-answering''': [],
    '''summarization''': [],
    '''zero-shot-classification''': [],
    '''table-to-text''': [],
    '''feature-extraction''': [],
    '''other''': [],
    '''multiple-choice''': [],
    '''text-classification''': [],
    '''text-to-image''': [],
    '''text2text-generation''': [],
    '''zero-shot-image-classification''': [],
    '''tabular-classification''': [],
    '''tabular-regression''': [],
    '''image-to-image''': [],
    '''tabular-to-text''': [],
    '''unconditional-image-generation''': [],
    '''text-retrieval''': [],
    '''text-to-speech''': [],
    '''object-detection''': [],
    '''audio-to-audio''': [],
    '''text-generation''': [],
    '''conversational''': [],
    '''table-question-answering''': [],
    '''visual-question-answering''': [],
    '''image-to-text''': [],
    '''reinforcement-learning''': [],
    '''voice-activity-detection''': [],
    '''time-series-forecasting''': [],
    '''document-question-answering''': [],
}


if __name__ == "__main__":
					from argparse import ArgumentParser

					lowerCamelCase_				: Tuple    	=							ArgumentParser(usage='''Validate the yaml metadata block of a README.md file.''')
					ap.add_argument('''readme_filepath''')
					lowerCamelCase_				: str    	=							ap.parse_args()

					lowerCamelCase_				: List[str]    	=							Path(args.readme_filepath)
					lowerCamelCase_				: Tuple    	=							DatasetMetadata.from_readme(readme_filepath)
					print(dataset_metadata)
					dataset_metadata.to_readme(readme_filepath)
