
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 			UpperCamelCase_			:




							'''simple docstring'''

							def __init__(				self     ,						UpperCamelCase     ,						UpperCamelCase=13     ,						UpperCamelCase=32     ,						UpperCamelCase=2     ,						UpperCamelCase=3     ,						UpperCamelCase=16     ,						UpperCamelCase=[1, 2, 1]     ,						UpperCamelCase=[2, 2, 4]     ,						UpperCamelCase=2     ,						UpperCamelCase=2.0     ,						UpperCamelCase=True     ,						UpperCamelCase=0.0     ,						UpperCamelCase=0.0     ,						UpperCamelCase=0.1     ,						UpperCamelCase="gelu"     ,						UpperCamelCase=False     ,						UpperCamelCase=True     ,						UpperCamelCase=0.02     ,						UpperCamelCase=1E-5     ,						UpperCamelCase=True     ,						UpperCamelCase=None     ,						UpperCamelCase=True     ,						UpperCamelCase=10     ,						UpperCamelCase=8     ,						) ->	Optional[int]:
										UpperCamelCase__    :		Optional[int]     =				parent
										UpperCamelCase__    :		Union[str, Any]     =				batch_size
										UpperCamelCase__    :		Optional[Any]     =				image_size
										UpperCamelCase__    :		Any     =				patch_size
										UpperCamelCase__    :		int     =				num_channels
										UpperCamelCase__    :		List[Any]     =				embed_dim
										UpperCamelCase__    :		Optional[Any]     =				depths
										UpperCamelCase__    :		Dict     =				num_heads
										UpperCamelCase__    :		List[Any]     =				window_size
										UpperCamelCase__    :		Optional[int]     =				mlp_ratio
										UpperCamelCase__    :		Union[str, Any]     =				qkv_bias
										UpperCamelCase__    :		List[Any]     =				hidden_dropout_prob
										UpperCamelCase__    :		Optional[Any]     =				attention_probs_dropout_prob
										UpperCamelCase__    :		Union[str, Any]     =				drop_path_rate
										UpperCamelCase__    :		List[Any]     =				hidden_act
										UpperCamelCase__    :		Any     =				use_absolute_embeddings
										UpperCamelCase__    :		Any     =				patch_norm
										UpperCamelCase__    :		int     =				layer_norm_eps
										UpperCamelCase__    :		Optional[Any]     =				initializer_range
										UpperCamelCase__    :		Optional[Any]     =				is_training
										UpperCamelCase__    :		str     =				scope
										UpperCamelCase__    :		List[Any]     =				use_labels
										UpperCamelCase__    :		str     =				type_sequence_label_size
										UpperCamelCase__    :		int     =				encoder_stride
							def        lowerCAmelCase__ (				self) ->	Optional[int]:
										UpperCamelCase__    :		int     =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

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

										return config, pixel_values, labels
							def        lowerCAmelCase__ (				self) ->	Tuple:
										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        lowerCAmelCase__ (				self     ,						UpperCamelCase     ,						UpperCamelCase     ,						UpperCamelCase) ->	Optional[Any]:
										UpperCamelCase__    :		Tuple     =				DonutSwinModel(config=UpperCamelCase)
										model.to(UpperCamelCase)
										model.eval()
										UpperCamelCase__    :		List[str]     =				model(UpperCamelCase)

										UpperCamelCase__    :		str     =				((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths) - 1))
										UpperCamelCase__    :		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        lowerCAmelCase__ (				self) ->	Optional[Any]:
										UpperCamelCase__    :		Union[str, Any]     =				self.prepare_config_and_inputs()
										(
										    UpperCamelCase__			
										)    :		int     =				config_and_inputs
										UpperCamelCase__    :		Any     =				{'pixel_values': pixel_values}
										return config, inputs_dict



