





import inspect
import unittest

from transformers import ConvNextVaConfig
from transformers.models.auto import get_values
from transformers.models.auto.modeling_auto import MODEL_FOR_BACKBONE_MAPPING_NAMES, MODEL_MAPPING_NAMES
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_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import ConvNextVaBackbone, ConvNextVaForImageClassification, ConvNextVaModel
    from transformers.models.convnextva.modeling_convnextva import CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor



class 					A__   :





      def __init__( self       ,							A_       ,							A_=13       ,							A_=32       ,							A_=3       ,							A_=4       ,							A_=[10, 20, 30, 40]       ,							A_=[2, 2, 3, 2]       ,							A_=True       ,							A_=True       ,							A_=37       ,							A_="gelu"       ,							A_=10       ,							A_=0.02       ,							A_=["stage2", "stage3", "stage4"]       ,							A_=[2, 3, 4]       ,							A_=None       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   parent
       UpperCamelCase	:					Union[str, Any]       	=   batch_size
       UpperCamelCase	:					int       	=   image_size
       UpperCamelCase	:					Optional[Any]       	=   num_channels
       UpperCamelCase	:					Tuple       	=   num_stages
       UpperCamelCase	:					List[Any]       	=   hidden_sizes
       UpperCamelCase	:					Union[str, Any]       	=   depths
       UpperCamelCase	:					List[str]       	=   is_training
       UpperCamelCase	:					str       	=   use_labels
       UpperCamelCase	:					Optional[Any]       	=   intermediate_size
       UpperCamelCase	:					Optional[Any]       	=   hidden_act
       UpperCamelCase	:					Any       	=   num_labels
       UpperCamelCase	:					str       	=   initializer_range
       UpperCamelCase	:					Optional[Any]       	=   out_features
       UpperCamelCase	:					Any       	=   out_indices
       UpperCamelCase	:					List[Any]       	=   scope





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					List[str]       	=   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

       UpperCamelCase	:					Any       	=   None
       if self.use_labels:
        UpperCamelCase	:					Optional[Any]       	=   ids_tensor([self.batch_size]       ,							self.num_labels   )

       UpperCamelCase	:					Optional[int]       	=   self.get_config()

       return config, pixel_values, labels





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       return ConvNextVaConfig(
           num_channels=self.num_channels       ,							hidden_sizes=self.hidden_sizes       ,							depths=self.depths       ,							num_stages=self.num_stages       ,							hidden_act=self.hidden_act       ,							is_decoder=A_       ,							initializer_range=self.initializer_range       ,							out_features=self.out_features       ,							out_indices=self.out_indices       ,							num_labels=self.num_labels       ,							)





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   ConvNextVaModel(config=A_   )
       model.to(A_   )
       model.eval()
       UpperCamelCase	:					List[str]       	=   model(A_   )
       # expected last hidden states: B, C, H // 32, W // 32
       self.parent.assertEqual(
           result.last_hidden_state.shape       ,							(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)       ,							)





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   ConvNextVaForImageClassification(A_   )
       model.to(A_   )
       model.eval()
       UpperCamelCase	:					Optional[Any]       	=   model(A_       ,							labels=A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.num_labels)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					Union[str, Any]       	=   ConvNextVaBackbone(config=A_   )
       model.to(A_   )
       model.eval()
       UpperCamelCase	:					Tuple       	=   model(A_   )

       # verify hidden states
       self.parent.assertEqual(len(result.feature_maps   )       ,							len(config.out_features   )   )
       self.parent.assertListEqual(list(result.feature_maps[0].shape   )       ,							[self.batch_size, self.hidden_sizes[1], 4, 4]   )

       # 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
       UpperCamelCase	:					Union[str, Any]       	=   None
       UpperCamelCase	:					Optional[Any]       	=   ConvNextVaBackbone(config=A_   )
       model.to(A_   )
       model.eval()
       UpperCamelCase	:					Optional[int]       	=   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.hidden_sizes[-1], 1, 1]   )

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





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Dict       	=   self.prepare_config_and_inputs()
       UpperCamelCase	:					Any       	=   config_and_inputs
       UpperCamelCase	:					Union[str, Any]       	=   {"""pixel_values""": pixel_values}
       return config, inputs_dict





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   self.prepare_config_and_inputs()
       UpperCamelCase	:					List[Any]       	=   config_and_inputs
       UpperCamelCase	:					Tuple       	=   {"""pixel_values""": pixel_values, """labels""": labels}
       return config, inputs_dict



