



"""simple docstring"""
import inspect
import unittest

from transformers import DataaVecVisionConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, 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, _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 (
						    MODEL_MAPPING,
						    DataaVecVisionForImageClassification,
						    DataaVecVisionForSemanticSegmentation,
						    DataaVecVisionModel,
						)
						from transformers.models.dataavec.modeling_dataavec_vision import DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import BeitImageProcessor
class 							lowerCamelCase    :






					'''simple docstring'''





					def __init__(			self:				str		,	snake_case:				Union[str, Any]		,	snake_case:				Optional[int]=100		,	snake_case:				Tuple=13		,	snake_case:				Optional[Any]=30		,	snake_case:				Dict=2		,	snake_case:				int=3		,	snake_case:				Union[str, Any]=True		,	snake_case:				List[str]=True		,	snake_case:				List[Any]=32		,	snake_case:				Optional[Any]=4		,	snake_case:				str=4		,	snake_case:				Any=37		,	snake_case:				str="gelu"		,	snake_case:				Union[str, Any]=0.1		,	snake_case:				int=0.1		,	snake_case:				Optional[int]=10		,	snake_case:				Any=0.0_2		,	snake_case:				str=3		,	snake_case:				List[str]=None		,	snake_case:				int=[0, 1, 2, 3]		,	)  ->      Union[str, Any]:
											snake_case_					:Tuple		=     parent
											snake_case_					:Union[str, Any]		=     100
											snake_case_					:Union[str, Any]		=     batch_size
											snake_case_					:Union[str, Any]		=     image_size
											snake_case_					:Tuple		=     patch_size
											snake_case_					:Any		=     num_channels
											snake_case_					:str		=     is_training
											snake_case_					:Dict		=     use_labels
											snake_case_					:Any		=     hidden_size
											snake_case_					:Optional[Any]		=     num_hidden_layers
											snake_case_					:Optional[Any]		=     num_attention_heads
											snake_case_					:str		=     intermediate_size
											snake_case_					:List[str]		=     hidden_act
											snake_case_					:Any		=     hidden_dropout_prob
											snake_case_					:List[str]		=     attention_probs_dropout_prob
											snake_case_					:int		=     type_sequence_label_size
											snake_case_					:Optional[int]		=     initializer_range
											snake_case_					:Optional[Any]		=     scope
											snake_case_					:Any		=     out_indices
											snake_case_					:Dict		=     num_labels

											# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
											snake_case_					:Tuple		=     (image_size // patch_size) ** 2
											snake_case_					:Tuple		=     num_patches + 1






					def   lowerCAmelCase_							(			self:				str		)  ->      List[Any]:
											snake_case_					:str		=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

											snake_case_					:Union[str, Any]		=     None
											snake_case_					:Optional[Any]		=     None
											if self.use_labels:
																	snake_case_					:Tuple		=     ids_tensor([self.batch_size]		,	self.type_sequence_label_size		)
																	snake_case_					:List[Any]		=     ids_tensor([self.batch_size, self.image_size, self.image_size]		,	self.num_labels		)

											snake_case_					:Tuple		=     self.get_config()

											return config, pixel_values, labels, pixel_labels






					def   lowerCAmelCase_							(			self:				Union[str, Any]		)  ->      Dict:
											return DataaVecVisionConfig(
											    vocab_size=self.vocab_size		,	image_size=self.image_size		,	patch_size=self.patch_size		,	num_channels=self.num_channels		,	hidden_size=self.hidden_size		,	num_hidden_layers=self.num_hidden_layers		,	num_attention_heads=self.num_attention_heads		,	intermediate_size=self.intermediate_size		,	hidden_act=self.hidden_act		,	hidden_dropout_prob=self.hidden_dropout_prob		,	attention_probs_dropout_prob=self.attention_probs_dropout_prob		,	is_decoder=snake_case		,	initializer_range=self.initializer_range		,	out_indices=self.out_indices		,	)






					def   lowerCAmelCase_							(			self:				Optional[int]		,	snake_case:				Optional[int]		,	snake_case:				int		,	snake_case:				str		,	snake_case:				Dict		)  ->      int:
											snake_case_					:Union[str, Any]		=     DataaVecVisionModel(config=snake_case		)
											model.to(snake_case		)
											model.eval()
											snake_case_					:Optional[int]		=     model(snake_case		)
											# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
											snake_case_					:Dict		=     (self.image_size // self.patch_size) ** 2
											self.parent.assertEqual(result.last_hidden_state.shape		,	(self.batch_size, num_patches + 1, self.hidden_size)		)






					def   lowerCAmelCase_							(			self:				Optional[Any]		,	snake_case:				Tuple		,	snake_case:				Union[str, Any]		,	snake_case:				Tuple		,	snake_case:				List[Any]		)  ->      Tuple:
											snake_case_					:Union[str, Any]		=     self.type_sequence_label_size
											snake_case_					:int		=     DataaVecVisionForImageClassification(snake_case		)
											model.to(snake_case		)
											model.eval()
											snake_case_					:Any		=     model(snake_case		,	labels=snake_case		)
											self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.type_sequence_label_size)		)






					def   lowerCAmelCase_							(			self:				List[Any]		,	snake_case:				Any		,	snake_case:				Tuple		,	snake_case:				List[Any]		,	snake_case:				Optional[int]		)  ->      Tuple:
											snake_case_					:Tuple		=     self.num_labels
											snake_case_					:List[str]		=     DataaVecVisionForSemanticSegmentation(snake_case		)
											model.to(snake_case		)
											model.eval()
											snake_case_					:Tuple		=     model(snake_case		)
											self.parent.assertEqual(
											    result.logits.shape		,	(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)		)
											snake_case_					:List[str]		=     model(snake_case		,	labels=snake_case		)
											self.parent.assertEqual(
											    result.logits.shape		,	(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)		)






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