@require_torch
class 			UpperCamelCase_			(						__UpperCamelCase	,     __UpperCamelCase	,     unittest.TestCase   ):




							'''simple docstring'''

							UpperCamelCase_					   =							(DonutSwinModel,) if is_torch_available() else ()
							UpperCamelCase_					   =							{"""feature-extraction""": DonutSwinModel} if is_torch_available() else {}
							UpperCamelCase_					   =							True

							UpperCamelCase_					   =							False
							UpperCamelCase_					   =							False
							UpperCamelCase_					   =							False
							def        lowerCAmelCase__ (				self) ->	Optional[Any]:
										UpperCamelCase__    :		str     =				DonutSwinModelTester(self)
										UpperCamelCase__    :		str     =				ConfigTester(self     ,						config_class=UpperCamelCase     ,						embed_dim=37)
							def        lowerCAmelCase__ (				self) ->	Any:
										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        lowerCAmelCase__ (				self) ->	Optional[int]:
										return
							def        lowerCAmelCase__ (				self) ->	Tuple:
										UpperCamelCase__    :		str     =				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*UpperCamelCase)
							def        lowerCAmelCase__ (				self) ->	List[str]:
										# DonutSwin does not use inputs_embeds
										pass
							def        lowerCAmelCase__ (				self) ->	List[Any]:
										UpperCamelCase__    :		str     =				self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													UpperCamelCase__    :		Union[str, Any]     =				model_class(UpperCamelCase)
													self.assertIsInstance(model.get_input_embeddings()     ,						(nn.Module))
													UpperCamelCase__    :		Optional[int]     =				model.get_output_embeddings()
													self.assertTrue(x is None or isinstance(UpperCamelCase     ,						nn.Linear))
							def        lowerCAmelCase__ (				self) ->	List[Any]:
										UpperCamelCase__    :		List[Any]     =				self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													UpperCamelCase__    :		Union[str, Any]     =				model_class(UpperCamelCase)
													UpperCamelCase__    :		Any     =				inspect.signature(model.forward)
													# signature.parameters is an OrderedDict => so arg_names order is deterministic
													UpperCamelCase__    :		int     =				[*signature.parameters.keys()]

													UpperCamelCase__    :		Optional[int]     =				['pixel_values']
													self.assertListEqual(arg_names[:1]     ,						UpperCamelCase)
							def        lowerCAmelCase__ (				self) ->	Tuple:
										UpperCamelCase__    :		Tuple     =				self.model_tester.prepare_config_and_inputs_for_common()
										UpperCamelCase__    :		Union[str, Any]     =				True

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

													# check that output_attentions also work using config
													del inputs_dict["output_attentions"]
													UpperCamelCase__    :		Optional[int]     =				True
													UpperCamelCase__    :		Optional[int]     =				config.window_size**2
													UpperCamelCase__    :		Optional[Any]     =				model_class(UpperCamelCase)
													model.to(UpperCamelCase)
													model.eval()
													with torch.no_grad():
																UpperCamelCase__    :		Dict     =				model(**self._prepare_for_class(UpperCamelCase     ,						UpperCamelCase))
													UpperCamelCase__    :		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]     ,						)
													UpperCamelCase__    :		Union[str, Any]     =				len(UpperCamelCase)

													# Check attention is always last and order is fine
													UpperCamelCase__    :		str     =				True
													UpperCamelCase__    :		str     =				True
													UpperCamelCase__    :		Dict     =				model_class(UpperCamelCase)
													model.to(UpperCamelCase)
													model.eval()
													with torch.no_grad():
																UpperCamelCase__    :		Any     =				model(**self._prepare_for_class(UpperCamelCase     ,						UpperCamelCase))

													if hasattr(self.model_tester     ,						'num_hidden_states_types'):
																UpperCamelCase__    :		int     =				self.model_tester.num_hidden_states_types
													else:
																# also another +1 for reshaped_hidden_states
																UpperCamelCase__    :		Dict     =				2
													self.assertEqual(out_len + added_hidden_states     ,						len(UpperCamelCase))

													UpperCamelCase__    :		List[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        lowerCAmelCase__ (				self     ,						UpperCamelCase     ,						UpperCamelCase     ,						UpperCamelCase     ,						UpperCamelCase) ->	Union[str, Any]:
										UpperCamelCase__    :		Dict     =				model_class(UpperCamelCase)
										model.to(UpperCamelCase)
										model.eval()

										with torch.no_grad():
													UpperCamelCase__    :		str     =				model(**self._prepare_for_class(UpperCamelCase     ,						UpperCamelCase))

										UpperCamelCase__    :		Dict     =				outputs.hidden_states

										UpperCamelCase__    :		Union[str, Any]     =				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
										UpperCamelCase__    :		Optional[int]     =				(
										    config.patch_size
										    if isinstance(config.patch_size     ,						collections.abc.Iterable)
										    else (config.patch_size, config.patch_size)
										)

										UpperCamelCase__    :		Optional[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]     ,						)

										UpperCamelCase__    :		Tuple     =				outputs.reshaped_hidden_states
										self.assertEqual(len(UpperCamelCase)     ,						UpperCamelCase)

										UpperCamelCase__    :		str     =				reshaped_hidden_states[0].shape
										UpperCamelCase__    :		int     =				(
										    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        lowerCAmelCase__ (				self) ->	List[Any]:
										UpperCamelCase__    :		List[str]     =				self.model_tester.prepare_config_and_inputs_for_common()

										UpperCamelCase__    :		Optional[Any]     =				(
										    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:
													UpperCamelCase__    :		List[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"]
													UpperCamelCase__    :		Optional[int]     =				True

													self.check_hidden_states_output(UpperCamelCase     ,						UpperCamelCase     ,						UpperCamelCase     ,						UpperCamelCase)
							def        lowerCAmelCase__ (				self) ->	Dict:
										UpperCamelCase__    :		Dict     =				self.model_tester.prepare_config_and_inputs_for_common()
										UpperCamelCase__    :		Optional[int]     =				3

										UpperCamelCase__    :		Tuple     =				(
										    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)
										)
										UpperCamelCase__    :		Union[str, Any]     =				(
										    config.patch_size
										    if isinstance(config.patch_size     ,						collections.abc.Iterable)
										    else (config.patch_size, config.patch_size)
										)

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

										for model_class in self.all_model_classes:
													UpperCamelCase__    :		Dict     =				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"]
													UpperCamelCase__    :		Dict     =				True
													self.check_hidden_states_output(UpperCamelCase     ,						UpperCamelCase     ,						UpperCamelCase     ,						(padded_height, padded_width))
							@slow
							def        lowerCAmelCase__ (				self) ->	List[Any]:
										for model_name in DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													UpperCamelCase__    :		List[str]     =				DonutSwinModel.from_pretrained(UpperCamelCase)
													self.assertIsNotNone(UpperCamelCase)






							def        lowerCAmelCase__ (				self) ->	List[str]:
										UpperCamelCase__    :		Any     =				self.model_tester.prepare_config_and_inputs_for_common()

										UpperCamelCase__    :		Any     =				_config_zero_init(UpperCamelCase)
										for model_class in self.all_model_classes:
													UpperCamelCase__    :		List[Any]     =				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"""     ,						)





from typing import Optional

from ..utils.logging import get_logger
from .audio_classificiation import AudioClassification
from .automatic_speech_recognition import AutomaticSpeechRecognition
from .base import TaskTemplate
from .image_classification import ImageClassification
from .language_modeling import LanguageModeling
from .question_answering import QuestionAnsweringExtractive
from .summarization import Summarization
from .text_classification import TextClassification


UpperCAmelCase__				:  Tuple       					=				[
    '''AutomaticSpeechRecognition''',
    '''AudioClassification''',
    '''ImageClassification''',
    '''LanguageModeling''',
    '''QuestionAnsweringExtractive''',
    '''Summarization''',
    '''TaskTemplate''',
    '''TextClassification''',
]

UpperCAmelCase__				:  Dict       					=				get_logger(__name__)


UpperCAmelCase__				:  Dict       					=				{
    AutomaticSpeechRecognition.task: AutomaticSpeechRecognition,
    AudioClassification.task: AudioClassification,
    ImageClassification.task: ImageClassification,
    LanguageModeling.task: LanguageModeling,
    QuestionAnsweringExtractive.task: QuestionAnsweringExtractive,
    Summarization.task: Summarization,
    TextClassification.task: TextClassification,
}



def       _lowercase					(      __SCREAMING_SNAKE_CASE )	->     Optional[TaskTemplate]:
			UpperCamelCase__    :		Any     =				task_template_dict.get('task' )
			if task_name is None:
						logger.warning(F"""Couldn't find template for task '{task_name}'. Available templates: {list(__SCREAMING_SNAKE_CASE )}""" )
						return None
			UpperCamelCase__    :		Any     =				NAMEaTEMPLATE.get(__SCREAMING_SNAKE_CASE )
			return template.from_dict(__SCREAMING_SNAKE_CASE )



