


"""simple docstring"""




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


__magic_name__	:   Union[str, Any]          =							[
    'AutomaticSpeechRecognition',
    'AudioClassification',
    'ImageClassification',
    'LanguageModeling',
    'QuestionAnsweringExtractive',
    'Summarization',
    'TaskTemplate',
    'TextClassification',
]

__magic_name__	:   Any          =							get_logger(__name__)


__magic_name__	:   List[Any]          =							{
    AutomaticSpeechRecognition.task: AutomaticSpeechRecognition,
    AudioClassification.task: AudioClassification,
    ImageClassification.task: ImageClassification,
    LanguageModeling.task: LanguageModeling,
    QuestionAnsweringExtractive.task: QuestionAnsweringExtractive,
    Summarization.task: Summarization,
    TextClassification.task: TextClassification,
}





def 			a_   (  lowercase__   :dict    ):
			__lowerCamelCase     :   Dict        =			task_template_dict.get("""task"""    )
			if task_name is None:
						logger.warning(f'Couldn\'t find template for task \'{task_name}\'. Available templates: {list(lowercase__    )}'    )
						return None
			__lowerCamelCase     :   List[str]        =			NAMEaTEMPLATE.get(lowercase__    )
			return template.from_dict(lowercase__    )








"""simple docstring"""




import inspect
import unittest

from transformers import SwinaSRConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_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 SwinaSRForImageSuperResolution, SwinaSRModel
					from transformers.models.swinasr.modeling_swinasr import SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
					from PIL import Image

					from transformers import SwinaSRImageProcessor




