
'''simple docstring'''






import collections
import inspect
import unittest

from transformers import FocalNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
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 (
     FocalNetBackbone,
     FocalNetForImageClassification,
     FocalNetForMaskedImageModeling,
     FocalNetModel,
 )
 from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor

class 			UpperCAmelCase__    :



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



    def 	lowerCamelCase_	(						self   :							Any  ):
          _lowerCamelCase					: List[str]		    = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

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

          _lowerCamelCase					: Optional[Any]		    = self.get_config()

          return config, pixel_values, labels



    def 	lowerCamelCase_	(						self   :							Union[str, Any]  ):
          return FocalNetConfig(
              image_size=self.image_size,patch_size=self.patch_size,num_channels=self.num_channels,embed_dim=self.embed_dim,hidden_sizes=self.hidden_sizes,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,out_features=self.out_features,out_indices=self.out_indices,)



    def 	lowerCamelCase_	(						self   :							int,__A   :							Union[str, Any],__A   :							Tuple,__A   :							List[Any]  ):
          _lowerCamelCase					: Optional[Any]		    = FocalNetModel(config=__A  )
          model.to(__A  )
          model.eval()
          _lowerCamelCase					: Optional[Any]		    = model(__A  )

          _lowerCamelCase					: Optional[Any]		    = ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths  ) - 1))
          _lowerCamelCase					: Union[str, 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   :							int,__A   :							Optional[int],__A   :							int,__A   :							Optional[int]  ):
          _lowerCamelCase					: Any		    = FocalNetBackbone(config=__A  )
          model.to(__A  )
          model.eval()
          _lowerCamelCase					: List[str]		    = model(__A  )

          # verify feature maps
          self.parent.assertEqual(len(result.feature_maps  ),len(config.out_features  )  )
          self.parent.assertListEqual(list(result.feature_maps[0].shape  ),[self.batch_size, self.image_size, 8, 8]  )

          # verify channels
          self.parent.assertEqual(len(model.channels  ),len(config.out_features  )  )
          self.parent.assertListEqual(model.channels,config.hidden_sizes[:-1]  )

          # verify backbone works with out_features=None
          _lowerCamelCase					: List[str]		    = None
          _lowerCamelCase					: List[str]		    = FocalNetBackbone(config=__A  )
          model.to(__A  )
          model.eval()
          _lowerCamelCase					: str		    = model(__A  )

          # verify feature maps
          self.parent.assertEqual(len(result.feature_maps  ),1  )
          self.parent.assertListEqual(list(result.feature_maps[0].shape  ),[self.batch_size, self.image_size * 2, 4, 4]  )

          # verify channels
          self.parent.assertEqual(len(model.channels  ),1  )
          self.parent.assertListEqual(model.channels,[config.hidden_sizes[-1]]  )



    def 	lowerCamelCase_	(						self   :							Optional[int],__A   :							Optional[int],__A   :							Dict,__A   :							Dict  ):
          _lowerCamelCase					: List[Any]		    = FocalNetForMaskedImageModeling(config=__A  )
          model.to(__A  )
          model.eval()
          _lowerCamelCase					: List[str]		    = model(__A  )
          self.parent.assertEqual(
              result.reconstruction.shape,(self.batch_size, self.num_channels, self.image_size, self.image_size)  )

          # test greyscale images
          _lowerCamelCase					: Dict		    = 1
          _lowerCamelCase					: Any		    = FocalNetForMaskedImageModeling(__A  )
          model.to(__A  )
          model.eval()

          _lowerCamelCase					: Dict		    = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
          _lowerCamelCase					: Optional[int]		    = model(__A  )
          self.parent.assertEqual(result.reconstruction.shape,(self.batch_size, 1, self.image_size, self.image_size)  )



    def 	lowerCamelCase_	(						self   :							List[Any],__A   :							Union[str, Any],__A   :							List[Any],__A   :							Optional[Any]  ):
          _lowerCamelCase					: Union[str, Any]		    = self.type_sequence_label_size
          _lowerCamelCase					: Optional[Any]		    = FocalNetForImageClassification(__A  )
          model.to(__A  )
          model.eval()
          _lowerCamelCase					: Optional[int]		    = model(__A,labels=__A  )
          self.parent.assertEqual(result.logits.shape,(self.batch_size, self.type_sequence_label_size)  )

          # test greyscale images
          _lowerCamelCase					: str		    = 1
          _lowerCamelCase					: str		    = FocalNetForImageClassification(__A  )
          model.to(__A  )
          model.eval()

          _lowerCamelCase					: str		    = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
          _lowerCamelCase					: List[Any]		    = model(__A  )
          self.parent.assertEqual(result.logits.shape,(self.batch_size, self.type_sequence_label_size)  )



    def 	lowerCamelCase_	(						self   :							Optional[int]  ):
          _lowerCamelCase					: int		    = self.prepare_config_and_inputs()

          _lowerCamelCase    ,		_lowerCamelCase    ,		_lowerCamelCase					: Union[str, Any]		    = config_and_inputs
          _lowerCamelCase					: Union[str, Any]		    = {"pixel_values": pixel_values}
          return config, inputs_dict


