



import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
						import tensorflow as tf

						from transformers import TensorFlowBenchmark, TensorFlowBenchmarkArguments



@require_tf
class   UpperCAmelCase_    (       unittest.TestCase):







							def    snake_case__      (   self,			__a):

														'''simple docstring'''

														for model_result in results.values():
																					for batch_size, sequence_length in zip(model_result["bs"],			model_result["ss"]):
																												_lowerCAmelCase     :							Tuple    =    model_result["result"][batch_size][sequence_length]
																												self.assertIsNotNone(__a)







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Dict    =    "sshleifer/tiny-gpt2"
														_lowerCAmelCase     :							Tuple    =    TensorFlowBenchmarkArguments(
														    models=[MODEL_ID],			training=__a,			inference=__a,			sequence_lengths=[8],			batch_sizes=[1],			eager_mode=__a,			multi_process=__a,			)
														_lowerCAmelCase     :							str    =    TensorFlowBenchmark(__a)
														_lowerCAmelCase     :							List[str]    =    benchmark.run()
														self.check_results_dict_not_empty(results.time_inference_result)
														self.check_results_dict_not_empty(results.memory_inference_result)







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Tuple    =    "sgugger/tiny-distilbert-classification"
														_lowerCAmelCase     :							Union[str, Any]    =    TensorFlowBenchmarkArguments(
														    models=[MODEL_ID],			training=__a,			inference=__a,			sequence_lengths=[8],			batch_sizes=[1],			multi_process=__a,			only_pretrain_model=__a,			)
														_lowerCAmelCase     :							Tuple    =    TensorFlowBenchmark(__a)
														_lowerCAmelCase     :							Tuple    =    benchmark.run()
														self.check_results_dict_not_empty(results.time_inference_result)
														self.check_results_dict_not_empty(results.memory_inference_result)







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Union[str, Any]    =    "sshleifer/tiny-gpt2"
														_lowerCAmelCase     :							Optional[Any]    =    TensorFlowBenchmarkArguments(
														    models=[MODEL_ID],			training=__a,			inference=__a,			sequence_lengths=[8],			batch_sizes=[1],			multi_process=__a,			)
														_lowerCAmelCase     :							List[Any]    =    TensorFlowBenchmark(__a)
														_lowerCAmelCase     :							Optional[int]    =    benchmark.run()
														self.check_results_dict_not_empty(results.time_inference_result)
														self.check_results_dict_not_empty(results.memory_inference_result)







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							str    =    "sshleifer/tiny-gpt2"
														_lowerCAmelCase     :							Dict    =    AutoConfig.from_pretrained(__a)
														_lowerCAmelCase     :							Tuple    =    TensorFlowBenchmarkArguments(
														    models=[MODEL_ID],			training=__a,			inference=__a,			sequence_lengths=[8],			batch_sizes=[1],			eager_mode=__a,			multi_process=__a,			)
														_lowerCAmelCase     :							Any    =    TensorFlowBenchmark(__a,			[config])
														_lowerCAmelCase     :							Optional[Any]    =    benchmark.run()
														self.check_results_dict_not_empty(results.time_inference_result)
														self.check_results_dict_not_empty(results.memory_inference_result)







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Union[str, Any]    =    "sshleifer/tiny-gpt2"
														_lowerCAmelCase     :							int    =    AutoConfig.from_pretrained(__a)
														_lowerCAmelCase     :							List[Any]    =    TensorFlowBenchmarkArguments(
														    models=[MODEL_ID],			training=__a,			inference=__a,			sequence_lengths=[8],			batch_sizes=[1],			multi_process=__a,			)
														_lowerCAmelCase     :							Any    =    TensorFlowBenchmark(__a,			[config])
														_lowerCAmelCase     :							Optional[Any]    =    benchmark.run()
														self.check_results_dict_not_empty(results.time_inference_result)
														self.check_results_dict_not_empty(results.memory_inference_result)







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Optional[Any]    =    "sshleifer/tiny-gpt2"
														_lowerCAmelCase     :							List[Any]    =    TensorFlowBenchmarkArguments(
														    models=[MODEL_ID],			training=__a,			inference=__a,			sequence_lengths=[8],			batch_sizes=[1],			multi_process=__a,			)
														_lowerCAmelCase     :							List[Any]    =    TensorFlowBenchmark(__a)
														_lowerCAmelCase     :							List[str]    =    benchmark.run()
														self.check_results_dict_not_empty(results.time_train_result)
														self.check_results_dict_not_empty(results.memory_train_result)







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Dict    =    "sshleifer/tiny-gpt2"
														_lowerCAmelCase     :							Optional[Any]    =    AutoConfig.from_pretrained(__a)
														_lowerCAmelCase     :							Dict    =    TensorFlowBenchmarkArguments(
														    models=[MODEL_ID],			training=__a,			inference=__a,			sequence_lengths=[8],			batch_sizes=[1],			multi_process=__a,			)
														_lowerCAmelCase     :							str    =    TensorFlowBenchmark(__a,			[config])
														_lowerCAmelCase     :							int    =    benchmark.run()
														self.check_results_dict_not_empty(results.time_train_result)
														self.check_results_dict_not_empty(results.memory_train_result)







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Tuple    =    "patrickvonplaten/t5-tiny-random"
														_lowerCAmelCase     :							List[str]    =    AutoConfig.from_pretrained(__a)
														_lowerCAmelCase     :							Tuple    =    TensorFlowBenchmarkArguments(
														    models=[MODEL_ID],			training=__a,			inference=__a,			sequence_lengths=[8],			batch_sizes=[1],			multi_process=__a,			)
														_lowerCAmelCase     :							Optional[Any]    =    TensorFlowBenchmark(__a,			configs=[config])
														_lowerCAmelCase     :							str    =    benchmark.run()
														self.check_results_dict_not_empty(results.time_inference_result)
														self.check_results_dict_not_empty(results.memory_inference_result)







							@unittest.skipIf(is_tf_available() and len(tf.config.list_physical_devices("GPU")) == 0,			"Cannot do xla on CPU.")
							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							str    =    "sshleifer/tiny-gpt2"
														_lowerCAmelCase     :							str    =    TensorFlowBenchmarkArguments(
														    models=[MODEL_ID],			training=__a,			inference=__a,			sequence_lengths=[8],			batch_sizes=[1],			use_xla=__a,			multi_process=__a,			)
														_lowerCAmelCase     :							Tuple    =    TensorFlowBenchmark(__a)
														_lowerCAmelCase     :							int    =    benchmark.run()
														self.check_results_dict_not_empty(results.time_inference_result)
														self.check_results_dict_not_empty(results.memory_inference_result)







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							List[Any]    =    "sshleifer/tiny-gpt2"
														with tempfile.TemporaryDirectory() as tmp_dir:
																					_lowerCAmelCase     :							Optional[Any]    =    TensorFlowBenchmarkArguments(
																					    models=[MODEL_ID],			inference=__a,			save_to_csv=__a,			sequence_lengths=[8],			batch_sizes=[1],			inference_time_csv_file=os.path.join(__a,			"inf_time.csv"),			inference_memory_csv_file=os.path.join(__a,			"inf_mem.csv"),			env_info_csv_file=os.path.join(__a,			"env.csv"),			multi_process=__a,			)
																					_lowerCAmelCase     :							List[str]    =    TensorFlowBenchmark(__a)
																					benchmark.run()
																					self.assertTrue(Path(os.path.join(__a,			"inf_time.csv")).exists())
																					self.assertTrue(Path(os.path.join(__a,			"inf_mem.csv")).exists())
																					self.assertTrue(Path(os.path.join(__a,			"env.csv")).exists())







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Dict    =    "sshleifer/tiny-gpt2"

														def _check_summary_is_not_empty(__a):
																					self.assertTrue(hasattr(__a,			"sequential"))
																					self.assertTrue(hasattr(__a,			"cumulative"))
																					self.assertTrue(hasattr(__a,			"current"))
																					self.assertTrue(hasattr(__a,			"total"))

														with tempfile.TemporaryDirectory() as tmp_dir:
																					_lowerCAmelCase     :							str    =    TensorFlowBenchmarkArguments(
																					    models=[MODEL_ID],			inference=__a,			sequence_lengths=[8],			batch_sizes=[1],			log_filename=os.path.join(__a,			"log.txt"),			log_print=__a,			trace_memory_line_by_line=__a,			eager_mode=__a,			multi_process=__a,			)
																					_lowerCAmelCase     :							List[Any]    =    TensorFlowBenchmark(__a)
																					_lowerCAmelCase     :							Tuple    =    benchmark.run()
																					_check_summary_is_not_empty(result.inference_summary)
																					self.assertTrue(Path(os.path.join(__a,			"log.txt")).exists())