class  __snake_case      :




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

										# here we set some attributes to make tests pass
										__lowerCamelCase        =			len(A_      )
										__lowerCamelCase        =			embed_dim
										__lowerCamelCase        =			(image_size // patch_size) ** 2




							def      __a      (    self:						Tuple      ):
										__lowerCamelCase        =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

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

										__lowerCamelCase        =			self.get_config()

										return config, pixel_values, labels




							def      __a      (    self:						int      ):
										return SwinaSRConfig(
										    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    ,       upscale=self.upscale    ,       )




							def      __a      (    self:						str    ,       A_:						str    ,       A_:						Optional[int]    ,       A_:						List[str]      ):
										__lowerCamelCase        =			SwinaSRModel(config=A_      )
										model.to(A_      )
										model.eval()
										__lowerCamelCase        =			model(A_      )

										self.parent.assertEqual(
										    result.last_hidden_state.shape    ,       (self.batch_size, self.embed_dim, self.image_size, self.image_size)      )




							def      __a      (    self:						Any    ,       A_:						Tuple    ,       A_:						Any    ,       A_:						Dict      ):
										__lowerCamelCase        =			SwinaSRForImageSuperResolution(A_      )
										model.to(A_      )
										model.eval()
										__lowerCamelCase        =			model(A_      )

										__lowerCamelCase        =			self.image_size * self.upscale
										self.parent.assertEqual(
										    result.reconstruction.shape    ,       (self.batch_size, self.num_channels, expected_image_size, expected_image_size)      )




							def      __a      (    self:						Tuple      ):
										__lowerCamelCase        =			self.prepare_config_and_inputs()
										__lowerCamelCase    ,__lowerCamelCase    ,__lowerCamelCase        =			config_and_inputs
										__lowerCamelCase        =			{"""pixel_values""": pixel_values}
										return config, inputs_dict




@require_torch
class  __snake_case      (lowerCamelCase ,		lowerCamelCase ,		unittest.TestCase     ):
							__a				   =					(SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
							__a				   =					{'''feature-extraction''': SwinaSRModel} if is_torch_available() else {}

							__a				   =					False
							__a				   =					False
							__a				   =					False
							__a				   =					False
							__a				   =					False




							def      __a      (    self:						List[str]      ):
										__lowerCamelCase        =			SwinaSRModelTester(self      )
										__lowerCamelCase        =			ConfigTester(self    ,       config_class=A_    ,       embed_dim=37      )




							def      __a      (    self:						Optional[Any]      ):
										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      __a      (    self:						Tuple      ):
										__lowerCamelCase        =			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*A_      )




							def      __a      (    self:						Any      ):
										__lowerCamelCase        =			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_super_resolution(*A_      )




							@unittest.skip(reason="""Got `CUDA error: misaligned address` with PyTorch 2.0.0."""      )
							def      __a      (    self:						Union[str, Any]      ):
										pass




							@unittest.skip(reason="""Swin2SR does not use inputs_embeds"""      )
							def      __a      (    self:						str      ):
										pass




							@unittest.skip(reason="""Swin2SR does not support training yet"""      )
							def      __a      (    self:						int      ):
										pass




							@unittest.skip(reason="""Swin2SR does not support training yet"""      )
							def      __a      (    self:						Optional[int]      ):
										pass




							def      __a      (    self:						List[Any]      ):
										__lowerCamelCase    ,__lowerCamelCase        =			self.model_tester.prepare_config_and_inputs_for_common()

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




							def      __a      (    self:						Optional[Any]      ):
										__lowerCamelCase    ,__lowerCamelCase        =			self.model_tester.prepare_config_and_inputs_for_common()

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

													__lowerCamelCase        =			["""pixel_values"""]
													self.assertListEqual(arg_names[:1]    ,       A_      )




							@slow
							def      __a      (    self:						List[Any]      ):
										for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													__lowerCamelCase        =			SwinaSRModel.from_pretrained(A_      )
													self.assertIsNotNone(A_      )




							def      __a      (    self:						Optional[int]      ):
										__lowerCamelCase    ,__lowerCamelCase        =			self.model_tester.prepare_config_and_inputs_for_common()

										__lowerCamelCase        =			_config_zero_init(A_      )
										for model_class in self.all_model_classes:
													__lowerCamelCase        =			model_class(config=A_      )
													for name, param in model.named_parameters():
																if "logit_scale" in name:
																			continue
																if 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'    ,       )




							def      __a      (    self:						Any      ):
										__lowerCamelCase    ,__lowerCamelCase        =			self.model_tester.prepare_config_and_inputs_for_common()
										__lowerCamelCase        =			True

										for model_class in self.all_model_classes:
													__lowerCamelCase        =			True
													__lowerCamelCase        =			False
													__lowerCamelCase        =			True
													__lowerCamelCase        =			model_class(A_      )
													model.to(A_      )
													model.eval()
													with torch.no_grad():
																__lowerCamelCase        =			model(**self._prepare_for_class(A_    ,       A_      )      )
													__lowerCamelCase        =			outputs.attentions
													__lowerCamelCase        =			len(self.model_tester.depths      )
													self.assertEqual(len(A_      )    ,       A_      )

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

													self.assertListEqual(
													    list(attentions[0].shape[-3:]      )    ,       [self.model_tester.num_heads[0], window_size_squared, window_size_squared]    ,       )
													__lowerCamelCase        =			len(A_      )

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

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

													__lowerCamelCase        =			outputs.attentions

													self.assertEqual(len(A_      )    ,       A_      )

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




@require_vision
@require_torch
@slow
class  __snake_case      (unittest.TestCase     ):




							def      __a      (    self:						List[Any]      ):
										__lowerCamelCase        =			SwinaSRImageProcessor()
										__lowerCamelCase        =			SwinaSRForImageSuperResolution.from_pretrained("""caidas/swin2SR-classical-sr-x2-64"""      ).to(A_      )

										__lowerCamelCase        =			Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""      )
										__lowerCamelCase        =			processor(images=A_    ,       return_tensors="""pt"""      ).to(A_      )

										# forward pass
										with torch.no_grad():
													__lowerCamelCase        =			model(**A_      )

										# verify the logits
										__lowerCamelCase        =			torch.Size([1, 3, 9_76, 12_96]      )
										self.assertEqual(outputs.reconstruction.shape    ,       A_      )
										__lowerCamelCase        =			torch.tensor(
										    [[0.5_458, 0.5_546, 0.5_638], [0.5_526, 0.5_565, 0.5_651], [0.5_396, 0.5_426, 0.5_621]]      ).to(A_      )
										self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]    ,       A_    ,       atol=1E-4      )      )




