

import collections
import inspect
import unittest

from transformers import DonutSwinConfig
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import is_torch_available

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


if is_torch_available():
  import torch
  from torch import nn

  from transformers import DonutSwinModel
  from transformers.models.donut.modeling_donut_swin import DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST






class 							A_			:


     def __init__(							self      :		Union[str, Any]      ,  UpperCAmelCase      :		str      ,  UpperCAmelCase      :		Any=1_3      ,  UpperCAmelCase      :		Optional[Any]=3_2      ,  UpperCAmelCase      :		Optional[int]=2      ,  UpperCAmelCase      :		Optional[int]=3      ,  UpperCAmelCase      :		List[Any]=1_6      ,  UpperCAmelCase      :		Dict=[1, 2, 1]      ,  UpperCAmelCase      :		List[str]=[2, 2, 4]      ,  UpperCAmelCase      :		Tuple=2      ,  UpperCAmelCase      :		Optional[Any]=2.0      ,  UpperCAmelCase      :		List[Any]=True      ,  UpperCAmelCase      :		str=0.0      ,  UpperCAmelCase      :		List[str]=0.0      ,  UpperCAmelCase      :		Tuple=0.1      ,  UpperCAmelCase      :		List[Any]="gelu"      ,  UpperCAmelCase      :		str=False      ,  UpperCAmelCase      :		Union[str, Any]=True      ,  UpperCAmelCase      :		Optional[Any]=0.02      ,  UpperCAmelCase      :		Any=1E-5      ,  UpperCAmelCase      :		List[Any]=True      ,  UpperCAmelCase      :		Optional[Any]=None      ,  UpperCAmelCase      :		Union[str, Any]=True      ,  UpperCAmelCase      :		Any=1_0      ,  UpperCAmelCase      :		Optional[int]=8      ,  )							->     int:
      __lowerCAmelCase:       Optional[Any]            =				parent
      __lowerCAmelCase:       Any            =				batch_size
      __lowerCAmelCase:       Dict            =				image_size
      __lowerCAmelCase:       int            =				patch_size
      __lowerCAmelCase:       str            =				num_channels
      __lowerCAmelCase:       Any            =				embed_dim
      __lowerCAmelCase:       Optional[int]            =				depths
      __lowerCAmelCase:       Optional[int]            =				num_heads
      __lowerCAmelCase:       Optional[int]            =				window_size
      __lowerCAmelCase:       Any            =				mlp_ratio
      __lowerCAmelCase:       List[Any]            =				qkv_bias
      __lowerCAmelCase:       List[str]            =				hidden_dropout_prob
      __lowerCAmelCase:       str            =				attention_probs_dropout_prob
      __lowerCAmelCase:       Union[str, Any]            =				drop_path_rate
      __lowerCAmelCase:       Union[str, Any]            =				hidden_act
      __lowerCAmelCase:       Any            =				use_absolute_embeddings
      __lowerCAmelCase:       Dict            =				patch_norm
      __lowerCAmelCase:       Union[str, Any]            =				layer_norm_eps
      __lowerCAmelCase:       Optional[int]            =				initializer_range
      __lowerCAmelCase:       List[str]            =				is_training
      __lowerCAmelCase:       Any            =				scope
      __lowerCAmelCase:       List[Any]            =				use_labels
      __lowerCAmelCase:       List[Any]            =				type_sequence_label_size
      __lowerCAmelCase:       Union[str, Any]            =				encoder_stride


     def   UpperCAmelCase  (							self      :		Dict )							->     Dict:
      __lowerCAmelCase:       str            =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

      __lowerCAmelCase:       int            =				None
      if self.use_labels:
       __lowerCAmelCase:       Any            =				ids_tensor([self.batch_size]      ,  self.type_sequence_label_size )

      __lowerCAmelCase:       Union[str, Any]            =				self.get_config()

      return config, pixel_values, labels


     def   UpperCAmelCase  (							self      :		Dict )							->     Any:
      return DonutSwinConfig(
          image_size=self.image_size      ,  patch_size=self.patch_size      ,  num_channels=self.num_channels      ,  embed_dim=self.embed_dim      ,  depths=self.depths      ,  num_heads=self.num_heads      ,  window_size=self.window_size      ,  mlp_ratio=self.mlp_ratio      ,  qkv_bias=self.qkv_bias      ,  hidden_dropout_prob=self.hidden_dropout_prob      ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,  drop_path_rate=self.drop_path_rate      ,  hidden_act=self.hidden_act      ,  use_absolute_embeddings=self.use_absolute_embeddings      ,  path_norm=self.patch_norm      ,  layer_norm_eps=self.layer_norm_eps      ,  initializer_range=self.initializer_range      ,  encoder_stride=self.encoder_stride      ,  )


     def   UpperCAmelCase  (							self      :		Dict      ,  UpperCAmelCase      :		Union[str, Any]      ,  UpperCAmelCase      :		str      ,  UpperCAmelCase      :		Dict )							->     Any:
      __lowerCAmelCase:       int            =				DonutSwinModel(config=UpperCAmelCase )
      model.to(UpperCAmelCase )
      model.eval()
      __lowerCAmelCase:       Dict            =				model(UpperCAmelCase )

      __lowerCAmelCase:       Dict            =				((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths ) - 1))
      __lowerCAmelCase:       List[Any]            =				int(config.embed_dim * 2 ** (len(config.depths ) - 1) )

      self.parent.assertEqual(result.last_hidden_state.shape      ,  (self.batch_size, expected_seq_len, expected_dim) )


     def   UpperCAmelCase  (							self      :		Any )							->     Dict:
      __lowerCAmelCase:       str            =				self.prepare_config_and_inputs()
      (
          (
          __lowerCAmelCase
      )     ,							(
          __lowerCAmelCase
      )     ,							(
          __lowerCAmelCase
      )     ,							
      ):       Any            =				config_and_inputs
      __lowerCAmelCase:       str            =				{'pixel_values': pixel_values}
      return config, inputs_dict







