




import inspect
import unittest

from transformers import SegformerConfig, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    MODEL_MAPPING,
		    SegformerForImageClassification,
		    SegformerForSemanticSegmentation,
		    SegformerModel,
		)
		from transformers.models.segformer.modeling_segformer import SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import SegformerImageProcessor

class 		__lowerCAmelCase	(	__magic_name__ ):
							def       lowerCamelCase__							(						self						:List[str]				):







								'''simple docstring'''







								a										=							self.config_class(**self.inputs_dict				)
								self.parent.assertTrue(hasattr(__magic_name__		,      """hidden_sizes"""				)				)
								self.parent.assertTrue(hasattr(__magic_name__		,      """num_attention_heads"""				)				)
								self.parent.assertTrue(hasattr(__magic_name__		,      """num_encoder_blocks"""				)				)







class 		__lowerCAmelCase	:
							def __init__(						self						:Union[str, Any]		,      __magic_name__						:Tuple		,      __magic_name__						:Union[str, Any]=13		,      __magic_name__						:List[Any]=64		,      __magic_name__						:Any=3		,      __magic_name__						:Optional[int]=4		,      __magic_name__						:List[Any]=[2, 2, 2, 2]		,      __magic_name__						:Tuple=[8, 4, 2, 1]		,      __magic_name__						:str=[16, 32, 64, 128]		,      __magic_name__						:Union[str, Any]=[1, 4, 8, 16]		,      __magic_name__						:Any=[1, 2, 4, 8]		,      __magic_name__						:Any=True		,      __magic_name__						:Union[str, Any]=True		,      __magic_name__						:List[Any]="gelu"		,      __magic_name__						:Dict=0.1		,      __magic_name__						:int=0.1		,      __magic_name__						:Union[str, Any]=0.02		,      __magic_name__						:Optional[Any]=3		,      __magic_name__						:List[Any]=None		,      ):







								'''simple docstring'''







								a										=							parent
								a										=							batch_size
								a										=							image_size
								a										=							num_channels
								a										=							num_encoder_blocks
								a										=							sr_ratios
								a										=							depths
								a										=							hidden_sizes
								a										=							downsampling_rates
								a										=							num_attention_heads
								a										=							is_training
								a										=							use_labels
								a										=							hidden_act
								a										=							hidden_dropout_prob
								a										=							attention_probs_dropout_prob
								a										=							initializer_range
								a										=							num_labels
								a										=							scope
							def       lowerCamelCase__							(						self						:Optional[Any]				):







								'''simple docstring'''







								a										=							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

								a										=							None
								if self.use_labels:
									a										=							ids_tensor([self.batch_size, self.image_size, self.image_size]		,      self.num_labels				)

								a										=							self.get_config()
								return config, pixel_values, labels
							def       lowerCamelCase__							(						self						:List[str]				):







								'''simple docstring'''







								return SegformerConfig(
								    image_size=self.image_size		,      num_channels=self.num_channels		,      num_encoder_blocks=self.num_encoder_blocks		,      depths=self.depths		,      hidden_sizes=self.hidden_sizes		,      num_attention_heads=self.num_attention_heads		,      hidden_act=self.hidden_act		,      hidden_dropout_prob=self.hidden_dropout_prob		,      attention_probs_dropout_prob=self.attention_probs_dropout_prob		,      initializer_range=self.initializer_range		,      )
							def       lowerCamelCase__							(						self						:List[str]		,      __magic_name__						:Union[str, Any]		,      __magic_name__						:Tuple		,      __magic_name__						:str				):







								'''simple docstring'''







								a										=							SegformerModel(config=__magic_name__				)
								model.to(__magic_name__				)
								model.eval()
								a										=							model(__magic_name__				)
								a										=							a										=							self.image_size // (self.downsampling_rates[-1] * 2)
								self.parent.assertEqual(
								    result.last_hidden_state.shape		,      (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)				)
							def       lowerCamelCase__							(						self						:List[str]		,      __magic_name__						:Dict		,      __magic_name__						:Dict		,      __magic_name__						:List[Any]				):







								'''simple docstring'''







								a										=							self.num_labels
								a										=							SegformerForSemanticSegmentation(__magic_name__				)
								model.to(__magic_name__				)
								model.eval()
								a										=							model(__magic_name__				)
								self.parent.assertEqual(
								    result.logits.shape		,      (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)				)
								a										=							model(__magic_name__		,      labels=__magic_name__				)
								self.parent.assertEqual(
								    result.logits.shape		,      (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)				)
								self.parent.assertGreater(result.loss		,      0.0				)
							def       lowerCamelCase__							(						self						:Tuple		,      __magic_name__						:List[Any]		,      __magic_name__						:List[Any]		,      __magic_name__						:Optional[int]				):







								'''simple docstring'''







								a										=							1
								a										=							SegformerForSemanticSegmentation(config=__magic_name__				)
								model.to(__magic_name__				)
								model.eval()
								a										=							torch.randint(0		,      1		,      (self.batch_size, self.image_size, self.image_size)				).to(__magic_name__				)
								a										=							model(__magic_name__		,      labels=__magic_name__				)
								self.parent.assertGreater(result.loss		,      0.0				)





							def       lowerCamelCase__							(						self						:Optional[Any]				):







								'''simple docstring'''







								a										=							self.prepare_config_and_inputs()
								a				, a				, a										=							config_and_inputs
								a										=							{"""pixel_values""": pixel_values}
								return config, inputs_dict