@require_torch
class 							lowerCamelCase    (     _lowerCAmelCase       ,		_lowerCAmelCase       ,		unittest.TestCase							):






					'''simple docstring'''





					_A						:     List[str]		=	(
					    (DataaVecVisionModel, DataaVecVisionForImageClassification, DataaVecVisionForSemanticSegmentation)
					    if is_torch_available()
					    else ()
					)
					_A						:     Union[str, Any]		=	(
					    {
					        """feature-extraction""": DataaVecVisionModel,
					        """image-classification""": DataaVecVisionForImageClassification,
					        """image-segmentation""": DataaVecVisionForSemanticSegmentation,
					    }
					    if is_torch_available()
					    else {}
					)

					_A						:     List[Any]		=	False
					_A						:     int		=	False
					_A						:     List[str]		=	False






					def   lowerCAmelCase_							(			self:				Any		)  ->      Any:
											snake_case_					:Union[str, Any]		=     DataaVecVisionModelTester(self		)
											snake_case_					:Union[str, Any]		=     ConfigTester(
											    self		,	config_class=snake_case		,	has_text_modality=snake_case		,	hidden_size=37		)






					def   lowerCAmelCase_							(			self:				Union[str, Any]		)  ->      Dict:
											self.config_tester.run_common_tests()






					def   lowerCAmelCase_							(			self:				int		)  ->      Union[str, Any]:
											# Data2VecVision does not use inputs_embeds
											pass






					@require_torch_multi_gpu
					@unittest.skip(
					    reason="""Data2VecVision has some layers using `add_module` which doesn't work well with `nn.DataParallel`"""		)
					def   lowerCAmelCase_							(			self:				Any		)  ->      Optional[Any]:
											pass






					def   lowerCAmelCase_							(			self:				List[Any]		)  ->      Optional[int]:
											snake_case_,			snake_case_					:Optional[Any]		=     self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	snake_case_					:Dict		=     model_class(snake_case		)
																	self.assertIsInstance(model.get_input_embeddings()		,	(nn.Module)		)
																	snake_case_					:Optional[int]		=     model.get_output_embeddings()
																	self.assertTrue(x is None or isinstance(snake_case		,	nn.Linear		)		)






					def   lowerCAmelCase_							(			self:				List[str]		)  ->      Dict:
											snake_case_,			snake_case_					:str		=     self.model_tester.prepare_config_and_inputs_for_common()

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

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






					def   lowerCAmelCase_							(			self:				int		)  ->      List[Any]:
											snake_case_					:int		=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*snake_case		)






					def   lowerCAmelCase_							(			self:				Optional[Any]		)  ->      Any:
											snake_case_					:Tuple		=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_segmentation(*snake_case		)






					def   lowerCAmelCase_							(			self:				Optional[int]		)  ->      List[str]:
											if not self.model_tester.is_training:
																	return

											snake_case_,			snake_case_					:List[str]		=     self.model_tester.prepare_config_and_inputs_for_common()
											snake_case_					:Optional[Any]		=     True

											for model_class in self.all_model_classes:
																	if model_class in [*get_values(snake_case		)]:
																							continue

																	snake_case_					:List[str]		=     model_class(snake_case		)
																	model.to(snake_case		)
																	model.train()
																	snake_case_					:Optional[Any]		=     self._prepare_for_class(snake_case		,	snake_case		,	return_labels=snake_case		)
																	snake_case_					:int		=     model(**snake_case		).loss
																	loss.backward()






					def   lowerCAmelCase_							(			self:				Tuple		)  ->      Tuple:
											snake_case_,			snake_case_					:List[str]		=     self.model_tester.prepare_config_and_inputs_for_common()
											if not self.model_tester.is_training:
																	return

											snake_case_					:List[Any]		=     False
											snake_case_					:Optional[int]		=     True

											for model_class in self.all_model_classes:
																	if model_class in [*get_values(snake_case		)] or not model_class.supports_gradient_checkpointing:
																							continue
																	# TODO: remove the following 3 lines once we have a MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING
																	# this can then be incorporated into _prepare_for_class in test_modeling_common.py
																	elif model_class.__name__ == "Data2VecVisionForSemanticSegmentation":
																							snake_case_,			snake_case_,			snake_case_,			snake_case_					:Tuple		=     inputs_dict["""pixel_values"""].shape
																							snake_case_					:List[Any]		=     torch.zeros(
																							    [self.model_tester.batch_size, height, width]		,	device=snake_case		).long()
																	snake_case_					:Dict		=     model_class(snake_case		)
																	model.gradient_checkpointing_enable()
																	model.to(snake_case		)
																	model.train()
																	snake_case_					:Any		=     self._prepare_for_class(snake_case		,	snake_case		,	return_labels=snake_case		)
																	snake_case_					:List[str]		=     model(**snake_case		).loss
																	loss.backward()






					def   lowerCAmelCase_							(			self:				List[str]		)  ->      Any:
											snake_case_,			snake_case_					:Dict		=     self.model_tester.prepare_config_and_inputs_for_common()

											snake_case_					:Union[str, Any]		=     _config_zero_init(snake_case		)
											for model_class in self.all_model_classes:
																	snake_case_					:List[Any]		=     model_class(config=snake_case		)
																	for name, param in model.named_parameters():
																							# we skip lambda parameters as these require special initial values
																							# determined by config.layer_scale_init_value
																							if "lambda" 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   lowerCAmelCase_							(			self:				List[str]		,	snake_case:				int		,	snake_case:				int		,	snake_case:				int		,	snake_case:				List[str]=2E-4		,	snake_case:				Any="outputs"		,	snake_case:				Tuple=None		)  ->      int:
											# We override with a slightly higher tol value, as semseg models tend to diverge a bit more
											super().check_pt_tf_outputs(snake_case		,	snake_case		,	snake_case		,	snake_case		,	snake_case		,	snake_case		)






					def   lowerCAmelCase_							(			self:				List[str]		)  ->      Optional[Any]:
											snake_case_					:str		=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*snake_case		)






					@slow
					def   lowerCAmelCase_							(			self:				Dict		)  ->      str:
											for model_name in DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	snake_case_					:Tuple		=     DataaVecVisionModel.from_pretrained(snake_case		)
																	self.assertIsNotNone(snake_case		)

