





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      :					List[Any]      ,				UpperCamelCase      :					Dict							):


    '''simple docstring'''





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







  def UpperCamelCase_ (      self      :					Optional[int]							):


    '''simple docstring'''





    pass




def     lowerCamelCase_	(      lowerCAmelCase:	str    ,							lowerCAmelCase:	str    ,							lowerCAmelCase:	str				)->		Dict:
  # load longformer model from model identifier
  _snake_case   :				Tuple														=	LongformerModel.from_pretrained(lowerCAmelCase				)
  _snake_case   :				Tuple														=	LightningModel(lowerCAmelCase				)

  _snake_case   :				int														=	torch.load(lowerCAmelCase    ,							map_location=torch.device('cpu'				)				)
  lightning_model.load_state_dict(ckpt['state_dict']				)

  # init longformer question answering model
  _snake_case   :				Dict														=	LongformerForQuestionAnswering.from_pretrained(lowerCAmelCase				)

  # 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(lowerCAmelCase				)

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


if __name__ == "__main__":
      lowerCAmelCase_						   =    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."""
      )
      lowerCAmelCase_						   =    parser.parse_args()
      convert_longformer_qa_checkpoint_to_pytorch(
          args.longformer_model, args.longformer_question_answering_ckpt_path, args.pytorch_dump_folder_path
      )











def     lowerCamelCase_	(      lowerCAmelCase:	bytes				)->		str:
  return "".join([hex(lowerCAmelCase				)[2:].zfill(2				).upper() for byte in list(lowerCAmelCase				)]				)


def     lowerCamelCase_	(      lowerCAmelCase:	str				)->		bytes:
  # Check data validity, following RFC3548
  # https://www.ietf.org/rfc/rfc3548.txt
  if (len(lowerCAmelCase				) % 2) != 0:
    raise ValueError(
        'Base16 encoded data is invalid:\nData does not have an even number of hex digits.'				)
  # Check the character set - the standard base16 alphabet
  # is uppercase according to RFC3548 section 6
  if not set(lowerCAmelCase				) <= set('0123456789ABCDEF'				):
    raise ValueError(
        'Base16 encoded data is invalid:\nData is not uppercase hex or it contains invalid characters.'				)
  # For every two hexadecimal digits (= a byte), turn it into an integer.
  # Then, string the result together into bytes, and return it.
  return bytes(int(data[i] + data[i + 1]    ,							16				) for i in range(0    ,							len(lowerCAmelCase				)    ,							2				)				)


if __name__ == "__main__":
      import doctest

      doctest.testmod()




