






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       :



      """simple docstring"""


      def __init__( self   :     str						, __lowercase   :     Dict						, __lowercase   :     Optional[Any]=13						, __lowercase   :     str=32						, __lowercase   :     Union[str, Any]=3						, __lowercase   :     int=4						, __lowercase   :     int=[10, 20, 30, 40]						, __lowercase   :     List[str]=[2, 2, 3, 2]						, __lowercase   :     Dict=True						, __lowercase   :     List[Any]=True						, __lowercase   :     str=37						, __lowercase   :     Tuple="gelu"						, __lowercase   :     int=10						, __lowercase   :     Dict=0.02						, __lowercase   :     Dict=["stage2", "stage3", "stage4"]						, __lowercase   :     Optional[int]=[2, 3, 4]						, __lowercase   :     List[Any]=None						, )				-> Dict:
        __UpperCAmelCase				:      List[str]	 =    parent
        __UpperCAmelCase				:      str	 =    batch_size
        __UpperCAmelCase				:      Optional[Any]	 =    image_size
        __UpperCAmelCase				:      Union[str, Any]	 =    num_channels
        __UpperCAmelCase				:      Any	 =    num_stages
        __UpperCAmelCase				:      str	 =    hidden_sizes
        __UpperCAmelCase				:      Optional[int]	 =    depths
        __UpperCAmelCase				:      List[str]	 =    is_training
        __UpperCAmelCase				:      Optional[int]	 =    use_labels
        __UpperCAmelCase				:      str	 =    intermediate_size
        __UpperCAmelCase				:      int	 =    hidden_act
        __UpperCAmelCase				:      List[Any]	 =    num_labels
        __UpperCAmelCase				:      Tuple	 =    initializer_range
        __UpperCAmelCase				:      List[str]	 =    out_features
        __UpperCAmelCase				:      Tuple	 =    out_indices
        __UpperCAmelCase				:      List[Any]	 =    scope
      def       UpperCAmelCase ( self   :     List[str]   )				-> int:
        __UpperCAmelCase				:      List[Any]	 =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

        __UpperCAmelCase				:      Union[str, Any]	 =    None
        if self.use_labels:
          __UpperCAmelCase				:      Tuple	 =    ids_tensor([self.batch_size]						, self.num_labels   )

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

        return config, pixel_values, labels
      def       UpperCAmelCase ( self   :     Dict   )				-> str:
        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=__lowercase						, initializer_range=self.initializer_range						, out_features=self.out_features						, out_indices=self.out_indices						, num_labels=self.num_labels						, )
      def       UpperCAmelCase ( self   :     Tuple						, __lowercase   :     Tuple						, __lowercase   :     Union[str, Any]						, __lowercase   :     Any   )				-> List[Any]:
        __UpperCAmelCase				:      Tuple	 =    ConvNextVaModel(config=__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      Optional[int]	 =    model(__lowercase   )
        # 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   :     Optional[int]						, __lowercase   :     List[Any]						, __lowercase   :     Optional[Any]						, __lowercase   :     List[Any]   )				-> Tuple:
        __UpperCAmelCase				:      Optional[int]	 =    ConvNextVaForImageClassification(__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      Union[str, Any]	 =    model(__lowercase						, labels=__lowercase   )
        self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_labels)   )
      def       UpperCAmelCase ( self   :     Union[str, Any]						, __lowercase   :     Optional[int]						, __lowercase   :     Tuple						, __lowercase   :     List[str]   )				-> Tuple:
        __UpperCAmelCase				:      int	 =    ConvNextVaBackbone(config=__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      List[Any]	 =    model(__lowercase   )

        # 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				:      List[str]	 =    None
        __UpperCAmelCase				:      List[str]	 =    ConvNextVaBackbone(config=__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      List[str]	 =    model(__lowercase   )

        # 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   :     List[str]   )				-> List[Any]:
        __UpperCAmelCase				:      Optional[int]	 =    self.prepare_config_and_inputs()
        __UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				:      Optional[int]	 =    config_and_inputs
        __UpperCAmelCase				:      List[str]	 =    {"""pixel_values""": pixel_values}
        return config, inputs_dict




      def       UpperCAmelCase ( self   :     Tuple   )				-> Tuple:
        __UpperCAmelCase				:      List[str]	 =    self.prepare_config_and_inputs()
        __UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				:      Union[str, Any]	 =    config_and_inputs
        __UpperCAmelCase				:      List[str]	 =    {"""pixel_values""": pixel_values, """labels""": labels}
        return config, inputs_dict


@require_torch
class        a       (					lowercase__				,			lowercase__				,			unittest.TestCase  ):



      """simple docstring"""


      a    :					int							=      (
          (
              ConvNextVaModel,
              ConvNextVaForImageClassification,
              ConvNextVaBackbone,
          )
          if is_torch_available()
          else ()
      )
      a    :					List[Any]							=      (
          {'feature-extraction': ConvNextVaModel, 'image-classification': ConvNextVaForImageClassification}
          if is_torch_available()
          else {}
      )

      a    :					List[str]							=      False
      a    :					Union[str, Any]							=      False
      a    :					Optional[Any]							=      False
      a    :					Optional[Any]							=      False
      a    :					Tuple							=      False
      def       UpperCAmelCase ( self   :     Optional[Any]   )				-> Optional[Any]:
        __UpperCAmelCase				:      Tuple	 =    ConvNextVaModelTester(self   )
        __UpperCAmelCase				:      str	 =    ConfigTester(self						, config_class=__lowercase						, has_text_modality=__lowercase						, hidden_size=37   )
      def       UpperCAmelCase ( self   :     int   )				-> Dict:
        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   :     Any   )				-> str:
        return
      @unittest.skip(reason="""ConvNextV2 does not use inputs_embeds"""   )
      def       UpperCAmelCase ( self   :     Optional[int]   )				-> Optional[int]:
        pass
      @unittest.skip(reason="""ConvNextV2 does not support input and output embeddings"""   )
      def       UpperCAmelCase ( self   :     Union[str, Any]   )				-> List[Any]:
        pass
      @unittest.skip(reason="""ConvNextV2 does not use feedforward chunking"""   )
      def       UpperCAmelCase ( self   :     Tuple   )				-> str:
        pass
      def       UpperCAmelCase ( self   :     List[str]   )				-> Optional[Any]:
        if not self.model_tester.is_training:
          return

        for model_class in self.all_model_classes:
          __UpperCAmelCase ,					__UpperCAmelCase				:      List[Any]	 =    self.model_tester.prepare_config_and_inputs_with_labels()
          __UpperCAmelCase				:      int	 =    True

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

          __UpperCAmelCase				:      List[str]	 =    model_class(__lowercase   )
          model.to(__lowercase   )
          model.train()
          __UpperCAmelCase				:      Dict	 =    self._prepare_for_class(__lowercase						, __lowercase						, return_labels=__lowercase   )
          __UpperCAmelCase				:      str	 =    model(**__lowercase   ).loss
          loss.backward()
      def       UpperCAmelCase ( self   :     Optional[int]   )				-> List[str]:
        if not self.model_tester.is_training:
          return

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

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

          __UpperCAmelCase				:      Tuple	 =    model_class(__lowercase   )
          model.to(__lowercase   )
          model.gradient_checkpointing_enable()
          model.train()
          __UpperCAmelCase				:      int	 =    self._prepare_for_class(__lowercase						, __lowercase						, return_labels=__lowercase   )
          __UpperCAmelCase				:      int	 =    model(**__lowercase   ).loss
          loss.backward()
      def       UpperCAmelCase ( self   :     str   )				-> List[Any]:
        __UpperCAmelCase ,					__UpperCAmelCase				:      List[str]	 =    self.model_tester.prepare_config_and_inputs_for_common()

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

          __UpperCAmelCase				:      Optional[int]	 =    ["""pixel_values"""]
          self.assertListEqual(arg_names[:1]						, __lowercase   )
      def       UpperCAmelCase ( self   :     Union[str, Any]   )				-> Union[str, Any]:
        __UpperCAmelCase				:      Tuple	 =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowercase   )
      def       UpperCAmelCase ( self   :     Dict   )				-> List[Any]:
        def check_hidden_states_output(__lowercase   :     List[Any]						, __lowercase   :     Optional[Any]						, __lowercase   :     Dict   ):
          __UpperCAmelCase				:      str	 =    model_class(__lowercase   )
          model.to(__lowercase   )
          model.eval()

          with torch.no_grad():
            __UpperCAmelCase				:      str	 =    model(**self._prepare_for_class(__lowercase						, __lowercase   )   )

          __UpperCAmelCase				:      List[Any]	 =    outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

          __UpperCAmelCase				:      Optional[int]	 =    self.model_tester.num_stages
          self.assertEqual(len(__lowercase   )						, 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 ,					__UpperCAmelCase				:      Optional[Any]	 =    self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          __UpperCAmelCase				:      Dict	 =    True
          check_hidden_states_output(__lowercase						, __lowercase						, __lowercase   )

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          __UpperCAmelCase				:      Optional[int]	 =    True

          check_hidden_states_output(__lowercase						, __lowercase						, __lowercase   )
      def       UpperCAmelCase ( self   :     Dict   )				-> str:
        __UpperCAmelCase				:      Any	 =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*__lowercase   )




      @slow
      def       UpperCAmelCase ( self   :     Union[str, Any]   )				-> Union[str, Any]:
        for model_name in CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          __UpperCAmelCase				:      Union[str, Any]	 =    ConvNextVaModel.from_pretrained(__lowercase   )
          self.assertIsNotNone(__lowercase   )






