





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


UpperCAmelCase	:	Tuple		     =       {
    """configuration_falcon""": ["""FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP""", """FalconConfig"""],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCAmelCase	:	List[Any]		     =       [
         """FALCON_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """FalconForCausalLM""",
         """FalconModel""",
         """FalconPreTrainedModel""",
         """FalconForSequenceClassification""",
         """FalconForTokenClassification""",
         """FalconForQuestionAnswering""",
     ]


if TYPE_CHECKING:
     from .configuration_falcon import FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP, FalconConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_falcon import (
              FALCON_PRETRAINED_MODEL_ARCHIVE_LIST,
              FalconForCausalLM,
              FalconForQuestionAnswering,
              FalconForSequenceClassification,
              FalconForTokenClassification,
              FalconModel,
              FalconPreTrainedModel,
          )


else:
     import sys

     UpperCAmelCase	:	List[str]		     =       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)












import argparse

import pytorch_lightning as pl
import torch
from torch import nn

from transformers import LongformerForQuestionAnswering, LongformerModel

class 		__lowerCAmelCase	(	pl.LightningModule ):
							def __init__(						self						:Union[str, Any]		,      __magic_name__						:Optional[int]				):







								'''simple docstring'''







								super().__init__()
								a										=							model
								a										=							2
								a										=							nn.Linear(self.model.config.hidden_size		,      self.num_labels				)





							def       lowerCamelCase__							(						self						:Dict				):







								'''simple docstring'''







								pass







def 		__A		(	__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase    )      ->       Optional[int]:
	# load longformer model from model identifier
	a										=							LongformerModel.from_pretrained(__lowerCamelCase    )
	a										=							LightningModel(__lowerCamelCase    )

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

	# init longformer question answering model
	a										=							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__":
		__UpperCamelCase					:       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."
		)
		__UpperCamelCase					:       List[Any]    					=      parser.parse_args()
		convert_longformer_qa_checkpoint_to_pytorch(
		    args.longformer_model, args.longformer_question_answering_ckpt_path, args.pytorch_dump_folder_path
		)