def   A_    (     ):







						'''simple docstring'''


						snake_case_					:Dict		=     Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""      )
						return image

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






					'''simple docstring'''





					@cached_property
					def   lowerCAmelCase_							(			self:				List[str]		)  ->      List[Any]:
											return (
											    BeitImageProcessor.from_pretrained("""facebook/data2vec-vision-base-ft1k"""		) if is_vision_available() else None
											)






					@slow
					def   lowerCAmelCase_							(			self:				Tuple		)  ->      Dict:
											snake_case_					:Optional[int]		=     DataaVecVisionForImageClassification.from_pretrained("""facebook/data2vec-vision-base-ft1k"""		).to(
											    snake_case		)

											snake_case_					:Union[str, Any]		=     self.default_image_processor
											snake_case_					:Union[str, Any]		=     prepare_img()
											snake_case_					:Dict		=     image_processor(images=snake_case		,	return_tensors="""pt"""		).to(snake_case		)

											# forward pass
											with torch.no_grad():
																	snake_case_					:str		=     model(**snake_case		)
											snake_case_					:Union[str, Any]		=     outputs.logits

											# verify the logits
											snake_case_					:Union[str, Any]		=     torch.Size((1, 1_000)		)
											self.assertEqual(logits.shape		,	snake_case		)

											snake_case_					:int		=     torch.tensor([0.3_2_7_7, -0.1_3_9_5, 0.0_9_1_1]		).to(snake_case		)

											self.assertTrue(torch.allclose(logits[0, :3]		,	snake_case		,	atol=1E-4		)		)

											snake_case_					:List[str]		=     [model.config.labelaid[i] for i in ["""remote control, remote""", """tabby, tabby cat"""]]
											self.assertEqual(logits[0].topk(2		).indices.cpu().tolist()		,	snake_case		)