def 							lowerCamelCase__ (       ):
  __UpperCAmelCase				:      List[Any]	 =    Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""						)
  return image







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



      """simple docstring"""


      @cached_property
      def       UpperCAmelCase ( self   :     List[Any]   )				-> Optional[Any]:
        return AutoImageProcessor.from_pretrained("""facebook/convnextv2-tiny-1k-224"""   ) if is_vision_available() else None




      @slow
      def       UpperCAmelCase ( self   :     Dict   )				-> int:
        __UpperCAmelCase				:      Optional[int]	 =    ConvNextVaForImageClassification.from_pretrained("""facebook/convnextv2-tiny-1k-224"""   ).to(__lowercase   )

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

        # forward pass
        with torch.no_grad():
          __UpperCAmelCase				:      Union[str, Any]	 =    model(**__lowercase   )

        # verify the logits
        __UpperCAmelCase				:      Dict	 =    torch.Size((1, 1000)   )
        self.assertEqual(outputs.logits.shape						, __lowercase   )

        __UpperCAmelCase				:      int	 =    torch.tensor([0.9_996, 0.1_966, -0.4_386]   ).to(__lowercase   )
        self.assertTrue(torch.allclose(outputs.logits[0, :3]						, __lowercase						, atol=1e-4   )   )





import contextlib
from multiprocessing import Pool, RLock

