'''simple docstring'''

import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
  import jax

  from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
  from transformers.models.bert.modeling_flax_bert import FlaxBertModel
  from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel



@require_flax
class 			UpperCAmelCase  (						unittest.TestCase					):





       '''simple docstring'''



       @slow
       def _lowerCAmelCase(						self )  ->  List[str]:
        for model_name in ["bert-base-cased", "bert-large-uncased"]:
         with self.subTest(__lowerCAmelCase ):
          lowercase__				:    Dict  =					AutoConfig.from_pretrained(__lowerCAmelCase )
          self.assertIsNotNone(__lowerCAmelCase )
          self.assertIsInstance(__lowerCAmelCase	,   __lowerCAmelCase )

          lowercase__				:    Optional[Any]  =					FlaxAutoModel.from_pretrained(__lowerCAmelCase )
          self.assertIsNotNone(__lowerCAmelCase )
          self.assertIsInstance(__lowerCAmelCase	,   __lowerCAmelCase )

       @slow
       def _lowerCAmelCase(						self )  ->  int:
        for model_name in ["roberta-base", "roberta-large"]:
         with self.subTest(__lowerCAmelCase ):
          lowercase__				:    str  =					AutoConfig.from_pretrained(__lowerCAmelCase )
          self.assertIsNotNone(__lowerCAmelCase )
          self.assertIsInstance(__lowerCAmelCase	,   __lowerCAmelCase )

          lowercase__				:    Tuple  =					FlaxAutoModel.from_pretrained(__lowerCAmelCase )
          self.assertIsNotNone(__lowerCAmelCase )
          self.assertIsInstance(__lowerCAmelCase	,   __lowerCAmelCase )

       @slow
       def _lowerCAmelCase(						self )  ->  Dict:
        for model_name in ["bert-base-cased", "bert-large-uncased"]:
         lowercase__				:    Tuple  =					AutoTokenizer.from_pretrained(__lowerCAmelCase )
         lowercase__				:    Any  =					FlaxBertModel.from_pretrained(__lowerCAmelCase )
         lowercase__				:    Dict  =					tokenizer('''Do you support jax jitted function?'''	,   return_tensors=TensorType.JAX )

         @jax.jit
         def eval(**__lowerCAmelCase ):
          return model(**__lowerCAmelCase )

         eval(**__lowerCAmelCase ).block_until_ready()

       @slow
       def _lowerCAmelCase(						self )  ->  Tuple:
        for model_name in ["roberta-base", "roberta-large"]:
         lowercase__				:    List[Any]  =					AutoTokenizer.from_pretrained(__lowerCAmelCase )
         lowercase__				:    Tuple  =					FlaxRobertaModel.from_pretrained(__lowerCAmelCase )
         lowercase__				:    List[str]  =					tokenizer('''Do you support jax jitted function?'''	,   return_tensors=TensorType.JAX )

         @jax.jit
         def eval(**__lowerCAmelCase ):
          return model(**__lowerCAmelCase )

         eval(**__lowerCAmelCase ).block_until_ready()

       def _lowerCAmelCase(						self )  ->  Tuple:
        with self.assertRaisesRegex(
            __lowerCAmelCase	,   '''bert-base is not a local folder and is not a valid model identifier''' ):
         lowercase__				:    Dict  =					FlaxAutoModel.from_pretrained('''bert-base''' )

       def _lowerCAmelCase(						self )  ->  Tuple:
        with self.assertRaisesRegex(
            __lowerCAmelCase	,   r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)''' ):
         lowercase__				:    str  =					FlaxAutoModel.from_pretrained(__lowerCAmelCase	,   revision='''aaaaaa''' )

       def _lowerCAmelCase(						self )  ->  Optional[Any]:
        with self.assertRaisesRegex(
            __lowerCAmelCase	,   '''hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack'''	,   ):
         lowercase__				:    Optional[int]  =					FlaxAutoModel.from_pretrained('''hf-internal-testing/config-no-model''' )


       def _lowerCAmelCase(						self )  ->  int:
        with self.assertRaisesRegex(__lowerCAmelCase	,   '''Use `from_pt=True` to load this model''' ):
         lowercase__				:    Optional[Any]  =					FlaxAutoModel.from_pretrained('''hf-internal-testing/tiny-bert-pt-only''' )


'''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
  )