@require_torch
class 			UpperCAmelCase__    (							A    ,						A    ,						unittest.TestCase			):
    lowerCAmelCase_					     =		(
        (
            FocalNetModel,
            FocalNetForImageClassification,
            FocalNetForMaskedImageModeling,
            FocalNetBackbone,
        )
        if is_torch_available()
        else ()
    )
    lowerCAmelCase_					     =		(
        {'feature-extraction': FocalNetModel, 'image-classification': FocalNetForImageClassification}
        if is_torch_available()
        else {}
    )
    lowerCAmelCase_					     =		False

    lowerCAmelCase_					     =		False
    lowerCAmelCase_					     =		False
    lowerCAmelCase_					     =		False
    lowerCAmelCase_					     =		False



    def 	lowerCamelCase_	(						self   :							int  ):
          _lowerCamelCase					: Optional[int]		    = FocalNetModelTester(self  )
          _lowerCamelCase					: int		    = ConfigTester(self,config_class=__A,embed_dim=3_7,has_text_modality=__A  )



    def 	lowerCamelCase_	(						self   :							Union[str, 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   :							List[str]  ):
          return



    def 	lowerCamelCase_	(						self   :							Any  ):
          _lowerCamelCase					: Any		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*__A  )



    def 	lowerCamelCase_	(						self   :							int  ):
          _lowerCamelCase					: Dict		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_backbone(*__A  )



    def 	lowerCamelCase_	(						self   :							Union[str, Any]  ):
          _lowerCamelCase					: Any		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_image_modeling(*__A  )



    def 	lowerCamelCase_	(						self   :							int  ):
          _lowerCamelCase					: Optional[Any]		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*__A  )



    @unittest.skip(reason="FocalNet does not use inputs_embeds"  )
    def 	lowerCamelCase_	(						self   :							Optional[int]  ):
          pass



    @unittest.skip(reason="FocalNet does not use feedforward chunking"  )
    def 	lowerCamelCase_	(						self   :							List[str]  ):
          pass



    def 	lowerCamelCase_	(						self   :							List[str]  ):
          _lowerCamelCase    ,		_lowerCamelCase					: List[str]		    = self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes[:-1]:
                _lowerCamelCase					: str		    = model_class(__A  )
                self.assertIsInstance(model.get_input_embeddings(),(nn.Module)  )
                _lowerCamelCase					: Any		    = model.get_output_embeddings()
                self.assertTrue(x is None or isinstance(__A,nn.Linear  )  )



    def 	lowerCamelCase_	(						self   :							List[Any]  ):
          _lowerCamelCase    ,		_lowerCamelCase					: Union[str, Any]		    = self.model_tester.prepare_config_and_inputs_for_common()

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

                _lowerCamelCase					: Union[str, Any]		    = ["pixel_values"]
                self.assertListEqual(arg_names[:1],__A  )



    def 	lowerCamelCase_	(						self   :							Tuple,__A   :							Any,__A   :							List[Any],__A   :							str,__A   :							Any  ):
          _lowerCamelCase					: Union[str, Any]		    = model_class(__A  )
          model.to(__A  )
          model.eval()

          with torch.no_grad():
                _lowerCamelCase					: Optional[int]		    = model(**self._prepare_for_class(__A,__A  )  )

          _lowerCamelCase					: Optional[int]		    = outputs.hidden_states

          _lowerCamelCase					: int		    = getattr(
              self.model_tester,"expected_num_hidden_layers",len(self.model_tester.depths  ) + 1  )
          self.assertEqual(len(__A  ),__A  )

          # FocalNet has a different seq_length
          _lowerCamelCase					: Optional[Any]		    = (
              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					: Any		    = outputs.reshaped_hidden_states
          self.assertEqual(len(__A  ),__A  )

          _lowerCamelCase    ,		_lowerCamelCase    ,		_lowerCamelCase    ,		_lowerCamelCase					: Tuple		    = reshaped_hidden_states[0].shape
          _lowerCamelCase					: List[str]		    = (
              reshaped_hidden_states[0].view(__A,__A,height * width  ).permute(0,2,1  )
          )
          self.assertListEqual(
              list(reshaped_hidden_states.shape[-2:]  ),[num_patches, self.model_tester.embed_dim],)



    def 	lowerCamelCase_	(						self   :							Union[str, Any]  ):
          _lowerCamelCase    ,		_lowerCamelCase					: Tuple		    = self.model_tester.prepare_config_and_inputs_for_common()

          _lowerCamelCase					: 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[:-1]:
                _lowerCamelCase					: List[Any]		    = True
                self.check_hidden_states_output(__A,__A,__A,__A  )

                # check that output_hidden_states also work using config
                del inputs_dict["output_hidden_states"]
                _lowerCamelCase					: List[Any]		    = True

                self.check_hidden_states_output(__A,__A,__A,__A  )



    def 	lowerCamelCase_	(						self   :							Optional[Any]  ):
          _lowerCamelCase    ,		_lowerCamelCase					: Dict		    = self.model_tester.prepare_config_and_inputs_for_common()
          _lowerCamelCase					: Tuple		    = 3

          _lowerCamelCase					: Optional[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)
          )
          _lowerCamelCase					: Tuple		    = (
              config.patch_size
              if isinstance(config.patch_size,collections.abc.Iterable  )
              else (config.patch_size, config.patch_size)
          )

          _lowerCamelCase					: Any		    = image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
          _lowerCamelCase					: int		    = image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

          for model_class in self.all_model_classes[:-1]:
                _lowerCamelCase					: List[Any]		    = True
                self.check_hidden_states_output(__A,__A,__A,(padded_height, padded_width)  )

                # check that output_hidden_states also work using config
                del inputs_dict["output_hidden_states"]
                _lowerCamelCase					: Optional[Any]		    = True
                self.check_hidden_states_output(__A,__A,__A,(padded_height, padded_width)  )



    @slow
    def 	lowerCamelCase_	(						self   :							Tuple  ):
          for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                _lowerCamelCase					: Dict		    = FocalNetModel.from_pretrained(__A  )
                self.assertIsNotNone(__A  )



    def 	lowerCamelCase_	(						self   :							Tuple  ):
          _lowerCamelCase    ,		_lowerCamelCase					: int		    = self.model_tester.prepare_config_and_inputs_for_common()

          _lowerCamelCase					: Optional[Any]		    = _config_zero_init(__A  )
          for model_class in self.all_model_classes:
                _lowerCamelCase					: Any		    = model_class(config=__A  )
                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',)


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



    @cached_property
    def 	lowerCamelCase_	(						self   :							Union[str, Any]  ):
          # TODO update organization
          return AutoImageProcessor.from_pretrained("microsoft/focalnet-tiny"  ) if is_vision_available() else None



    @slow
    def 	lowerCamelCase_	(						self   :							Union[str, Any]  ):
          _lowerCamelCase					: Any		    = FocalNetForImageClassification.from_pretrained("microsoft/focalnet-tiny"  ).to(__A  )
          _lowerCamelCase					: int		    = self.default_image_processor

          _lowerCamelCase					: Union[str, Any]		    = Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"  )
          _lowerCamelCase					: Dict		    = image_processor(images=__A,return_tensors="pt"  ).to(__A  )

          # forward pass
          with torch.no_grad():
                _lowerCamelCase					: Dict		    = model(**__A  )

          # verify the logits
          _lowerCamelCase					: List[Any]		    = torch.Size((1, 1_0_0_0)  )
          self.assertEqual(outputs.logits.shape,__A  )
          _lowerCamelCase					: List[str]		    = torch.tensor([0.2166, -0.4368, 0.2191]  ).to(__A  )
          self.assertTrue(torch.allclose(outputs.logits[0, :3],__A,atol=1e-4  )  )
          self.assertTrue(outputs.logits.argmax(dim=-1  ).item(),2_8_1  )