@require_torch
class 		__lowerCAmelCase	(	__magic_name__      ,       __magic_name__      ,       unittest.TestCase ):
							UpperCamelCase__					     =  (
							    (
							        SegformerModel,
							        SegformerForSemanticSegmentation,
							        SegformerForImageClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							UpperCamelCase__					     =  (
							    {
							        '''feature-extraction''': SegformerModel,
							        '''image-classification''': SegformerForImageClassification,
							        '''image-segmentation''': SegformerForSemanticSegmentation,
							    }
							    if is_torch_available()
							    else {}
							)

							UpperCamelCase__					     =  True
							UpperCamelCase__					     =  False
							UpperCamelCase__					     =  False
							UpperCamelCase__					     =  False
							def       lowerCamelCase__							(						self						:Any				):







								'''simple docstring'''







								a										=							SegformerModelTester(self				)
								a										=							SegformerConfigTester(self		,      config_class=__magic_name__				)
							def       lowerCamelCase__							(						self						:Dict				):







								'''simple docstring'''







								self.config_tester.run_common_tests()
							def       lowerCamelCase__							(						self						:Any				):







								'''simple docstring'''







								a										=							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__magic_name__				)
							def       lowerCamelCase__							(						self						:List[Any]				):







								'''simple docstring'''







								a										=							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_binary_image_segmentation(*__magic_name__				)
							def       lowerCamelCase__							(						self						:Union[str, Any]				):







								'''simple docstring'''







								a										=							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_segmentation(*__magic_name__				)
							@unittest.skip("""SegFormer does not use inputs_embeds"""				)
							def       lowerCamelCase__							(						self						:Any				):







								'''simple docstring'''







								pass
							@unittest.skip("""SegFormer does not have get_input_embeddings method and get_output_embeddings methods"""				)
							def       lowerCamelCase__							(						self						:int				):







								'''simple docstring'''







								pass
							def       lowerCamelCase__							(						self						:List[Any]				):







								'''simple docstring'''







								a				, a										=							self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
									a										=							model_class(__magic_name__				)
									a										=							inspect.signature(model.forward				)
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									a										=							[*signature.parameters.keys()]

									a										=							["""pixel_values"""]
									self.assertListEqual(arg_names[:1]		,      __magic_name__				)
							def       lowerCamelCase__							(						self						:int				):







								'''simple docstring'''







								a				, a										=							self.model_tester.prepare_config_and_inputs_for_common()
								a										=							True

								for model_class in self.all_model_classes:
									a										=							True
									a										=							False
									a										=							True
									a										=							model_class(__magic_name__				)
									model.to(__magic_name__				)
									model.eval()
									with torch.no_grad():
										a										=							model(**self._prepare_for_class(__magic_name__		,      __magic_name__				)				)
									a										=							outputs.attentions

									a										=							sum(self.model_tester.depths				)
									self.assertEqual(len(__magic_name__				)		,      __magic_name__				)

									# check that output_attentions also work using config
									del inputs_dict["output_attentions"]
									a										=							True
									a										=							model_class(__magic_name__				)
									model.to(__magic_name__				)
									model.eval()
									with torch.no_grad():
										a										=							model(**self._prepare_for_class(__magic_name__		,      __magic_name__				)				)
									a										=							outputs.attentions

									self.assertEqual(len(__magic_name__				)		,      __magic_name__				)

									# verify the first attentions (first block, first layer)
									a										=							(self.model_tester.image_size // 4) ** 2
									a										=							(self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
									self.assertListEqual(
									    list(attentions[0].shape[-3:]				)		,      [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]		,      )

									# verify the last attentions (last block, last layer)
									a										=							(self.model_tester.image_size // 32) ** 2
									a										=							(self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
									self.assertListEqual(
									    list(attentions[-1].shape[-3:]				)		,      [self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]		,      )
									a										=							len(__magic_name__				)

									# Check attention is always last and order is fine
									a										=							True
									a										=							True
									a										=							model_class(__magic_name__				)
									model.to(__magic_name__				)
									model.eval()
									with torch.no_grad():
										a										=							model(**self._prepare_for_class(__magic_name__		,      __magic_name__				)				)

									self.assertEqual(out_len + 1		,      len(__magic_name__				)				)

									a										=							outputs.attentions

									self.assertEqual(len(__magic_name__				)		,      __magic_name__				)
									# verify the first attentions (first block, first layer)
									a										=							(self.model_tester.image_size // 4) ** 2
									a										=							(self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
									self.assertListEqual(
									    list(self_attentions[0].shape[-3:]				)		,      [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]		,      )
							def       lowerCamelCase__							(						self						:Any				):







								'''simple docstring'''







								def check_hidden_states_output(__magic_name__						:Tuple		,      __magic_name__						:Optional[Any]		,      __magic_name__						:Union[str, Any]				):
									a										=							model_class(__magic_name__				)
									model.to(__magic_name__				)
									model.eval()

									with torch.no_grad():
										a										=							model(**self._prepare_for_class(__magic_name__		,      __magic_name__				)				)

									a										=							outputs.hidden_states

									a										=							self.model_tester.num_encoder_blocks
									self.assertEqual(len(__magic_name__				)		,      __magic_name__				)

									# verify the first hidden states (first block)
									self.assertListEqual(
									    list(hidden_states[0].shape[-3:]				)		,      [
									        self.model_tester.hidden_sizes[0],
									        self.model_tester.image_size // 4,
									        self.model_tester.image_size // 4,
									    ]		,      )

								a				, a										=							self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
									a										=							True
									check_hidden_states_output(__magic_name__		,      __magic_name__		,      __magic_name__				)

									# check that output_hidden_states also work using config
									del inputs_dict["output_hidden_states"]
									a										=							True

									check_hidden_states_output(__magic_name__		,      __magic_name__		,      __magic_name__				)
							def       lowerCamelCase__							(						self						:Tuple				):







								'''simple docstring'''







								if not self.model_tester.is_training:
									return

								a				, a										=							self.model_tester.prepare_config_and_inputs_for_common()
								a										=							True

								for model_class in self.all_model_classes:
									if model_class in get_values(__magic_name__				):
										continue

									a										=							model_class(__magic_name__				)
									model.to(__magic_name__				)
									model.train()
									a										=							self._prepare_for_class(__magic_name__		,      __magic_name__		,      return_labels=__magic_name__				)
									a										=							model(**__magic_name__				).loss
									loss.backward()
							@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""				)
							def       lowerCamelCase__							(						self						:Dict				):







								'''simple docstring'''







								pass





							@slow
							def       lowerCamelCase__							(						self						:int				):







								'''simple docstring'''







								for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									a										=							SegformerModel.from_pretrained(__magic_name__				)
									self.assertIsNotNone(__magic_name__				)







def 		__A		(	)      ->       str:
	a										=							Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""    )
	return image



@require_torch
class 		__lowerCAmelCase	(	unittest.TestCase ):
							@slow
							def       lowerCamelCase__							(						self						:Union[str, Any]				):







								'''simple docstring'''







								a										=							SegformerImageProcessor(
								    image_scale=(512, 512)		,      keep_ratio=__magic_name__		,      align=__magic_name__		,      do_random_crop=__magic_name__				)
								a										=							SegformerForSemanticSegmentation.from_pretrained("""nvidia/segformer-b0-finetuned-ade-512-512"""				).to(
								    __magic_name__				)

								a										=							prepare_img()
								a										=							image_processor(images=__magic_name__		,      return_tensors="""pt"""				)
								a										=							encoded_inputs.pixel_values.to(__magic_name__				)

								with torch.no_grad():
									a										=							model(__magic_name__				)

								a										=							torch.Size((1, model.config.num_labels, 128, 128)				)
								self.assertEqual(outputs.logits.shape		,      __magic_name__				)

								a										=							torch.tensor(
								    [
								        [[-4.6310, -5.5232, -6.2356], [-5.1921, -6.1444, -6.5996], [-5.4424, -6.2790, -6.7574]],
								        [[-12.1391, -13.3122, -13.9554], [-12.8732, -13.9352, -14.3563], [-12.9438, -13.8226, -14.2513]],
								        [[-12.5134, -13.4686, -14.4915], [-12.8669, -14.4343, -14.7758], [-13.2523, -14.5819, -15.0694]],
								    ]				).to(__magic_name__				)
								self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]		,      __magic_name__		,      atol=1E-4				)				)
							@slow
							def       lowerCamelCase__							(						self						:Any				):







								'''simple docstring'''







								a										=							SegformerImageProcessor(
								    image_scale=(512, 512)		,      keep_ratio=__magic_name__		,      align=__magic_name__		,      do_random_crop=__magic_name__				)
								a										=							SegformerForSemanticSegmentation.from_pretrained(
								    """nvidia/segformer-b1-finetuned-cityscapes-1024-1024"""				).to(__magic_name__				)

								a										=							prepare_img()
								a										=							image_processor(images=__magic_name__		,      return_tensors="""pt"""				)
								a										=							encoded_inputs.pixel_values.to(__magic_name__				)

								with torch.no_grad():
									a										=							model(__magic_name__				)

								a										=							torch.Size((1, model.config.num_labels, 128, 128)				)
								self.assertEqual(outputs.logits.shape		,      __magic_name__				)

								a										=							torch.tensor(
								    [
								        [[-13.5748, -13.9111, -12.6500], [-14.3500, -15.3683, -14.2328], [-14.7532, -16.0424, -15.6087]],
								        [[-17.1651, -15.8725, -12.9653], [-17.2580, -17.3718, -14.8223], [-16.6058, -16.8783, -16.7452]],
								        [[-3.6456, -3.0209, -1.4203], [-3.0797, -3.1959, -2.0000], [-1.8757, -1.9217, -1.6997]],
								    ]				).to(__magic_name__				)
								self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]		,      __magic_name__		,      atol=1E-1				)				)





							@slow
							def       lowerCamelCase__							(						self						:List[Any]				):







								'''simple docstring'''







								a										=							SegformerImageProcessor(
								    image_scale=(512, 512)		,      keep_ratio=__magic_name__		,      align=__magic_name__		,      do_random_crop=__magic_name__				)
								a										=							SegformerForSemanticSegmentation.from_pretrained("""nvidia/segformer-b0-finetuned-ade-512-512"""				).to(
								    __magic_name__				)

								a										=							prepare_img()
								a										=							image_processor(images=__magic_name__		,      return_tensors="""pt"""				)
								a										=							encoded_inputs.pixel_values.to(__magic_name__				)

								with torch.no_grad():
									a										=							model(__magic_name__				)

								a										=							outputs.logits.detach().cpu()

								a										=							image_processor.post_process_semantic_segmentation(outputs=__magic_name__		,      target_sizes=[(500, 300)]				)
								a										=							torch.Size((500, 300)				)
								self.assertEqual(segmentation[0].shape		,      __magic_name__				)

								a										=							image_processor.post_process_semantic_segmentation(outputs=__magic_name__				)
								a										=							torch.Size((128, 128)				)
								self.assertEqual(segmentation[0].shape		,      __magic_name__				)







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
		)

