



"""simple docstring"""





import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, 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_FOR_IMAGE_CLASSIFICATION_MAPPING,
      MODEL_MAPPING,
      LevitForImageClassification,
      LevitForImageClassificationWithTeacher,
      LevitModel,
  )
  from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import LevitImageProcessor
class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def 	lowercase__							(				self     :    Optional[int]       )					->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.config_class(**self.inputs_dict       )
             self.parent.assertTrue(hasattr(_UpperCAmelCase						,							"hidden_sizes"       )       )
             self.parent.assertTrue(hasattr(_UpperCAmelCase						,							"num_attention_heads"       )       )



class  lowercase__  :







       '''simple docstring'''


       def __init__(				self     :    Optional[int]						,							_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    int=13						,							_UpperCAmelCase     :    Union[str, Any]=64						,							_UpperCAmelCase     :    List[Any]=3						,							_UpperCAmelCase     :    Any=3						,							_UpperCAmelCase     :    Union[str, Any]=2						,							_UpperCAmelCase     :    List[str]=1						,							_UpperCAmelCase     :    Any=16						,							_UpperCAmelCase     :    Dict=[128, 256, 384]						,							_UpperCAmelCase     :    int=[4, 6, 8]						,							_UpperCAmelCase     :    Dict=[2, 3, 4]						,							_UpperCAmelCase     :    Union[str, Any]=[16, 16, 16]						,							_UpperCAmelCase     :    int=0						,							_UpperCAmelCase     :    Tuple=[2, 2, 2]						,							_UpperCAmelCase     :    Tuple=[2, 2, 2]						,							_UpperCAmelCase     :    Dict=0.02						,							_UpperCAmelCase     :    Union[str, Any]=True						,							_UpperCAmelCase     :    Dict=True						,							_UpperCAmelCase     :    str=2						,							)					->		Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				parent
             UpperCAmelCase_     					=				batch_size
             UpperCAmelCase_     					=				image_size
             UpperCAmelCase_     					=				num_channels
             UpperCAmelCase_     					=				kernel_size
             UpperCAmelCase_     					=				stride
             UpperCAmelCase_     					=				padding
             UpperCAmelCase_     					=				hidden_sizes
             UpperCAmelCase_     					=				num_attention_heads
             UpperCAmelCase_     					=				depths
             UpperCAmelCase_     					=				key_dim
             UpperCAmelCase_     					=				drop_path_rate
             UpperCAmelCase_     					=				patch_size
             UpperCAmelCase_     					=				attention_ratio
             UpperCAmelCase_     					=				mlp_ratio
             UpperCAmelCase_     					=				initializer_range
             UpperCAmelCase_     					=				[
                 ["Subsample", key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
                 ["Subsample", key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
             ]
             UpperCAmelCase_     					=				is_training
             UpperCAmelCase_     					=				use_labels
             UpperCAmelCase_     					=				num_labels
             UpperCAmelCase_     					=				initializer_range



       def 	lowercase__							(				self     :    List[Any]       )					->		Union[str, Any]:


             '''simple docstring'''




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

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

             UpperCAmelCase_     					=				self.get_config()
             return config, pixel_values, labels



       def 	lowercase__							(				self     :    List[str]       )					->		Union[str, Any]:


             '''simple docstring'''




             return LevitConfig(
                 image_size=self.image_size						,							num_channels=self.num_channels						,							kernel_size=self.kernel_size						,							stride=self.stride						,							padding=self.padding						,							patch_size=self.patch_size						,							hidden_sizes=self.hidden_sizes						,							num_attention_heads=self.num_attention_heads						,							depths=self.depths						,							key_dim=self.key_dim						,							drop_path_rate=self.drop_path_rate						,							mlp_ratio=self.mlp_ratio						,							attention_ratio=self.attention_ratio						,							initializer_range=self.initializer_range						,							down_ops=self.down_ops						,							)



       def 	lowercase__							(				self     :    List[Any]						,							_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    List[Any]						,							_UpperCAmelCase     :    Optional[int]       )					->		Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				LevitModel(config=_UpperCAmelCase       )
             model.to(_UpperCAmelCase       )
             model.eval()
             UpperCAmelCase_     					=				model(_UpperCAmelCase       )
             UpperCAmelCase_     					=				(self.image_size, self.image_size)
             UpperCAmelCase_					,       UpperCAmelCase_     					=				image_size[0], image_size[1]
             for _ in range(4       ):
                   UpperCAmelCase_     					=				floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1       )
                   UpperCAmelCase_     					=				floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1       )
             self.parent.assertEqual(
                 result.last_hidden_state.shape						,							(self.batch_size, ceil(height / 4       ) * ceil(width / 4       ), self.hidden_sizes[-1])						,							)



       def 	lowercase__							(				self     :    Any						,							_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    List[str]       )					->		List[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.num_labels
             UpperCAmelCase_     					=				LevitForImageClassification(_UpperCAmelCase       )
             model.to(_UpperCAmelCase       )
             model.eval()
             UpperCAmelCase_     					=				model(_UpperCAmelCase						,							labels=_UpperCAmelCase       )
             self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.num_labels)       )







       def 	lowercase__							(				self     :    Union[str, Any]       )					->		List[Any]:


             '''simple docstring'''




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