from tqdm.auto import tqdm

from ..utils import experimental, logging


A_			:					Any    			=     logging.get_logger(__name__)



class        _a		:


					'''simple docstring'''






					UpperCAmelCase__:  str    				=			None



@experimental
def 		UpperCamelCase   (lowercase_:					int       , lowercase_:					Any       , lowercase_:					str       , lowercase_:					str       , lowercase_:					List[str]       , lowercase_:					List[Any]       , lowercase_:					Union[str, Any]   )	-> Union[str, Any]:
	if ParallelBackendConfig.backend_name is None:
		return _map_with_multiprocessing_pool(
		    lowercase_       , lowercase_       , lowercase_       , lowercase_       , lowercase_       , lowercase_       , lowercase_   )

	return _map_with_joblib(lowercase_       , lowercase_       , lowercase_       , lowercase_       , lowercase_       , lowercase_       , lowercase_   )

def 		UpperCamelCase   (lowercase_:					Tuple       , lowercase_:					List[Any]       , lowercase_:					Tuple       , lowercase_:					Tuple       , lowercase_:					List[str]       , lowercase_:					Dict       , lowercase_:					Tuple   )	-> Tuple:
	A__				:       Union[str, Any]	  =						num_proc if num_proc <= len(lowercase_   ) else len(lowercase_   )
	A__				:       str	  =						[]  # We organize the splits ourselve (contiguous splits)
	for index in range(lowercase_   ):
		A__				:       Tuple	  =						len(lowercase_   ) // num_proc
		A__				:       Tuple	  =						len(lowercase_   ) % num_proc
		A__				:       Optional[Any]	  =						div * index + min(lowercase_       , lowercase_   )
		A__				:       Union[str, Any]	  =						start + div + (1 if index < mod else 0)
		split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc)   )

	if len(lowercase_   ) != sum(len(i[1]   ) for i in split_kwds   ):
		raise ValueError(
		    f"""Error dividing inputs iterable among processes. """
		    f"""Total number of objects {len(lowercase_   )}, """
		    f"""length: {sum(len(i[1]   ) for i in split_kwds   )}"""   )

	logger.info(
	    f"""Spawning {num_proc} processes for {len(lowercase_   )} objects in slices of {[len(i[1]   ) for i in split_kwds]}"""   )
	A__  ,  A__				:       Optional[int]	  =						None, None
	if not disable_tqdm:
		A__  ,  A__				:       List[str]	  =						(RLock(),), tqdm.set_lock
	with Pool(lowercase_       , initargs=lowercase_       , initializer=lowercase_   ) as pool:
		A__				:       Optional[Any]	  =						pool.map(lowercase_       , lowercase_   )
	logger.info(f"""Finished {num_proc} processes"""   )
	A__				:       Tuple	  =						[obj for proc_res in mapped for obj in proc_res]
	logger.info(f"""Unpacked {len(lowercase_   )} objects"""   )

	return mapped

def 		UpperCamelCase   (lowercase_:					Union[str, Any]       , lowercase_:					Dict       , lowercase_:					str       , lowercase_:					Union[str, Any]       , lowercase_:					List[Any]       , lowercase_:					Dict       , lowercase_:					Optional[Any]   )	-> Union[str, Any]:
	# progress bar is not yet supported for _map_with_joblib, because tqdm couldn't accurately be applied to joblib,
	# and it requires monkey-patching joblib internal classes which is subject to change
	import joblib

	with joblib.parallel_backend(ParallelBackendConfig.backend_name       , n_jobs=lowercase_   ):
		return joblib.Parallel()(
		    joblib.delayed(lowercase_   )((function, obj, types, None, True, None)   ) for obj in iterable   )


@experimental
@contextlib.contextmanager
def 		UpperCamelCase   (lowercase_:					str   )	-> str:
	A__				:       int	  =						backend_name

	if backend_name == "spark":
		from joblibspark import register_spark

		register_spark()

		# TODO: call create_cache_and_write_probe if "download" in steps
		# TODO: raise NotImplementedError when Dataset.map etc is called

	try:
		yield
	finally:
		A__				:       Optional[Any]	  =						None