@require_torch
class 							A_			(		snake_case__							,     snake_case__							,     unittest.TestCase ):
     _lowercase      :					List[str]       =     (DonutSwinModel,) if is_torch_available() else ()
     _lowercase      :					List[Any]       =     {'feature-extraction': DonutSwinModel} if is_torch_available() else {}
     _lowercase      :					str       =     True

     _lowercase      :					Optional[Any]       =     False
     _lowercase      :					Any       =     False
     _lowercase      :					Dict       =     False


     def   UpperCAmelCase  (							self      :		Any )							->     Optional[Any]:
      __lowerCAmelCase:       Optional[int]            =				DonutSwinModelTester(self )
      __lowerCAmelCase:       Dict            =				ConfigTester(self      ,  config_class=UpperCAmelCase      ,  embed_dim=3_7 )


     def   UpperCAmelCase  (							self      :		List[Any] )							->     int:
      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   UpperCAmelCase  (							self      :		Any )							->     Any:
      return


     def   UpperCAmelCase  (							self      :		List[Any] )							->     Any:
      __lowerCAmelCase:       Union[str, Any]            =				self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		Union[str, Any] )							->     Union[str, Any]:
      # DonutSwin does not use inputs_embeds
      pass


     def   UpperCAmelCase  (							self      :		Union[str, Any] )							->     Optional[Any]:
      __lowerCAmelCase     ,							__lowerCAmelCase:       List[Any]            =				self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
       __lowerCAmelCase:       str            =				model_class(UpperCAmelCase )
       self.assertIsInstance(model.get_input_embeddings()      ,  (nn.Module) )
       __lowerCAmelCase:       int            =				model.get_output_embeddings()
       self.assertTrue(x is None or isinstance(UpperCAmelCase      ,  nn.Linear ) )


     def   UpperCAmelCase  (							self      :		int )							->     str:
      __lowerCAmelCase     ,							__lowerCAmelCase:       int            =				self.model_tester.prepare_config_and_inputs_for_common()

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

       __lowerCAmelCase:       Any            =				['pixel_values']
       self.assertListEqual(arg_names[:1]      ,  UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		Union[str, Any] )							->     Dict:
      __lowerCAmelCase     ,							__lowerCAmelCase:       List[str]            =				self.model_tester.prepare_config_and_inputs_for_common()
      __lowerCAmelCase:       List[str]            =				True

      for model_class in self.all_model_classes:
       __lowerCAmelCase:       List[Any]            =				True
       __lowerCAmelCase:       Any            =				False
       __lowerCAmelCase:       Dict            =				True
       __lowerCAmelCase:       Optional[int]            =				model_class(UpperCAmelCase )
       model.to(UpperCAmelCase )
       model.eval()
       with torch.no_grad():
        __lowerCAmelCase:       Union[str, Any]            =				model(**self._prepare_for_class(UpperCAmelCase      ,  UpperCAmelCase ) )
       __lowerCAmelCase:       Union[str, Any]            =				outputs.attentions
       __lowerCAmelCase:       Optional[int]            =				len(self.model_tester.depths )
       self.assertEqual(len(UpperCAmelCase )      ,  UpperCAmelCase )

       # check that output_attentions also work using config
       del inputs_dict["output_attentions"]
       __lowerCAmelCase:       str            =				True
       __lowerCAmelCase:       Dict            =				config.window_size**2
       __lowerCAmelCase:       Optional[Any]            =				model_class(UpperCAmelCase )
       model.to(UpperCAmelCase )
       model.eval()
       with torch.no_grad():
        __lowerCAmelCase:       Optional[Any]            =				model(**self._prepare_for_class(UpperCAmelCase      ,  UpperCAmelCase ) )
       __lowerCAmelCase:       List[Any]            =				outputs.attentions
       self.assertEqual(len(UpperCAmelCase )      ,  UpperCAmelCase )

       self.assertListEqual(
           list(attentions[0].shape[-3:] )      ,  [self.model_tester.num_heads[0], window_size_squared, window_size_squared]      ,  )
       __lowerCAmelCase:       List[Any]            =				len(UpperCAmelCase )

       # Check attention is always last and order is fine
       __lowerCAmelCase:       Dict            =				True
       __lowerCAmelCase:       List[str]            =				True
       __lowerCAmelCase:       Union[str, Any]            =				model_class(UpperCAmelCase )
       model.to(UpperCAmelCase )
       model.eval()
       with torch.no_grad():
        __lowerCAmelCase:       int            =				model(**self._prepare_for_class(UpperCAmelCase      ,  UpperCAmelCase ) )

       if hasattr(self.model_tester      ,  'num_hidden_states_types' ):
        __lowerCAmelCase:       List[str]            =				self.model_tester.num_hidden_states_types
       else:
        # also another +1 for reshaped_hidden_states
        __lowerCAmelCase:       Dict            =				2
       self.assertEqual(out_len + added_hidden_states      ,  len(UpperCAmelCase ) )

       __lowerCAmelCase:       Optional[Any]            =				outputs.attentions

       self.assertEqual(len(UpperCAmelCase )      ,  UpperCAmelCase )

       self.assertListEqual(
           list(self_attentions[0].shape[-3:] )      ,  [self.model_tester.num_heads[0], window_size_squared, window_size_squared]      ,  )


     def   UpperCAmelCase  (							self      :		Tuple      ,  UpperCAmelCase      :		List[str]      ,  UpperCAmelCase      :		List[str]      ,  UpperCAmelCase      :		Dict      ,  UpperCAmelCase      :		Any )							->     Tuple:
      __lowerCAmelCase:       Dict            =				model_class(UpperCAmelCase )
      model.to(UpperCAmelCase )
      model.eval()

      with torch.no_grad():
       __lowerCAmelCase:       int            =				model(**self._prepare_for_class(UpperCAmelCase      ,  UpperCAmelCase ) )

      __lowerCAmelCase:       Union[str, Any]            =				outputs.hidden_states

      __lowerCAmelCase:       Optional[int]            =				getattr(
          self.model_tester      ,  'expected_num_hidden_layers'      ,  len(self.model_tester.depths ) + 1 )
      self.assertEqual(len(UpperCAmelCase )      ,  UpperCAmelCase )

      # DonutSwin has a different seq_length
      __lowerCAmelCase:       Optional[int]            =				(
          config.patch_size
          if isinstance(config.patch_size      ,  collections.abc.Iterable )
          else (config.patch_size, config.patch_size)
      )

      __lowerCAmelCase:       List[Any]            =				(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

      self.assertListEqual(
          list(hidden_states[0].shape[-2:] )      ,  [num_patches, self.model_tester.embed_dim]      ,  )

      __lowerCAmelCase:       List[str]            =				outputs.reshaped_hidden_states
      self.assertEqual(len(UpperCAmelCase )      ,  UpperCAmelCase )

      __lowerCAmelCase     ,							__lowerCAmelCase     ,							__lowerCAmelCase     ,							__lowerCAmelCase:       Dict            =				reshaped_hidden_states[0].shape
      __lowerCAmelCase:       List[Any]            =				(
          reshaped_hidden_states[0].view(UpperCAmelCase      ,  UpperCAmelCase      ,  height * width ).permute(0      ,  2      ,  1 )
      )
      self.assertListEqual(
          list(reshaped_hidden_states.shape[-2:] )      ,  [num_patches, self.model_tester.embed_dim]      ,  )


     def   UpperCAmelCase  (							self      :		List[str] )							->     Dict:
      __lowerCAmelCase     ,							__lowerCAmelCase:       Any            =				self.model_tester.prepare_config_and_inputs_for_common()

      __lowerCAmelCase:       int            =				(
          self.model_tester.image_size
          if isinstance(self.model_tester.image_size      ,  collections.abc.Iterable )
          else (self.model_tester.image_size, self.model_tester.image_size)
      )

      for model_class in self.all_model_classes:
       __lowerCAmelCase:       str            =				True
       self.check_hidden_states_output(UpperCAmelCase      ,  UpperCAmelCase      ,  UpperCAmelCase      ,  UpperCAmelCase )

       # check that output_hidden_states also work using config
       del inputs_dict["output_hidden_states"]
       __lowerCAmelCase:       str            =				True

       self.check_hidden_states_output(UpperCAmelCase      ,  UpperCAmelCase      ,  UpperCAmelCase      ,  UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		int )							->     List[Any]:
      __lowerCAmelCase     ,							__lowerCAmelCase:       List[Any]            =				self.model_tester.prepare_config_and_inputs_for_common()
      __lowerCAmelCase:       Optional[int]            =				3

      __lowerCAmelCase:       Dict            =				(
          self.model_tester.image_size
          if isinstance(self.model_tester.image_size      ,  collections.abc.Iterable )
          else (self.model_tester.image_size, self.model_tester.image_size)
      )
      __lowerCAmelCase:       Tuple            =				(
          config.patch_size
          if isinstance(config.patch_size      ,  collections.abc.Iterable )
          else (config.patch_size, config.patch_size)
      )

      __lowerCAmelCase:       Dict            =				image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
      __lowerCAmelCase:       Union[str, Any]            =				image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

      for model_class in self.all_model_classes:
       __lowerCAmelCase:       List[Any]            =				True
       self.check_hidden_states_output(UpperCAmelCase      ,  UpperCAmelCase      ,  UpperCAmelCase      ,  (padded_height, padded_width) )

       # check that output_hidden_states also work using config
       del inputs_dict["output_hidden_states"]
       __lowerCAmelCase:       Union[str, Any]            =				True
       self.check_hidden_states_output(UpperCAmelCase      ,  UpperCAmelCase      ,  UpperCAmelCase      ,  (padded_height, padded_width) )


     @slow
     def   UpperCAmelCase  (							self      :		Optional[int] )							->     Optional[Any]:
      for model_name in DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       __lowerCAmelCase:       List[str]            =				DonutSwinModel.from_pretrained(UpperCAmelCase )
       self.assertIsNotNone(UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		Optional[Any] )							->     str:
      __lowerCAmelCase     ,							__lowerCAmelCase:       List[str]            =				self.model_tester.prepare_config_and_inputs_for_common()

      __lowerCAmelCase:       Union[str, Any]            =				_config_zero_init(UpperCAmelCase )
      for model_class in self.all_model_classes:
       __lowerCAmelCase:       int            =				model_class(config=UpperCAmelCase )
       for name, param in model.named_parameters():
        if "embeddings" not in name and param.requires_grad:
         self.assertIn(
             ((param.data.mean() * 1E9).round() / 1E9).item()      ,  [0.0, 1.0]      ,  msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''      ,  )








import torch

from transformers import PreTrainedModel

from .custom_configuration import CustomConfig, NoSuperInitConfig






class 							A_			(		snake_case__ ):
     _lowercase      :					Any       =     CustomConfig


     def __init__(							self      :		List[Any]      ,  UpperCAmelCase      :		Optional[Any] )							->     int:
      super().__init__(UpperCAmelCase )
      __lowerCAmelCase:       Tuple            =				torch.nn.Linear(config.hidden_size      ,  config.hidden_size )


     def   UpperCAmelCase  (							self      :		Optional[int]      ,  UpperCAmelCase      :		Tuple )							->     Union[str, Any]:
      return self.linear(UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		str      ,  UpperCAmelCase      :		str )							->     Optional[int]:
      pass







class 							A_			(		snake_case__ ):
     _lowercase      :					str       =     NoSuperInitConfig


     def __init__(							self      :		List[Any]      ,  UpperCAmelCase      :		Union[str, Any] )							->     int:
      super().__init__(UpperCAmelCase )
      __lowerCAmelCase:       int            =				torch.nn.Linear(config.attribute      ,  config.attribute )


     def   UpperCAmelCase  (							self      :		Tuple      ,  UpperCAmelCase      :		Tuple )							->     Tuple:
      return self.linear(UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		Optional[Any]      ,  UpperCAmelCase      :		Any )							->     str:
      pass





