


"""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							, _a      ):
				super().__init__()
				__a           =			model
				__a           =			2
				__a           =			nn.Linear(self.model.config.hidden_size							, self.num_labels      )





		def 			__UpperCAmelCase       (  self      ):
				pass







def        lowercase							(		lowerCAmelCase__ :     str    ,					lowerCAmelCase__ :     str    ,					lowerCAmelCase__ :     str			) ->      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__":
				lowercase_			     =  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."
				)
				lowercase_			     =  parser.parse_args()
				convert_longformer_qa_checkpoint_to_pytorch(
				    args.longformer_model, args.longformer_question_answering_ckpt_path, args.pytorch_dump_folder_path
				)










"""simple docstring"""




lowercase_			     =  "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"

def        lowercase							(		lowerCAmelCase__ :     bytes			) ->      bytes:
		# Make sure the supplied data is a bytes-like object
		if not isinstance(lowerCAmelCase__    ,					lowerCAmelCase__			):
				__a           =			f'''a bytes-like object is required, not \'{data.__class__.__name__}\''''
				raise TypeError(lowerCAmelCase__			)

		__a           =			''''''.join(bin(lowerCAmelCase__			)[2:].zfill(8			) for byte in data			)

		__a           =			len(lowerCAmelCase__			) % 6 != 0

		if padding_needed:
				# The padding that will be added later
				__a           =			b'''=''' * ((6 - len(lowerCAmelCase__			) % 6) // 2)

				# Append binary_stream with arbitrary binary digits (0's by default) to make its
				# length a multiple of 6.
				binary_stream += "0" * (6 - len(lowerCAmelCase__			) % 6)
		else:
				__a           =			b''''''

		# Encode every 6 binary digits to their corresponding Base64 character
		return (
		    "".join(
		        B64_CHARSET[int(binary_stream[index : index + 6]    ,					2			)]
		        for index in range(0    ,					len(lowerCAmelCase__			)    ,					6			)			).encode()
		    + padding
		)

def        lowercase							(		lowerCAmelCase__ :     str			) ->      bytes:
		# Make sure encoded_data is either a string or a bytes-like object
		if not isinstance(lowerCAmelCase__    ,					lowerCAmelCase__			) and not isinstance(lowerCAmelCase__    ,					lowerCAmelCase__			):
				__a           =			(
				    '''argument should be a bytes-like object or ASCII string, '''
				    f'''not \'{encoded_data.__class__.__name__}\''''
				)
				raise TypeError(lowerCAmelCase__			)

		# In case encoded_data is a bytes-like object, make sure it contains only
		# ASCII characters so we convert it to a string object
		if isinstance(lowerCAmelCase__    ,					lowerCAmelCase__			):
				try:
						__a           =			encoded_data.decode('''utf-8'''			)
				except UnicodeDecodeError:
						raise ValueError('''base64 encoded data should only contain ASCII characters'''			)

		__a           =			encoded_data.count('''='''			)

		# Check if the encoded string contains non base64 characters
		if padding:
				assert all(
				    char in B64_CHARSET for char in encoded_data[:-padding]			), "Invalid base64 character(s) found."
		else:
				assert all(
				    char in B64_CHARSET for char in encoded_data			), "Invalid base64 character(s) found."

		# Check the padding
		assert len(lowerCAmelCase__			) % 4 == 0 and padding < 3, "Incorrect padding"

		if padding:
				# Remove padding if there is one
				__a           =			encoded_data[:-padding]

				__a           =			''''''.join(
				    bin(B64_CHARSET.index(lowerCAmelCase__			)			)[2:].zfill(6			) for char in encoded_data			)[: -padding * 2]
		else:
				__a           =			''''''.join(
				    bin(B64_CHARSET.index(lowerCAmelCase__			)			)[2:].zfill(6			) for char in encoded_data			)

		__a           =			[
		    int(binary_stream[index : index + 8]    ,					2			)
		    for index in range(0    ,					len(lowerCAmelCase__			)    ,					8			)
		]

		return bytes(lowerCAmelCase__			)


if __name__ == "__main__":
				import doctest

				doctest.testmod()






