




"""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

_A													=      logging.get_logger(__name__)



def 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase				)						->    Union[str, Any]:
      def run_func(__UpperCAmelCase				):
            @wraps(__UpperCAmelCase				)
            def run_in_eager_mode(*__UpperCAmelCase    ,		**__UpperCAmelCase				):
                  return func(*__UpperCAmelCase    ,		**__UpperCAmelCase				)

            @wraps(__UpperCAmelCase				)
            @tf.function(experimental_compile=__UpperCAmelCase				)
            def run_in_graph_mode(*__UpperCAmelCase    ,		**__UpperCAmelCase				):
                  return func(*__UpperCAmelCase    ,		**__UpperCAmelCase				)

            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_				(			__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase				)						->    ["tf.Tensor"]:
      lowerCAmelCase__      :			List[Any]   =	random.Random()
      lowerCAmelCase__      :			Dict   =	[rng.randint(0    ,		vocab_size - 1				) for i in range(batch_size * sequence_length				)]
      return tf.constant(__UpperCAmelCase    ,		shape=(batch_size, sequence_length)    ,		dtype=tf.intaa				)


class _lowerCamelCase    (					a_						):
   _lowerCamelCase    :TensorFlowBenchmarkArguments
   _lowerCamelCase    :PretrainedConfig
   _lowerCamelCase    :str     				=     "TensorFlow"






   @property
   def 	_lowerCAmelCase   (		self       :		List[Any]   )			->						str:

         """simple docstring"""
         return tf.__version__






   def 	_lowerCAmelCase   (		self       :		Dict   ,     UpperCamelCase       :		str   ,     UpperCamelCase       :		int   ,     UpperCamelCase       :		int   )			->						float:

         """simple docstring"""
         lowerCAmelCase__      :			str   =	self.args.strategy
         if strategy is None:
               raise ValueError("""A device strategy has to be initialized before using TensorFlow."""   )
         lowerCAmelCase__      :			Union[str, Any]   =	self._prepare_inference_func(UpperCamelCase   ,     UpperCamelCase   ,     UpperCamelCase   )
         return self._measure_speed(_inference   )






   def 	_lowerCAmelCase   (		self       :		Dict   ,     UpperCamelCase       :		str   ,     UpperCamelCase       :		int   ,     UpperCamelCase       :		int   )			->						float:

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






   def 	_lowerCAmelCase   (		self       :		Union[str, Any]   ,     UpperCamelCase       :		str   ,     UpperCamelCase       :		int   ,     UpperCamelCase       :		int   )			->						[Memory, Optional[MemorySummary]]:

         """simple docstring"""
         # 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   )
         lowerCAmelCase__      :			Optional[int]   =	self.args.strategy
         if strategy is None:
               raise ValueError("""A device strategy has to be initialized before using TensorFlow."""   )
         lowerCAmelCase__      :			Any   =	self._prepare_inference_func(UpperCamelCase   ,     UpperCamelCase   ,     UpperCamelCase   )
         return self._measure_memory(_inference   )






   def 	_lowerCAmelCase   (		self       :		Any   ,     UpperCamelCase       :		str   ,     UpperCamelCase       :		int   ,     UpperCamelCase       :		int   )			->						[Memory, Optional[MemorySummary]]:

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

         lowerCAmelCase__      :			str   =	self._prepare_train_func(UpperCamelCase   ,     UpperCamelCase   ,     UpperCamelCase   )
         return self._measure_memory(_train   )






   def 	_lowerCAmelCase   (		self       :		int   ,     UpperCamelCase       :		str   ,     UpperCamelCase       :		int   ,     UpperCamelCase       :		int   )			->						Callable[[], None]:

         """simple docstring"""
         lowerCAmelCase__      :			int   =	self.config_dict[model_name]

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

         lowerCAmelCase__      :			Optional[Any]   =	(
             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:
                     lowerCAmelCase__      :			str   =	"""TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
                     lowerCAmelCase__      :			Dict   =	__import__("""transformers"""   ,     fromlist=[model_class]   )
                     lowerCAmelCase__      :			int   =	getattr(UpperCamelCase   ,     UpperCamelCase   )
                     lowerCAmelCase__      :			Tuple   =	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:
               lowerCAmelCase__      :			Tuple   =	TF_MODEL_MAPPING[config.__class__](UpperCamelCase   )

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

         lowerCAmelCase__      :			Optional[Any]   =	encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

         return _inference






   def 	_lowerCAmelCase   (		self       :		List[str]   ,     UpperCamelCase       :		str   ,     UpperCamelCase       :		int   ,     UpperCamelCase       :		int   )			->						Callable[[], None]:

         """simple docstring"""
         lowerCAmelCase__      :			List[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."""   )

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

         # encoder-decoder has vocab size saved differently
         lowerCAmelCase__      :			Dict   =	config.vocab_size if hasattr(UpperCamelCase   ,     """vocab_size"""   ) else config.encoder.vocab_size
         lowerCAmelCase__      :			Optional[int]   =	random_input_ids(UpperCamelCase   ,     UpperCamelCase   ,     UpperCamelCase   )

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

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

         lowerCAmelCase__      :			int   =	encoder_decoder_train if config.is_encoder_decoder else encoder_train

         return _train






   def 	_lowerCAmelCase   (		self       :		List[Any]   ,     UpperCamelCase       :		List[Any]   )			->						float:

         """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
                     lowerCAmelCase__      :			Any   =	timeit.repeat(
                         UpperCamelCase   ,     repeat=self.args.repeat   ,     number=10   ,     )

                     return min(UpperCamelCase   ) / 10.0
               except ResourceExhaustedError as e:
                     self.print_fn(f"""Doesn't fit on GPU. {e}"""   )






   def 	_lowerCAmelCase   (		self       :		Dict   ,     UpperCamelCase       :		Callable[[], None]   )			->						[Memory, MemorySummary]:

         """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."""   )
                           lowerCAmelCase__      :			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."""   )
                                 lowerCAmelCase__      :			int   =	"""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()
                                 lowerCAmelCase__      :			Optional[int]   =	nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx   )
                                 lowerCAmelCase__      :			int   =	nvml.nvmlDeviceGetMemoryInfo(UpperCamelCase   )
                                 lowerCAmelCase__      :			int   =	meminfo.used
                                 lowerCAmelCase__      :			Dict   =	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."""   )
                                 lowerCAmelCase__      :			Dict   =	None
                           else:
                                 lowerCAmelCase__      :			Tuple   =	measure_peak_memory_cpu(UpperCamelCase   )
                                 lowerCAmelCase__      :			Any   =	Memory(UpperCamelCase   ) if isinstance(UpperCamelCase   ,     UpperCamelCase   ) else memory_bytes
                     if self.args.trace_memory_line_by_line:
                           lowerCAmelCase__      :			int   =	stop_memory_tracing(UpperCamelCase   )
                           if memory is None:
                                 lowerCAmelCase__      :			str   =	summary.total
                     else:
                           lowerCAmelCase__      :			int   =	None

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











"""simple docstring"""


import os
from collections import deque

import torch
from torch.utils.data import Dataset


class _lowerCamelCase    (					a_						):






   def __init__(		self       :		Tuple   ,     UpperCamelCase       :		List[Any]=""   ,     UpperCamelCase       :		List[str]="train"   )			->						List[Any]:

         """simple docstring"""
         assert os.path.isdir(UpperCamelCase   )

         lowerCAmelCase__      :			Union[str, Any]   =	[]
         lowerCAmelCase__      :			List[Any]   =	os.listdir(UpperCamelCase   )
         for story_filename in story_filenames_list:
               if "summary" in story_filename:
                     continue
               lowerCAmelCase__      :			Any   =	os.path.join(UpperCamelCase   ,     UpperCamelCase   )
               if not os.path.isfile(UpperCamelCase   ):
                     continue
               self.documents.append(UpperCamelCase   )






   def __len__(		self       :		List[Any]   )			->						int:

         """simple docstring"""
         return len(self.documents   )






   def __getitem__(		self       :		str   ,     UpperCamelCase       :		Union[str, Any]   )			->						str:

         """simple docstring"""
         lowerCAmelCase__      :			Any   =	self.documents[idx]
         lowerCAmelCase__      :			List[Any]   =	document_path.split("""/"""   )[-1]
         with open(UpperCamelCase   ,     encoding="""utf-8"""   ) as source:
               lowerCAmelCase__      :			List[str]   =	source.read()
               lowerCAmelCase__ ,   lowerCAmelCase__      :			Union[str, Any]   =	process_story(UpperCamelCase   )
         return document_name, story_lines, summary_lines






def 	lowercase_				(			__UpperCAmelCase				)						->    Union[str, Any]:
      lowerCAmelCase__      :			Tuple   =	list(filter(lambda __UpperCAmelCase				: len(__UpperCAmelCase				) != 0    ,		[line.strip() for line in raw_story.split("""\n"""				)]				)				)

      # for some unknown reason some lines miss a period, add it
      lowerCAmelCase__      :			List[str]   =	[_add_missing_period(__UpperCAmelCase				) for line in nonempty_lines]

      # gather article lines
      lowerCAmelCase__      :			Optional[int]   =	[]
      lowerCAmelCase__      :			Optional[int]   =	deque(__UpperCAmelCase				)
      while True:
            try:
                  lowerCAmelCase__      :			List[Any]   =	lines.popleft()
                  if element.startswith("""@highlight"""				):
                        break
                  story_lines.append(__UpperCAmelCase				)
            except IndexError:
                  # if "@highlight" is absent from the file we pop
                  # all elements until there is None, raising an exception.
                  return story_lines, []

    # gather summary lines
      lowerCAmelCase__      :			Any   =	list(filter(lambda __UpperCAmelCase				: not t.startswith("""@highlight"""				)    ,		__UpperCAmelCase				)				)

      return story_lines, summary_lines



def 	lowercase_				(			__UpperCAmelCase				)						->    Optional[int]:
      lowerCAmelCase__      :			List[str]   =	[""".""", """!""", """?""", """...""", """'""", """`""", """\"""", """\u2019""", """\u2019""", """)"""]
      if line.startswith("""@highlight"""				):
            return line
      if line[-1] in END_TOKENS:
            return line
      return line + "."