import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available







def     A__	(       __A ,			__A ,			__A ,			__A=None						):







    '''simple docstring'''
    _lowerCamelCase  :		List[Any]			    =    None
    # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
    if str(weight.dtype						) == "torch.bfloat16":
        # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
        _lowerCamelCase  :		int			    =    weight.view(torch.intaa						)
        _lowerCamelCase  :		List[Any]			    =    """bfloat16"""
    _lowerCamelCase  :		Tuple			    =    weight.cpu().numpy()
    _lowerCamelCase  :		Tuple			    =    os.path.join(__A ,			F"""{weight_name}.dat"""						)
    if index is not None:
        if dtype is None:
            _lowerCamelCase  :		Union[str, Any]			    =    str(array.dtype						)
        _lowerCamelCase  :		Optional[Any]			    =    {"""dtype""": dtype, """shape""": list(array.shape						)}
    if array.ndim == 0:
        _lowerCamelCase  :		int			    =    array[None]
    _lowerCamelCase  :		str			    =    np.memmap(__A ,			dtype=array.dtype ,			mode="""w+""" ,			shape=array.shape						)
    _lowerCamelCase  :		Union[str, Any]			    =    array[:]
    file_array.flush()
    return index







def     A__	(       __A ,			__A						):







    '''simple docstring'''
    _lowerCamelCase  :		Tuple			    =    tuple(weight_info["""shape"""]						)
    if shape == ():
        # NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
        _lowerCamelCase  :		Union[str, Any]			    =    (1,)

    _lowerCamelCase  :		str			    =    weight_info["""dtype"""]
    if dtype == "bfloat16":
        # NumPy does not support bfloat16 so this was saved as a int16
        _lowerCamelCase  :		Optional[int]			    =    """int16"""

    _lowerCamelCase  :		str			    =    np.memmap(__A ,			dtype=__A ,			shape=__A ,			mode="""r"""						)

    if len(weight_info["""shape"""]						) == 0:
        _lowerCamelCase  :		List[str]			    =    weight[0]
    _lowerCamelCase  :		List[str]			    =    torch.tensor(__A						)
    if weight_info["dtype"] == "bfloat16":
        _lowerCamelCase  :		List[Any]			    =    weight.view(torch.bfloataa						)

    return weight







def     A__	(       __A ,			__A						):







    '''simple docstring'''
    if index is None or len(__A						) == 0:
        # Nothing to save
        return

    _lowerCamelCase  :		List[Any]			    =    os.path.join(__A ,			"""index.json"""						)
    if os.path.isfile(__A						):
        with open(__A ,			"""r""" ,			encoding="""utf-8"""						) as f:
            _lowerCamelCase  :		List[str]			    =    json.load(__A						)
    else:
        _lowerCamelCase  :		List[str]			    =    {}
    current_index.update(__A						)

    with open(__A ,			"""w""" ,			encoding="""utf-8"""						) as f:
        json.dump(__A ,			__A ,			indent=2						)







