


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

__UpperCamelCase				:				Tuple  				= logging.get_logger(__name__)

def 	_a     ( SCREAMING_SNAKE_CASE		: Optional[int]			,       SCREAMING_SNAKE_CASE		: Tuple   ):







 """simple docstring"""



 def run_func(SCREAMING_SNAKE_CASE		: Dict   ):
  @wraps(SCREAMING_SNAKE_CASE   )
  def run_in_eager_mode(*SCREAMING_SNAKE_CASE		: Dict			,       **SCREAMING_SNAKE_CASE		: List[str]   ):
   return func(*SCREAMING_SNAKE_CASE			,       **SCREAMING_SNAKE_CASE   )

  @wraps(SCREAMING_SNAKE_CASE   )
  @tf.function(experimental_compile=SCREAMING_SNAKE_CASE   )
  def run_in_graph_mode(*SCREAMING_SNAKE_CASE		: Dict			,       **SCREAMING_SNAKE_CASE		: List[Any]   ):
   return func(*SCREAMING_SNAKE_CASE			,       **SCREAMING_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 	_a     ( SCREAMING_SNAKE_CASE		: Any			,       SCREAMING_SNAKE_CASE		: str			,       SCREAMING_SNAKE_CASE		: int   ):







 """simple docstring"""



 UpperCamelCase__    :			str      				=      random.Random()
 UpperCamelCase__    :			Union[str, Any]      				=      [rng.randint(0			,       vocab_size - 1   ) for i in range(batch_size * sequence_length   )]
 return tf.constant(SCREAMING_SNAKE_CASE			,       shape=(batch_size, sequence_length)			,       dtype=tf.intaa   )






class 					__magic_name__		(   a__):
     A:						List[str]								=					4_2
     A:						Union[str, Any]								=					4_2
     A:						List[Any]								=					"TensorFlow"


     @property
     def    UpperCAmelCase__		(					self							:   Optional[Any]	)     ->					Optional[int]:





      '''simple docstring'''






      return tf.__version__


     def    UpperCAmelCase__		(					self							:   int							,		lowerCamelCase__							:   List[Any]							,		lowerCamelCase__							:   str							,		lowerCamelCase__							:   Optional[Any]	)     ->					float:





      '''simple docstring'''






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


     def    UpperCAmelCase__		(					self							:   Union[str, Any]							,		lowerCamelCase__							:   Dict							,		lowerCamelCase__							:   Optional[Any]							,		lowerCamelCase__							:   Any	)     ->					float:





      '''simple docstring'''






      UpperCamelCase__    :			Dict      				=      self.args.strategy
      if strategy is None:
       raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''	)
      UpperCamelCase__    :			Optional[int]      				=      self._prepare_train_func(__lowerCAmelCase							,		__lowerCAmelCase							,		__lowerCAmelCase	)
      return self._measure_speed(_train	)


     def    UpperCAmelCase__		(					self							:   Tuple							,		lowerCamelCase__							:   Union[str, Any]							,		lowerCamelCase__							:   Union[str, Any]							,		lowerCamelCase__							:   Any	)     ->					[Memory, Optional[MemorySummary]]:





      '''simple docstring'''






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


     def    UpperCAmelCase__		(					self							:   str							,		lowerCamelCase__							:   List[Any]							,		lowerCamelCase__							:   Optional[Any]							,		lowerCamelCase__							:   Dict	)     ->					[Memory, Optional[MemorySummary]]:





      '''simple docstring'''






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

      UpperCamelCase__    :			Tuple      				=      self._prepare_train_func(__lowerCAmelCase							,		__lowerCAmelCase							,		__lowerCAmelCase	)
      return self._measure_memory(_train	)


     def    UpperCAmelCase__		(					self							:   Dict							,		lowerCamelCase__							:   int							,		lowerCamelCase__							:   Tuple							,		lowerCamelCase__							:   str	)     ->					Callable[[], None]:





      '''simple docstring'''






      UpperCamelCase__    :			int      				=      self.config_dict[model_name]

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

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

      # encoder-decoder has vocab size saved differently
      UpperCamelCase__    :			Union[str, Any]      				=      config.vocab_size if hasattr(__lowerCAmelCase							,		'''vocab_size'''	) else config.encoder.vocab_size
      UpperCamelCase__    :			Optional[int]      				=      random_input_ids(__lowerCAmelCase							,		__lowerCAmelCase							,		__lowerCAmelCase	)

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

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

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

      return _inference


     def    UpperCAmelCase__		(					self							:   Any							,		lowerCamelCase__							:   Union[str, Any]							,		lowerCamelCase__							:   str							,		lowerCamelCase__							:   Optional[Any]	)     ->					Callable[[], None]:





      '''simple docstring'''






      UpperCamelCase__    :			Dict      				=      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__    :			Optional[Any]      				=      (
          hasattr(__lowerCAmelCase							,		'''architectures'''	)
          and isinstance(config.architectures							,		__lowerCAmelCase	)
          and len(config.architectures	) > 0
      )
      if not self.args.only_pretrain_model and has_model_class_in_config:
       try:
        UpperCamelCase__    :			Tuple      				=      '''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
        UpperCamelCase__    :			Dict      				=      __import__('''transformers'''							,		fromlist=[model_class]	)
        UpperCamelCase__    :			Dict      				=      getattr(__lowerCAmelCase							,		__lowerCAmelCase	)
        UpperCamelCase__    :			Optional[int]      				=      model_cls(__lowerCAmelCase	)
       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__    :			Union[str, Any]      				=      TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__lowerCAmelCase	)

      # encoder-decoder has vocab size saved differently
      UpperCamelCase__    :			int      				=      config.vocab_size if hasattr(__lowerCAmelCase							,		'''vocab_size'''	) else config.encoder.vocab_size
      UpperCamelCase__    :			List[str]      				=      random_input_ids(__lowerCAmelCase							,		__lowerCAmelCase							,		__lowerCAmelCase	)

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

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

      UpperCamelCase__    :			Any      				=      encoder_decoder_train if config.is_encoder_decoder else encoder_train

      return _train


     def    UpperCAmelCase__		(					self							:   Tuple							,		lowerCamelCase__							:   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(__lowerCAmelCase							,		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__    :			Any      				=      timeit.repeat(
            __lowerCAmelCase							,		repeat=self.args.repeat							,		number=10							,		)

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


     def    UpperCAmelCase__		(					self							:   Dict							,		lowerCamelCase__							:   Dict	)     ->					[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.'''	)
         UpperCamelCase__    :			Union[str, Any]      				=      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__    :			Tuple      				=      '''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__    :			Tuple      				=      nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
          UpperCamelCase__    :			str      				=      nvml.nvmlDeviceGetMemoryInfo(__lowerCAmelCase	)
          UpperCamelCase__    :			List[str]      				=      meminfo.used
          UpperCamelCase__    :			Optional[Any]      				=      Memory(__lowerCAmelCase	)
          # 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__    :			Dict      				=      None
         else:
          UpperCamelCase__    :			Dict      				=      measure_peak_memory_cpu(__lowerCAmelCase	)
          UpperCamelCase__    :			Tuple      				=      Memory(__lowerCAmelCase	) if isinstance(__lowerCAmelCase							,		__lowerCAmelCase	) else memory_bytes
        if self.args.trace_memory_line_by_line:
         UpperCamelCase__    :			List[str]      				=      stop_memory_tracing(__lowerCAmelCase	)
         if memory is None:
          UpperCamelCase__    :			Optional[Any]      				=      summary.total
        else:
         UpperCamelCase__    :			Union[str, Any]      				=      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 argparse

import pytorch_lightning as pl
import torch
from torch import nn

from transformers import LongformerForQuestionAnswering, LongformerModel



class 			UpperCAmelCase  (						pl.LightningModule					):





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase )  ->  List[str]:
        super().__init__()
        lowercase__				:    List[str]  =					model
        lowercase__				:    Dict  =					2
        lowercase__				:    Any  =					nn.Linear(self.model.config.hidden_size	,   self.num_labels )


       def _lowerCAmelCase(						self )  ->  str:
        pass


def 			__UpperCamelCase       (   UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase     ):
 # load longformer model from model identifier
 lowercase__				:    Dict  =					LongformerModel.from_pretrained(UpperCAmelCase     )
 lowercase__				:    List[str]  =					LightningModel(UpperCAmelCase     )

 lowercase__				:    List[Any]  =					torch.load(UpperCAmelCase		,     map_location=torch.device('''cpu'''     )     )
 lightning_model.load_state_dict(ckpt['''state_dict''']     )

 # init longformer question answering model
 lowercase__				:    Optional[int]  =					LongformerForQuestionAnswering.from_pretrained(UpperCAmelCase     )

 # transfer weights
 longformer_for_qa.longformer.load_state_dict(lightning_model.model.state_dict()     )
 longformer_for_qa.qa_outputs.load_state_dict(lightning_model.qa_outputs.state_dict()     )
 longformer_for_qa.eval()

 # save model
 longformer_for_qa.save_pretrained(UpperCAmelCase     )

 print(F"""Conversion successful. Model saved under {pytorch_dump_folder_path}"""     )


if __name__ == "__main__":
  __a: List[Any]          =   argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--longformer_model""",
      default=None,
      type=str,
      required=True,
      help="""model identifier of longformer. Should be either `longformer-base-4096` or `longformer-large-4096`.""",
  )
  parser.add_argument(
      """--longformer_question_answering_ckpt_path""",
      default=None,
      type=str,
      required=True,
      help="""Path the official PyTorch Lightning Checkpoint.""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
  )
  __a: Tuple          =   parser.parse_args()
  convert_longformer_qa_checkpoint_to_pytorch(
      args.longformer_model, args.longformer_question_answering_ckpt_path, args.pytorch_dump_folder_path
  )