@require_torch
class 					A__   (       UpperCAmelCase__       ,   UpperCAmelCase__       ,   unittest.TestCase			):
      _UpperCAmelCase	:Optional[int]						    =   (
          (
              ConvNextVaModel,
              ConvNextVaForImageClassification,
              ConvNextVaBackbone,
          )
          if is_torch_available()
          else ()
      )
      _UpperCAmelCase	:Tuple						    =   (
          {"feature-extraction": ConvNextVaModel, "image-classification": ConvNextVaForImageClassification}
          if is_torch_available()
          else {}
      )

      _UpperCAmelCase	:Any						    =   False
      _UpperCAmelCase	:Any						    =   False
      _UpperCAmelCase	:Optional[int]						    =   False
      _UpperCAmelCase	:Any						    =   False
      _UpperCAmelCase	:List[Any]						    =   False





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[int]       	=   ConvNextVaModelTester(self   )
       UpperCamelCase	:					List[str]       	=   ConfigTester(self       ,							config_class=A_       ,							has_text_modality=A_       ,							hidden_size=37   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       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   ):
       '''simple docstring'''




       return





      @unittest.skip(reason="ConvNextV2 does not use inputs_embeds"   )
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       pass





      @unittest.skip(reason="ConvNextV2 does not support input and output embeddings"   )
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       pass





      @unittest.skip(reason="ConvNextV2 does not use feedforward chunking"   )
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       pass





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       if not self.model_tester.is_training:
        return

       for model_class in self.all_model_classes:
        UpperCamelCase	:					Any       	=   self.model_tester.prepare_config_and_inputs_with_labels()
        UpperCamelCase	:					Dict       	=   True

        if model_class.__name__ in [
            *get_values(A_   ),
            *get_values(A_   ),
        ]:
         continue

        UpperCamelCase	:					Union[str, Any]       	=   model_class(A_   )
        model.to(A_   )
        model.train()
        UpperCamelCase	:					Optional[Any]       	=   self._prepare_for_class(A_       ,							A_       ,							return_labels=A_   )
        UpperCamelCase	:					List[Any]       	=   model(**A_   ).loss
        loss.backward()





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       if not self.model_tester.is_training:
        return

       for model_class in self.all_model_classes:
        UpperCamelCase	:					Any       	=   self.model_tester.prepare_config_and_inputs_with_labels()
        UpperCamelCase	:					List[str]       	=   False
        UpperCamelCase	:					Any       	=   True

        if (
            model_class.__name__
            in [*get_values(A_   ), *get_values(A_   )]
            or not model_class.supports_gradient_checkpointing
        ):
         continue

        UpperCamelCase	:					List[str]       	=   model_class(A_   )
        model.to(A_   )
        model.gradient_checkpointing_enable()
        model.train()
        UpperCamelCase	:					Union[str, Any]       	=   self._prepare_for_class(A_       ,							A_       ,							return_labels=A_   )
        UpperCamelCase	:					Any       	=   model(**A_   ).loss
        loss.backward()





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Union[str, Any]       	=   self.model_tester.prepare_config_and_inputs_for_common()

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

        UpperCamelCase	:					List[str]       	=   ["""pixel_values"""]
        self.assertListEqual(arg_names[:1]       ,							A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					List[str]       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       def check_hidden_states_output(A_       ,							A_       ,							A_   ):
        UpperCamelCase	:					str       	=   model_class(A_   )
        model.to(A_   )
        model.eval()

        with torch.no_grad():
         UpperCamelCase	:					Tuple       	=   model(**self._prepare_for_class(A_       ,							A_   )   )

        UpperCamelCase	:					Optional[Any]       	=   outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

        UpperCamelCase	:					str       	=   self.model_tester.num_stages
        self.assertEqual(len(A_   )       ,							expected_num_stages + 1   )

        # ConvNextV2's feature maps are of shape (batch_size, num_channels, height, width)
        self.assertListEqual(
            list(hidden_states[0].shape[-2:]   )       ,							[self.model_tester.image_size // 4, self.model_tester.image_size // 4]       ,							)

       UpperCamelCase	:					Tuple       	=   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCamelCase	:					List[Any]       	=   True
        check_hidden_states_output(A_       ,							A_       ,							A_   )

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

        check_hidden_states_output(A_       ,							A_       ,							A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					str       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*A_   )





      @slow
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       for model_name in CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCamelCase	:					Dict       	=   ConvNextVaModel.from_pretrained(A_   )
        self.assertIsNotNone(A_   )



def 		A_				(     )						-> List[str]:
 UpperCamelCase	:					Any       	=   Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"     )
 return image



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





      @cached_property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       return AutoImageProcessor.from_pretrained("facebook/convnextv2-tiny-1k-224"   ) if is_vision_available() else None





      @slow
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   ConvNextVaForImageClassification.from_pretrained("facebook/convnextv2-tiny-1k-224"   ).to(A_   )

       UpperCamelCase	:					Tuple       	=   self.default_image_processor
       UpperCamelCase	:					Optional[int]       	=   prepare_img()
       UpperCamelCase	:					List[Any]       	=   preprocessor(images=A_       ,							return_tensors="pt"   ).to(A_   )

       # forward pass
       with torch.no_grad():
        UpperCamelCase	:					List[Any]       	=   model(**A_   )

       # verify the logits
       UpperCamelCase	:					Any       	=   torch.Size((1, 1000)   )
       self.assertEqual(outputs.logits.shape       ,							A_   )

       UpperCamelCase	:					Optional[Any]       	=   torch.tensor([0.99_96, 0.19_66, -0.43_86]   ).to(A_   )
       self.assertTrue(torch.allclose(outputs.logits[0, :3]       ,							A_       ,							atol=1e-4   )   )




'''simple docstring'''







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(_UpperCAmelCase    ,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_					(  _UpperCAmelCase	:	Union[str, Any]    ,_UpperCAmelCase	:	Dict  )  ->			Optional[Any]:
       _a    : Tuple     					=_distribute_shards(**_UpperCAmelCase  )
       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_					(  _UpperCAmelCase	:	Tuple    ,_UpperCAmelCase	:	List[Any]    ,_UpperCAmelCase	:	Union[str, Any]  )  ->			List[str]:
       _a    : List[str]     					=_split_gen_kwargs(_UpperCAmelCase    ,_UpperCAmelCase  )
       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_					(  _UpperCAmelCase	:	Optional[Any]    ,_UpperCAmelCase	:	List[Any]  )  ->			Union[str, Any]:
       if expected is RuntimeError:
              with pytest.raises(_UpperCAmelCase  ):
                     _number_of_shards_in_gen_kwargs(_UpperCAmelCase  )
       else:
              _a    : Dict     					=_number_of_shards_in_gen_kwargs(_UpperCAmelCase  )
              assert out == expected