"""simple docstring"""





import argparse

import pytorch_lightning as pl
import torch
from torch import nn

from transformers import LongformerForQuestionAnswering, LongformerModel




class 					_lowerCAmelCase  (	pl.LightningModule   ):







 """simple docstring"""







 def __init__(			self							:			Optional[Any],					UpperCAmelCase__							:			str	):
    super().__init__()
    __lowercase								=						model
    __lowercase								=						2
    __lowercase								=						nn.Linear(self.model.config.hidden_size,					self.num_labels	)



 def 					_lowercase (			self							:			Optional[int]	):
    pass







def   _A  (  UpperCamelCase_     :       str,			UpperCamelCase_     :       str,			UpperCamelCase_     :       str)	->						str:
   '''simple docstring'''


   __lowercase								=						LongformerModel.from_pretrained(UpperCamelCase_)
   __lowercase								=						LightningModel(UpperCamelCase_)

   __lowercase								=						torch.load(UpperCamelCase_,			map_location=torch.device("cpu"))
   lightning_model.load_state_dict(ckpt["state_dict"])

   # init longformer question answering model
   __lowercase								=						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       =							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       =							parser.parse_args()
   convert_longformer_qa_checkpoint_to_pytorch(
       args.longformer_model, args.longformer_question_answering_ckpt_path, args.pytorch_dump_folder_path
   )