@require_torch
class  lowercase__  (       SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   unittest.TestCase ):







       '''simple docstring'''


       UpperCamelCase		=	(
           (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
           if is_torch_available()
           else ()
       )
       UpperCamelCase		=	(
           {
               '''feature-extraction''': LevitModel,
               '''image-classification''': (LevitForImageClassification, LevitForImageClassificationWithTeacher),
           }
           if is_torch_available()
           else {}
       )

       UpperCamelCase		=	False
       UpperCamelCase		=	False
       UpperCamelCase		=	False
       UpperCamelCase		=	False
       UpperCamelCase		=	False



       def 	lowercase__							(				self     :    List[str]       )					->		int:


             '''simple docstring'''




             UpperCAmelCase_     					=				LevitModelTester(self       )
             UpperCAmelCase_     					=				ConfigTester(self						,							config_class=_UpperCAmelCase						,							has_text_modality=_UpperCAmelCase						,							hidden_size=37       )



       def 	lowercase__							(				self     :    Tuple       )					->		int:


             '''simple docstring'''




             self.create_and_test_config_common_properties()
             self.config_tester.create_and_test_config_to_json_string()
             self.config_tester.create_and_test_config_to_json_file()
             self.config_tester.create_and_test_config_from_and_save_pretrained()
             self.config_tester.create_and_test_config_with_num_labels()
             self.config_tester.check_config_can_be_init_without_params()
             self.config_tester.check_config_arguments_init()



       def 	lowercase__							(				self     :    List[str]       )					->		List[Any]:


             '''simple docstring'''




             return



       @unittest.skip(reason="Levit does not use inputs_embeds"       )
       def 	lowercase__							(				self     :    Any       )					->		Optional[int]:


             '''simple docstring'''




             pass



       @unittest.skip(reason="Levit does not support input and output embeddings"       )
       def 	lowercase__							(				self     :    Optional[Any]       )					->		Union[str, Any]:


             '''simple docstring'''




             pass



       @unittest.skip(reason="Levit does not output attentions"       )
       def 	lowercase__							(				self     :    Optional[int]       )					->		Union[str, Any]:


             '''simple docstring'''




             pass



       def 	lowercase__							(				self     :    Optional[int]       )					->		Any:


             '''simple docstring'''




             UpperCAmelCase_					,       UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs_for_common()

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

                   UpperCAmelCase_     					=				["pixel_values"]
                   self.assertListEqual(arg_names[:1]						,							_UpperCAmelCase       )



       def 	lowercase__							(				self     :    str       )					->		Optional[Any]:


             '''simple docstring'''




             def check_hidden_states_output(_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    Dict       ):
                   UpperCAmelCase_     					=				model_class(_UpperCAmelCase       )
                   model.to(_UpperCAmelCase       )
                   model.eval()

                   with torch.no_grad():
                         UpperCAmelCase_     					=				model(**self._prepare_for_class(_UpperCAmelCase						,							_UpperCAmelCase       )       )

                   UpperCAmelCase_     					=				outputs.hidden_states

                   UpperCAmelCase_     					=				len(self.model_tester.depths       ) + 1
                   self.assertEqual(len(_UpperCAmelCase       )						,							_UpperCAmelCase       )

                   UpperCAmelCase_     					=				(self.model_tester.image_size, self.model_tester.image_size)
                   UpperCAmelCase_					,       UpperCAmelCase_     					=				image_size[0], image_size[1]
                   for _ in range(4       ):
                         UpperCAmelCase_     					=				floor(
                             (
                                 (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                                 / self.model_tester.stride
                             )
                             + 1       )
                         UpperCAmelCase_     					=				floor(
                             (
                                 (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                                 / self.model_tester.stride
                             )
                             + 1       )
                   # verify the first hidden states (first block)
                   self.assertListEqual(
                       list(hidden_states[0].shape[-2:]       )						,							[
                           height * width,
                           self.model_tester.hidden_sizes[0],
                       ]						,							)

             UpperCAmelCase_					,       UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   UpperCAmelCase_     					=				True
                   check_hidden_states_output(_UpperCAmelCase						,							_UpperCAmelCase						,							_UpperCAmelCase       )

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

                   check_hidden_states_output(_UpperCAmelCase						,							_UpperCAmelCase						,							_UpperCAmelCase       )



       @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."       )
       def 	lowercase__							(				self     :    Optional[Any]       )					->		List[Any]:


             '''simple docstring'''




             pass



       def 	lowercase__							(				self     :    str						,							_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    List[Any]						,							_UpperCAmelCase     :    Tuple=False       )					->		Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				super()._prepare_for_class(_UpperCAmelCase						,							_UpperCAmelCase						,							return_labels=_UpperCAmelCase       )

             if return_labels:
                   if model_class.__name__ == "LevitForImageClassificationWithTeacher":
                         del inputs_dict["labels"]

             return inputs_dict



       def 	lowercase__							(				self     :    Any       )					->		Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*_UpperCAmelCase       )



       def 	lowercase__							(				self     :    str       )					->		List[str]:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*_UpperCAmelCase       )



       def 	lowercase__							(				self     :    Union[str, Any]       )					->		Any:


             '''simple docstring'''




             if not self.model_tester.is_training:
                   return

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

             for model_class in self.all_model_classes:
                   # LevitForImageClassificationWithTeacher supports inference-only
                   if (
                       model_class in get_values(_UpperCAmelCase       )
                       or model_class.__name__ == "LevitForImageClassificationWithTeacher"
                   ):
                         continue
                   UpperCAmelCase_     					=				model_class(_UpperCAmelCase       )
                   model.to(_UpperCAmelCase       )
                   model.train()
                   UpperCAmelCase_     					=				self._prepare_for_class(_UpperCAmelCase						,							_UpperCAmelCase						,							return_labels=_UpperCAmelCase       )
                   UpperCAmelCase_     					=				model(**_UpperCAmelCase       ).loss
                   loss.backward()



       def 	lowercase__							(				self     :    Union[str, Any]       )					->		Dict:


             '''simple docstring'''




             UpperCAmelCase_					,       UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs_for_common()
             if not self.model_tester.is_training:
                   return

             UpperCAmelCase_     					=				False
             UpperCAmelCase_     					=				True

             for model_class in self.all_model_classes:
                   if model_class in get_values(_UpperCAmelCase       ) or not model_class.supports_gradient_checkpointing:
                         continue
                   # LevitForImageClassificationWithTeacher supports inference-only
                   if model_class.__name__ == "LevitForImageClassificationWithTeacher":
                         continue
                   UpperCAmelCase_     					=				model_class(_UpperCAmelCase       )
                   model.gradient_checkpointing_enable()
                   model.to(_UpperCAmelCase       )
                   model.train()
                   UpperCAmelCase_     					=				self._prepare_for_class(_UpperCAmelCase						,							_UpperCAmelCase						,							return_labels=_UpperCAmelCase       )
                   UpperCAmelCase_     					=				model(**_UpperCAmelCase       ).loss
                   loss.backward()



       def 	lowercase__							(				self     :    Dict       )					->		Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_					,       UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs_for_common()

             UpperCAmelCase_     					=				[
                 {"title": "multi_label_classification", "num_labels": 2, "dtype": torch.float},
                 {"title": "single_label_classification", "num_labels": 1, "dtype": torch.long},
                 {"title": "regression", "num_labels": 1, "dtype": torch.float},
             ]

             for model_class in self.all_model_classes:
                   if (
                       model_class
                       not in [
                           *get_values(_UpperCAmelCase       ),
                       ]
                       or model_class.__name__ == "LevitForImageClassificationWithTeacher"
                   ):
                         continue

                   for problem_type in problem_types:
                         with self.subTest(msg=F"""Testing {model_class} with {problem_type['title']}"""       ):
                               UpperCAmelCase_     					=				problem_type["title"]
                               UpperCAmelCase_     					=				problem_type["num_labels"]

                               UpperCAmelCase_     					=				model_class(_UpperCAmelCase       )
                               model.to(_UpperCAmelCase       )
                               model.train()

                               UpperCAmelCase_     					=				self._prepare_for_class(_UpperCAmelCase						,							_UpperCAmelCase						,							return_labels=_UpperCAmelCase       )

                               if problem_type["num_labels"] > 1:
                                     UpperCAmelCase_     					=				inputs["labels"].unsqueeze(1       ).repeat(1						,							problem_type["num_labels"]       )

                               UpperCAmelCase_     					=				inputs["labels"].to(problem_type["dtype"]       )

                               # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
                               # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
                               # they have the same size." which is a symptom something in wrong for the regression problem.
                               # See https://github.com/huggingface/transformers/issues/11780
                               with warnings.catch_warnings(record=_UpperCAmelCase       ) as warning_list:
                                     UpperCAmelCase_     					=				model(**_UpperCAmelCase       ).loss
                               for w in warning_list:
                                     if "Using a target size that is different to the input size" in str(w.message       ):
                                           raise ValueError(
                                               F"""Something is going wrong in the regression problem: intercepted {w.message}"""       )

                               loss.backward()







       @slow
       def 	lowercase__							(				self     :    List[str]       )					->		Optional[Any]:


             '''simple docstring'''




             for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   UpperCAmelCase_     					=				LevitModel.from_pretrained(_UpperCAmelCase       )
                   self.assertIsNotNone(_UpperCAmelCase       )



def   a__      ( ):
      UpperCAmelCase_     					=				Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png" )
      return image


@require_torch
@require_vision
class  lowercase__  (       unittest.TestCase ):







       '''simple docstring'''


       @cached_property
       def 	lowercase__							(				self     :    List[str]       )					->		str:


             '''simple docstring'''




             return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]       )







       @slow
       def 	lowercase__							(				self     :    Tuple       )					->		int:


             '''simple docstring'''




             UpperCAmelCase_     					=				LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]       ).to(
                 _UpperCAmelCase       )

             UpperCAmelCase_     					=				self.default_image_processor
             UpperCAmelCase_     					=				prepare_img()
             UpperCAmelCase_     					=				image_processor(images=_UpperCAmelCase						,							return_tensors="pt"       ).to(_UpperCAmelCase       )

             # forward pass
             with torch.no_grad():
                   UpperCAmelCase_     					=				model(**_UpperCAmelCase       )

             # verify the logits
             UpperCAmelCase_     					=				torch.Size((1, 1000)       )
             self.assertEqual(outputs.logits.shape						,							_UpperCAmelCase       )

             UpperCAmelCase_     					=				torch.tensor([1.0448, -0.3745, -1.8317]       ).to(_UpperCAmelCase       )

             self.assertTrue(torch.allclose(outputs.logits[0, :3]						,							_UpperCAmelCase						,							atol=1e-4       )       )