@require_torch
class 			UpperCAmelCase__    (							A    ,						unittest.TestCase			):
    lowerCAmelCase_					     =		(FocalNetBackbone,) if is_torch_available() else ()
    lowerCAmelCase_					     =		FocalNetConfig

    lowerCAmelCase_					     =		False



    def 	lowerCamelCase_	(						self   :							int  ):
          _lowerCamelCase					: int		    = FocalNetModelTester(self  )



import pytest

from datasets.utils.sharding import _distribute_shards, _number_of_shards_in_gen_kwargs, _split_gen_kwargs







@pytest.mark.parametrize(
    '''kwargs, expected'''							,   [
        ({'''num_shards''': 0, '''max_num_jobs''': 1}, []),
        ({'''num_shards''': 10, '''max_num_jobs''': 1}, [range(10				)]),
        ({'''num_shards''': 10, '''max_num_jobs''': 10}, [range(__lowerCAmelCase							,   i + 1				) for i in range(10				)]),
        ({'''num_shards''': 1, '''max_num_jobs''': 10}, [range(1				)]),
        ({'''num_shards''': 10, '''max_num_jobs''': 3}, [range(0							,   4				), range(4							,   7				), range(7							,   10				)]),
        ({'''num_shards''': 3, '''max_num_jobs''': 10}, [range(0							,   1				), range(1							,   2				), range(2							,   3				)]),
    ]							,   )
