




'''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      _a       :






     def __init__(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE=13      ,_SCREAMING_SNAKE_CASE=32      ,_SCREAMING_SNAKE_CASE=1      ,_SCREAMING_SNAKE_CASE=3      ,_SCREAMING_SNAKE_CASE=16      ,_SCREAMING_SNAKE_CASE=[1, 2, 1]      ,_SCREAMING_SNAKE_CASE=[2, 2, 4]      ,_SCREAMING_SNAKE_CASE=2      ,_SCREAMING_SNAKE_CASE=2.0      ,_SCREAMING_SNAKE_CASE=True      ,_SCREAMING_SNAKE_CASE=0.0      ,_SCREAMING_SNAKE_CASE=0.0      ,_SCREAMING_SNAKE_CASE=0.1      ,_SCREAMING_SNAKE_CASE="gelu"      ,_SCREAMING_SNAKE_CASE=False      ,_SCREAMING_SNAKE_CASE=True      ,_SCREAMING_SNAKE_CASE=0.0_2      ,_SCREAMING_SNAKE_CASE=1e-5      ,_SCREAMING_SNAKE_CASE=True      ,_SCREAMING_SNAKE_CASE=None      ,_SCREAMING_SNAKE_CASE=False      ,_SCREAMING_SNAKE_CASE=2      ,)							->    Optional[Any]:
          _snake_case							=   parent
          _snake_case							=   batch_size
          _snake_case							=   image_size
          _snake_case							=   patch_size
          _snake_case							=   num_channels
          _snake_case							=   embed_dim
          _snake_case							=   depths
          _snake_case							=   num_heads
          _snake_case							=   window_size
          _snake_case							=   mlp_ratio
          _snake_case							=   qkv_bias
          _snake_case							=   hidden_dropout_prob
          _snake_case							=   attention_probs_dropout_prob
          _snake_case							=   drop_path_rate
          _snake_case							=   hidden_act
          _snake_case							=   use_absolute_embeddings
          _snake_case							=   patch_norm
          _snake_case							=   layer_norm_eps
          _snake_case							=   initializer_range
          _snake_case							=   is_training
          _snake_case							=   scope
          _snake_case							=   use_labels
          _snake_case							=   upscale

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






     def        _lowercase			(				self					)							->    Any:
          _snake_case							=   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

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

          _snake_case							=   self.get_config()

          return config, pixel_values, labels






     def        _lowercase			(				self					)							->    Union[str, Any]:
          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        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)							->    Optional[Any]:
          _snake_case							=   SwinaSRModel(config=_SCREAMING_SNAKE_CASE					)
          model.to(_SCREAMING_SNAKE_CASE					)
          model.eval()
          _snake_case							=   model(_SCREAMING_SNAKE_CASE					)

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






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)							->    int:
          _snake_case							=   SwinaSRForImageSuperResolution(_SCREAMING_SNAKE_CASE					)
          model.to(_SCREAMING_SNAKE_CASE					)
          model.eval()
          _snake_case							=   model(_SCREAMING_SNAKE_CASE					)

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






     def        _lowercase			(				self					)							->    List[str]:
          _snake_case							=   self.prepare_config_and_inputs()
          _snake_case							,      _snake_case							,      _snake_case							=   config_and_inputs
          _snake_case							=   {"pixel_values": pixel_values}
          return config, inputs_dict







@require_torch
class      _a       (		__lowerCAmelCase      ,      __lowerCAmelCase      ,      unittest.TestCase							):
     SCREAMING_SNAKE_CASE_							:					Union[str, Any]					       =		(SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
     SCREAMING_SNAKE_CASE_							:					Optional[int]					       =		{"""feature-extraction""": SwinaSRModel} if is_torch_available() else {}

     SCREAMING_SNAKE_CASE_							:					str					       =		False
     SCREAMING_SNAKE_CASE_							:					List[str]					       =		False
     SCREAMING_SNAKE_CASE_							:					int					       =		False
     SCREAMING_SNAKE_CASE_							:					str					       =		False
     SCREAMING_SNAKE_CASE_							:					Tuple					       =		False






     def        _lowercase			(				self					)							->    Any:
          _snake_case							=   SwinaSRModelTester(self					)
          _snake_case							=   ConfigTester(self      ,config_class=_SCREAMING_SNAKE_CASE      ,embed_dim=37					)






     def        _lowercase			(				self					)							->    Tuple:
          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]:
          _snake_case							=   self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    List[str]:
          _snake_case							=   self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_super_resolution(*_SCREAMING_SNAKE_CASE					)






     @unittest.skip(reason="Got `CUDA error: misaligned address` with PyTorch 2.0.0."					)
     def        _lowercase			(				self					)							->    List[str]:
          pass






     @unittest.skip(reason="Swin2SR does not use inputs_embeds"					)
     def        _lowercase			(				self					)							->    Optional[Any]:
          pass






     @unittest.skip(reason="Swin2SR does not support training yet"					)
     def        _lowercase			(				self					)							->    str:
          pass






     @unittest.skip(reason="Swin2SR does not support training yet"					)
     def        _lowercase			(				self					)							->    List[str]:
          pass






     def        _lowercase			(				self					)							->    int:
          _snake_case							,      _snake_case							=   self.model_tester.prepare_config_and_inputs_for_common()

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






     def        _lowercase			(				self					)							->    str:
          _snake_case							,      _snake_case							=   self.model_tester.prepare_config_and_inputs_for_common()

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

               _snake_case							=   ["pixel_values"]
               self.assertListEqual(arg_names[:1]      ,_SCREAMING_SNAKE_CASE					)






     @slow
     def        _lowercase			(				self					)							->    Any:
          for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _snake_case							=   SwinaSRModel.from_pretrained(_SCREAMING_SNAKE_CASE					)
               self.assertIsNotNone(_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Dict:
          _snake_case							,      _snake_case							=   self.model_tester.prepare_config_and_inputs_for_common()

          _snake_case							=   _config_zero_init(_SCREAMING_SNAKE_CASE					)
          for model_class in self.all_model_classes:
               _snake_case							=   model_class(config=_SCREAMING_SNAKE_CASE					)
               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        _lowercase			(				self					)							->    List[str]:
          _snake_case							,      _snake_case							=   self.model_tester.prepare_config_and_inputs_for_common()
          _snake_case							=   True

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

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

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

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

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

               _snake_case							=   outputs.attentions

               self.assertEqual(len(_SCREAMING_SNAKE_CASE					)      ,_SCREAMING_SNAKE_CASE					)

               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      _a       (		unittest.TestCase							):






     def        _lowercase			(				self					)							->    List[Any]:
          _snake_case							=   SwinaSRImageProcessor()
          _snake_case							=   SwinaSRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64"					).to(_SCREAMING_SNAKE_CASE					)

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

          # forward pass
          with torch.no_grad():
               _snake_case							=   model(**_SCREAMING_SNAKE_CASE					)

          # verify the logits
          _snake_case							=   torch.Size([1, 3, 976, 1_296]					)
          self.assertEqual(outputs.reconstruction.shape      ,_SCREAMING_SNAKE_CASE					)
          _snake_case							=   torch.tensor(
              [[0.5_4_5_8, 0.5_5_4_6, 0.5_6_3_8], [0.5_5_2_6, 0.5_5_6_5, 0.5_6_5_1], [0.5_3_9_6, 0.5_4_2_6, 0.5_6_2_1]]					).to(_SCREAMING_SNAKE_CASE					)
          self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]      ,_SCREAMING_SNAKE_CASE      ,atol=1e-4					)					)








'''simple docstring'''





import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


UpperCamelCase_		:    List[str]									=		500000

UpperCamelCase_			,			UpperCamelCase_		:    Dict									=		os.path.split(__file__)
UpperCamelCase_		:    int									=		os.path.join(RESULTS_BASEPATH, '''results''', RESULTS_FILENAME.replace('''.py''', '''.json'''))





@get_duration
def 			__a			(     _UpperCamelCase:					datasets.Dataset					)		->				List[Any]:

     """simple docstring"""




     _snake_case							=   dataset.select(range(0      ,		len(_UpperCamelCase					)      ,		2					)					)





@get_duration
def 			__a			(     _UpperCamelCase:					datasets.Dataset					)		->				List[Any]:

     """simple docstring"""




     _snake_case							=   dataset.sort("numbers"					)





@get_duration
def 			__a			(     _UpperCamelCase:					datasets.Dataset					)		->				List[Any]:

     """simple docstring"""




     _snake_case							=   dataset.shuffle()





@get_duration
def 			__a			(     _UpperCamelCase:					datasets.Dataset					)		->				Any:

     """simple docstring"""




     _snake_case							=   dataset.train_test_split(0.1					)





@get_duration
def 			__a			(     _UpperCamelCase:					datasets.Dataset      ,		_UpperCamelCase:					int=10					)		->				List[str]:

     """simple docstring"""




     for shard_id in range(_UpperCamelCase					):
          _snake_case							=   dataset.shard(_UpperCamelCase      ,		_UpperCamelCase					)





def 			__a			(     )		->				int:

     """simple docstring"""




     _snake_case							=   {"num examples": SPEED_TEST_N_EXAMPLES}
     _snake_case							=   (select, sort, shuffle, train_test_split, shard)
     with tempfile.TemporaryDirectory() as tmp_dir:
          print("generating dataset"					)
          _snake_case							=   datasets.Features({"text": datasets.Value("string"					), "numbers": datasets.Value("float32"					)}					)
          _snake_case							=   generate_example_dataset(
              os.path.join(_UpperCamelCase      ,		"dataset.arrow"					)      ,		_UpperCamelCase      ,		num_examples=_UpperCamelCase					)
          print("Functions"					)
          for func in functions:
               print(func.__name__					)
               _snake_case							=   func(_UpperCamelCase					)

     with open(_UpperCamelCase      ,		"wb"					) as f:
          f.write(json.dumps(_UpperCamelCase					).encode("utf-8"					)					)


if __name__ == "__main__":  # useful to run the profiler
   benchmark_indices_mapping()