def 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase				)						->    List[Any]:
      if len(__UpperCAmelCase				) > block_size:
            return sequence[:block_size]
      else:
            sequence.extend([pad_token_id] * (block_size - len(__UpperCAmelCase				))				)
            return sequence



def 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase				)						->    Tuple:
      lowerCAmelCase__      :			List[str]   =	torch.ones_like(__UpperCAmelCase				)
      lowerCAmelCase__      :			List[str]   =	sequence == pad_token_id
      lowerCAmelCase__      :			str   =	0
      return mask



def 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase				)						->    List[Any]:
      lowerCAmelCase__      :			Tuple   =	[tokenizer.encode(__UpperCAmelCase				) for line in story_lines]
      lowerCAmelCase__      :			List[str]   =	[token for sentence in story_lines_token_ids for token in sentence]
      lowerCAmelCase__      :			int   =	[tokenizer.encode(__UpperCAmelCase				) for line in summary_lines]
      lowerCAmelCase__      :			Union[str, Any]   =	[token for sentence in summary_lines_token_ids for token in sentence]

      return story_token_ids, summary_token_ids



def 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase				)						->    List[str]:
      lowerCAmelCase__      :			Tuple   =	[]
      for sequence in batch:
            lowerCAmelCase__      :			Union[str, Any]   =	-1
            lowerCAmelCase__      :			List[str]   =	[]
            for s in sequence:
                  if s == separator_token_id:
                        sentence_num += 1
                  embeddings.append(sentence_num % 2				)
            batch_embeddings.append(__UpperCAmelCase				)
      return torch.tensor(__UpperCAmelCase				)