def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase				)    ->				Optional[int]:
  snake_case__						= _distribute_shards(**__lowerCAmelCase				)
  assert out == expected







@pytest.mark.parametrize(
    '''gen_kwargs, max_num_jobs, expected'''							,   [
        ({'''foo''': 0}, 10, [{'''foo''': 0}]),
        ({'''shards''': [0, 1, 2, 3]}, 1, [{'''shards''': [0, 1, 2, 3]}]),
        ({'''shards''': [0, 1, 2, 3]}, 4, [{'''shards''': [0]}, {'''shards''': [1]}, {'''shards''': [2]}, {'''shards''': [3]}]),
        ({'''shards''': [0, 1]}, 4, [{'''shards''': [0]}, {'''shards''': [1]}]),
        ({'''shards''': [0, 1, 2, 3]}, 2, [{'''shards''': [0, 1]}, {'''shards''': [2, 3]}]),
    ]							,   )
def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase				)    ->				Dict:
  snake_case__						= _split_gen_kwargs(__lowerCAmelCase							,   __lowerCAmelCase				)
  assert out == expected







@pytest.mark.parametrize(
    '''gen_kwargs, expected'''							,   [
        ({'''foo''': 0}, 1),
        ({'''shards''': [0]}, 1),
        ({'''shards''': [0, 1, 2, 3]}, 4),
        ({'''shards''': [0, 1, 2, 3], '''foo''': 0}, 4),
        ({'''shards''': [0, 1, 2, 3], '''other''': (0, 1)}, 4),
        ({'''shards''': [0, 1, 2, 3], '''shards2''': [0, 1]}, RuntimeError),
    ]							,   )
def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase				)    ->				List[Any]:
  if expected is RuntimeError:
    with pytest.raises(__lowerCAmelCase				):
      _number_of_shards_in_gen_kwargs(__lowerCAmelCase				)
  else:
    snake_case__						= _number_of_shards_in_gen_kwargs(__lowerCAmelCase				)
    assert out == expected