'''simple docstring'''


import argparse

import pytorch_lightning as pl
import torch
from torch import nn

from transformers import LongformerForQuestionAnswering, LongformerModel






class     A ( pl.LightningModule ):





							def __init__(   self					:				Dict  ,		__a					:				List[str]						)				->						Tuple:
								super().__init__()
								__UpperCAmelCase      					=      model
								__UpperCAmelCase      					=      2
								__UpperCAmelCase      					=      nn.Linear(self.model.config.hidden_size  ,		self.num_labels						)






							def    snake_case__    (   self					:				int						)				->						int:
								pass





def    lowerCAmelCase       (		UpperCamelCase__					: str     ,       UpperCamelCase__					: str     ,       UpperCamelCase__					: str							):


	"""simple docstring"""
	# load longformer model from model identifier
	__UpperCAmelCase      					=      LongformerModel.from_pretrained(UpperCamelCase__							)
	__UpperCAmelCase      					=      LightningModel(UpperCamelCase__							)

	__UpperCAmelCase      					=      torch.load(UpperCamelCase__     ,       map_location=torch.device('''cpu'''							)							)
	lightning_model.load_state_dict(ckpt['''state_dict''']							)

	# init longformer question answering model
	__UpperCAmelCase      					=      LongformerForQuestionAnswering.from_pretrained(UpperCamelCase__							)

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

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


if __name__ == "__main__":
					__lowerCAmelCase						:							Optional[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."
					)
					__lowerCAmelCase						:							List[str]    							=						parser.parse_args()
					convert_longformer_qa_checkpoint_to_pytorch(
					    args.longformer_model, args.longformer_question_answering_ckpt_path, args.pytorch_dump_folder_path
					)