def     A__	(       __A ,			__A						):







    '''simple docstring'''
    os.makedirs(__A ,			exist_ok=__A						)
    _lowerCamelCase  :		Dict			    =    {}
    for name, parameter in state_dict.items():
        _lowerCamelCase  :		Union[str, Any]			    =    offload_weight(__A ,			__A ,			__A ,			index=__A						)

    # Update index
    save_offload_index(__A ,			__A						)




class 						__snake_case						(			__lowerCAmelCase		):





      '''simple docstring'''






      def __init__(       self  :					str     ,	_UpperCamelCase  :					Mapping     ,	_UpperCamelCase  :					str)					->List[Any]:


          """simple docstring"""


          _lowerCamelCase  :		Union[str, Any]			    =    dataset
          _lowerCamelCase  :		Tuple			    =    prefix



      def __getitem__(       self  :					List[Any]     ,	_UpperCamelCase  :					Union[str, Any])					->Union[str, Any]:


          """simple docstring"""


          return self.dataset[F"""{self.prefix}{key}"""]



      def __iter__(       self  :					List[str])					->Any:


          """simple docstring"""


          return iter([key for key in self.dataset if key.startswith(self.prefix)])




      def __len__(       self  :					Any)					->int:


          """simple docstring"""


          return len(self.dataset)





class 						__snake_case						(			__lowerCAmelCase		):





      '''simple docstring'''






      def __init__(       self  :					Union[str, Any]     ,	_UpperCamelCase  :					Dict[str, torch.Tensor] = None     ,	_UpperCamelCase  :					Optional[Union[str, os.PathLike]] = None     ,	_UpperCamelCase  :					Mapping = None     ,	_UpperCamelCase  :					Union[str, Any]=None     ,	)					->Any:


          """simple docstring"""


          if state_dict is None and save_folder is None:
              raise ValueError("""Need either a `state_dict` or a `save_folder` containing offloaded weights.""")

          _lowerCamelCase  :		int			    =    {} if state_dict is None else state_dict
          _lowerCamelCase  :		Optional[int]			    =    save_folder
          if index is None and save_folder is not None:
              with open(os.path.join(_UpperCamelCase     ,	"""index.json""")) as f:
                  _lowerCamelCase  :		Union[str, Any]			    =    json.load(_UpperCamelCase)
          _lowerCamelCase  :		str			    =    {} if index is None else index
          _lowerCamelCase  :		Optional[int]			    =    list(self.state_dict.keys())
          self.all_keys.extend([key for key in self.index if key not in self.all_keys])
          _lowerCamelCase  :		str			    =    device



      def __getitem__(       self  :					Union[str, Any]     ,	_UpperCamelCase  :					str)					->Tuple:


          """simple docstring"""


          if key in self.state_dict:
              return self.state_dict[key]
          _lowerCamelCase  :		int			    =    self.index[key]
          if weight_info.get("""safetensors_file""") is not None:
              if not is_safetensors_available():
                  raise ImportError("""These offloaded weights require the use of safetensors: `pip install safetensors`.""")

              from safetensors import safe_open

              _lowerCamelCase  :		Any			    =    """cpu""" if self.device is None else self.device
              with safe_open(weight_info["""safetensors_file"""]     ,	framework="""pt"""     ,	device=_UpperCamelCase) as f:
                  _lowerCamelCase  :		Tuple			    =    f.get_tensor(weight_info.get("""weight_name"""     ,	_UpperCamelCase))

              if "dtype" in weight_info:
                  return tensor.to(getattr(_UpperCamelCase     ,	weight_info["""dtype"""]))
              else:
                  return tensor

          _lowerCamelCase  :		Optional[Any]			    =    os.path.join(self.save_folder     ,	F"""{key}.dat""")
          return load_offloaded_weight(_UpperCamelCase     ,	_UpperCamelCase)



      def __iter__(       self  :					Dict)					->List[Any]:


          """simple docstring"""


          return iter(self.all_keys)




      def __len__(       self  :					int)					->Any:


          """simple docstring"""


          return len(self.all_keys)







def     A__	(       __A ,			__A						):







    '''simple docstring'''
    _lowerCamelCase  :		str			    =    {}
    for module_name in submodule_names:
        # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
        # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
        result.update(
            {
                key: param
                for key, param in state_dict.items()
                if key == module_name or key.startswith(module_name + """."""						)
            }						)
    return result




