





"""simple docstring"""





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   ,							__a		) ->   List[str]:
								super().__init__()
								UpperCamelCase   =	model
								UpperCamelCase   =	2
								UpperCamelCase   =	nn.Linear(self.model.config.hidden_size   ,							self.num_labels		)




						def     snake_case_   (self		) ->   Tuple:
								pass



def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	LongformerModel.from_pretrained(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	LightningModel(_SCREAMING_SNAKE_CASE    )

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

		# init longformer question answering model
		UpperCamelCase   =	LongformerForQuestionAnswering.from_pretrained(_SCREAMING_SNAKE_CASE    )

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

		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
						)













"""simple docstring"""





import pytest

from datasets.parallel import ParallelBackendConfig, parallel_backend
from datasets.utils.py_utils import map_nested

from .utils import require_dill_gt_0_3_2, require_joblibspark, require_not_windows


def       a__					(					_SCREAMING_SNAKE_CASE    ):  # picklable for multiprocessing




		"""simple docstring"""


		return i + 1


@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
def       a__					(					):




		"""simple docstring"""


		with parallel_backend("spark"    ):
				assert ParallelBackendConfig.backend_name == "spark"

		UpperCamelCase   =	[1, 2, 3]
		with pytest.raises(_SCREAMING_SNAKE_CASE    ):
				with parallel_backend("unsupported backend"    ):
						map_nested(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	num_proc=2    )

		with pytest.raises(_SCREAMING_SNAKE_CASE    ):
				with parallel_backend("unsupported backend"    ):
						map_nested(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	num_proc=-1    )






@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
@pytest.mark.parametrize("num_proc"  ,	[2, -1]    )
def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	[1, 2]
		UpperCamelCase   =	{"a": 1, "b": 2}
		UpperCamelCase   =	{"a": [1, 2], "b": [3, 4]}
		UpperCamelCase   =	{"a": {"1": 1}, "b": 2}
		UpperCamelCase   =	{"a": 1, "b": 2, "c": 3, "d": 4}
		UpperCamelCase   =	[2, 3]
		UpperCamelCase   =	{"a": 2, "b": 3}
		UpperCamelCase   =	{"a": [2, 3], "b": [4, 5]}
		UpperCamelCase   =	{"a": {"1": 2}, "b": 3}
		UpperCamelCase   =	{"a": 2, "b": 3, "c": 4, "d": 5}

		with parallel_backend("spark"    ):
				assert map_nested(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	num_proc=_SCREAMING_SNAKE_CASE    ) == expected_map_nested_sa
				assert map_nested(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	num_proc=_SCREAMING_SNAKE_CASE    ) == expected_map_nested_sa
				assert map_nested(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	num_proc=_SCREAMING_SNAKE_CASE    ) == expected_map_nested_sa
				assert map_nested(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	num_proc=_SCREAMING_SNAKE_CASE    ) == expected_map_nested_sa
				assert map_nested(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	num_proc=_SCREAMING_SNAKE_CASE    ) == expected_map_nested_sa






