import importlib.metadata
import warnings
from copy import deepcopy

from packaging import version

from ..utils import logging
from .import_utils import is_accelerate_available, is_bitsandbytes_available


if is_bitsandbytes_available():
		import bitsandbytes as bnb
		import torch
		import torch.nn as nn

		from ..pytorch_utils import ConvaD

if is_accelerate_available():
		from accelerate import init_empty_weights
		from accelerate.utils import find_tied_parameters

lowerCamelCase__							     =					logging.get_logger(__name__)



def 				_lowerCamelCase(					__snake_case    ,    __snake_case    ,    __snake_case    ,    __snake_case=None    ,    __snake_case=None						)	-> Optional[Any]:



		if "." in tensor_name:
				__snake_case									=  tensor_name.split("."						)
				for split in splits[:-1]:
						__snake_case									=  getattr(__a    ,    __a						)
						if new_module is None:
								raise ValueError(f"""{module} has no attribute {split}."""						)
						__snake_case									=  new_module
				__snake_case									=  splits[-1]

		if tensor_name not in module._parameters and tensor_name not in module._buffers:
				raise ValueError(f"""{module} does not have a parameter or a buffer named {tensor_name}."""						)
		__snake_case									=  tensor_name in module._buffers
		__snake_case									=  getattr(__a    ,    __a						)

		if old_value.device == torch.device("meta"						) and device not in ["meta", torch.device("meta"						)] and value is None:
				raise ValueError(f"""{tensor_name} is on the meta device, we need a `value` to put in on {device}."""						)

		__snake_case									=  False
		__snake_case									=  False
		if is_buffer or not is_bitsandbytes_available():
				__snake_case									=  False
				__snake_case									=  False
		else:
				__snake_case									=  hasattr(bnb.nn    ,    "Params4bit"						) and isinstance(module._parameters[tensor_name]    ,    bnb.nn.Paramsabit						)
				__snake_case									=  isinstance(module._parameters[tensor_name]    ,    bnb.nn.IntaParams						)

		if is_abit or is_abit:
				__snake_case									=  module._parameters[tensor_name]
				if param.device.type != "cuda":
						if value is None:
								__snake_case									=  old_value.to(__a						)
						elif isinstance(__a    ,    torch.Tensor						):
								__snake_case									=  value.to("cpu"						)
								if value.dtype == torch.inta:
										__snake_case									=  version.parse(importlib.metadata.version("bitsandbytes"						)						) > version.parse(
										    "0.37.2"						)
										if not is_abit_serializable:
												raise ValueError(
												    "Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. "
												    "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."						)
						else:
								__snake_case									=  torch.tensor(__a    ,    device="cpu"						)

						# Support models using `Conv1D` in place of `nn.Linear` (e.g. gpt2) by transposing the weight matrix prior to quantization.
						# Since weights are saved in the correct "orientation", we skip transposing when loading.
						if issubclass(module.source_cls    ,    __a						) and fpaa_statistics is None:
								__snake_case									=  new_value.T

						__snake_case									=  old_value.__dict__
						if is_abit:
								__snake_case									=  bnb.nn.IntaParams(__a    ,    requires_grad=__a    ,    **__a						).to(__a						)
						elif is_abit:
								__snake_case									=  bnb.nn.Paramsabit(__a    ,    requires_grad=__a    ,    **__a						).to(__a						)

						__snake_case									=  new_value
						if fpaa_statistics is not None:
								setattr(module.weight    ,    "SCB"    ,    fpaa_statistics.to(__a						)						)

		else:
				if value is None:
						__snake_case									=  old_value.to(__a						)
				elif isinstance(__a    ,    torch.Tensor						):
						__snake_case									=  value.to(__a						)
				else:
						__snake_case									=  torch.tensor(__a    ,    device=__a						)

				if is_buffer:
						__snake_case									=  new_value
				else:
						__snake_case									=  nn.Parameter(__a    ,    requires_grad=old_value.requires_grad						)
						__snake_case									=  new_value



def 				_lowerCamelCase(					__snake_case    ,    __snake_case=None    ,    __snake_case=None    ,    __snake_case=None    ,    __snake_case=False						)	-> Optional[int]:



		for name, module in model.named_children():
				if current_key_name is None:
						__snake_case									=  []
				current_key_name.append(__a						)

				if (isinstance(__a    ,    nn.Linear						) or isinstance(__a    ,    __a						)) and name not in modules_to_not_convert:
						# Check if the current key is not in the `modules_to_not_convert`
						if not any(key in ".".join(__a						) for key in modules_to_not_convert						):
								with init_empty_weights():
										if isinstance(__a    ,    __a						):
												__snake_case									=  module.weight.shape
										else:
												__snake_case									=  module.in_features
												__snake_case									=  module.out_features

										if quantization_config.quantization_method() == "llm_int8":
												__snake_case									=  bnb.nn.LinearabitLt(
												    __a    ,    __a    ,    module.bias is not None    ,    has_fpaa_weights=quantization_config.llm_inta_has_fpaa_weight    ,    threshold=quantization_config.llm_inta_threshold    ,    )
												__snake_case									=  True
										else:
												if (
												    quantization_config.llm_inta_skip_modules is not None
												    and name in quantization_config.llm_inta_skip_modules
												):
														pass
												else:
														__snake_case									=  bnb.nn.Linearabit(
														    __a    ,    __a    ,    module.bias is not None    ,    quantization_config.bnb_abit_compute_dtype    ,    compress_statistics=quantization_config.bnb_abit_use_double_quant    ,    quant_type=quantization_config.bnb_abit_quant_type    ,    )
														__snake_case									=  True
                    # Store the module class in case we need to transpose the weight later
										__snake_case									=  type(__a						)
										# Force requires grad to False to avoid unexpected errors
										model._modules[name].requires_grad_(__a						)
				if len(list(module.children()						)						) > 0:
						__snake_case									=  _replace_with_bnb_linear(
						    __a    ,    __a    ,    __a    ,    __a    ,    has_been_replaced=__a    ,    )
				# Remove the last key for recursion
				current_key_name.pop(-1						)
		return model, has_been_replaced



def 				_lowerCamelCase(					__snake_case    ,    __snake_case=None    ,    __snake_case=None    ,    __snake_case=None						)	-> Optional[Any]:
		__snake_case									=  ['lm_head'] if modules_to_not_convert is None else modules_to_not_convert
		__snake_case									=  _replace_with_bnb_linear(
		    __a    ,    __a    ,    __a    ,    __a						)

		if not has_been_replaced:
				logger.warning(
				    "You are loading your model in 8bit or 4bit but no linear modules were found in your model."
				    " Please double check your model architecture, or submit an issue on github if you think this is"
				    " a bug."						)

		return model



def 				_lowerCamelCase(					*__snake_case    ,    **__snake_case						)	-> int:
		warnings.warn(
		    "`replace_8bit_linear` will be deprecated in a future version, please use `replace_with_bnb_linear` instead"    ,    __a    ,    )
		return replace_with_bnb_linear(*__a    ,    **__a						)



def 				_lowerCamelCase(					*__snake_case    ,    **__snake_case						)	-> Any:
		warnings.warn(
		    "`set_module_8bit_tensor_to_device` will be deprecated in a future version, please use `set_module_quantized_tensor_to_device` instead"    ,    __a    ,    )
		return set_module_quantized_tensor_to_device(*__a    ,    **__a						)





def 				_lowerCamelCase(					__snake_case						)	-> Union[str, Any]:
		__snake_case									=  deepcopy(__a						)  # this has 0 cost since it is done inside `init_empty_weights` context manager`
		tied_model.tie_weights()

		__snake_case									=  find_tied_parameters(__a						)
		# For compatibility with Accelerate < 0.18
		if isinstance(__a    ,    __a						):
				__snake_case									=  sum(list(tied_params.values()						)    ,    []						) + list(tied_params.keys()						)
		else:
				__snake_case									=  sum(__a    ,    []						)
		__snake_case									=  len(__a						) > 0

		# Check if it is a base model
		__snake_case									=  not hasattr(__a    ,    model.base_model_prefix						)

		# Ignore this for base models (BertModel, GPT2Model, etc.)
		if (not has_tied_params) and is_base_model:
				return []

		# otherwise they have an attached head
		__snake_case									=  list(model.named_children()						)
		__snake_case									=  [list_modules[-1][0]]

		# add last module together with tied weights
		__snake_case									=  set(__a						) - set(__a						)
		__snake_case									=  list(set(__a						)						) + list(__a						)

		# remove ".weight" from the keys
		__snake_case									=  ['.weight', '.bias']
		__snake_case									=  []
		for name in list_untouched:
				for name_to_remove in names_to_remove:
						if name_to_remove in name:
								__snake_case									=  name.replace(__a    ,    ""						)
				filtered_module_names.append(__a						)

		return filtered_module_names










"""simple docstring"""







import inspect
import unittest

from transformers import ViTMSNConfig
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 torch import nn

	from transformers import ViTMSNForImageClassification, ViTMSNModel
	from transformers.models.vit_msn.modeling_vit_msn import VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image

	from transformers import ViTImageProcessor







class 					A_:







			"""simple docstring"""



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

						# in ViT MSN, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
						_lowerCamelCase					:					Dict	       =    (image_size // patch_size) ** 2
						_lowerCamelCase					:					Any	       =    num_patches + 1
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Any	       =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

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

						_lowerCamelCase					:					Dict	       =    self.get_config()

						return config, pixel_values, labels
			def    _lowerCAmelCase		(	self							):
						return ViTMSNConfig(
						    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				,   initializer_range=self.initializer_range				,   )
			def    _lowerCAmelCase		(	self				,   A				,   A				,   A							):
						_lowerCamelCase					:					int	       =    ViTMSNModel(config=A							)
						model.to(A							)
						model.eval()
						_lowerCamelCase					:					List[Any]	       =    model(A							)
						self.parent.assertEqual(result.last_hidden_state.shape				,   (self.batch_size, self.seq_length, self.hidden_size)							)
			def    _lowerCAmelCase		(	self				,   A				,   A				,   A							):
						_lowerCamelCase					:					List[Any]	       =    self.type_sequence_label_size
						_lowerCamelCase					:					List[str]	       =    ViTMSNForImageClassification(A							)
						model.to(A							)
						model.eval()
						_lowerCamelCase					:					Optional[Any]	       =    model(A				,   labels=A							)
						print('Pixel and labels shape: {pixel_values.shape}, {labels.shape}'							)
						print('Labels: {labels}'							)
						self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.type_sequence_label_size)							)

						# test greyscale images
						_lowerCamelCase					:					Optional[Any]	       =    1
						_lowerCamelCase					:					List[Any]	       =    ViTMSNForImageClassification(A							)
						model.to(A							)
						model.eval()

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

			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					int	       =    self.prepare_config_and_inputs()
						_lowerCamelCase       ,							_lowerCamelCase       ,							_lowerCamelCase					:					Tuple	       =    config_and_inputs
						_lowerCamelCase					:					List[str]	       =    {'pixel_values': pixel_values}
						return config, inputs_dict







@require_torch
class 					A_(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_   ,   unittest.TestCase      ):







			"""simple docstring"""



			a_       :			Union[str, Any]												=			(ViTMSNModel, ViTMSNForImageClassification) if is_torch_available() else ()
			a_       :			Union[str, Any]												=			(
			    {"""feature-extraction""": ViTMSNModel, """image-classification""": ViTMSNForImageClassification}
			    if is_torch_available()
			    else {}
			)

			a_       :			Tuple												=			False
			a_       :			Union[str, Any]												=			False
			a_       :			Tuple												=			False
			a_       :			List[str]												=			False
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					List[str]	       =    ViTMSNModelTester(self							)
						_lowerCamelCase					:					List[str]	       =    ConfigTester(self				,   config_class=A				,   has_text_modality=A				,   hidden_size=37							)
			def    _lowerCAmelCase		(	self							):
						self.config_tester.run_common_tests()
			@unittest.skip(reason='ViTMSN does not use inputs_embeds'							)
			def    _lowerCAmelCase		(	self							):
						pass
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase       ,							_lowerCamelCase					:					Optional[Any]	       =    self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									_lowerCamelCase					:					List[str]	       =    model_class(A							)
									self.assertIsInstance(model.get_input_embeddings()				,   (nn.Module)							)
									_lowerCamelCase					:					Dict	       =    model.get_output_embeddings()
									self.assertTrue(x is None or isinstance(A				,   nn.Linear							)							)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase       ,							_lowerCamelCase					:					List[str]	       =    self.model_tester.prepare_config_and_inputs_for_common()

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

									_lowerCamelCase					:					Union[str, Any]	       =    ['pixel_values']
									self.assertListEqual(arg_names[:1]				,   A							)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					List[Any]	       =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*A							)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					List[Any]	       =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*A							)

			@slow
			def    _lowerCAmelCase		(	self							):
						for model_name in VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									_lowerCamelCase					:					Any	       =    ViTMSNModel.from_pretrained(A							)
									self.assertIsNotNone(A							)






def     UpperCAmelCase_					(				):




			'''simple docstring'''



			_lowerCamelCase					:					Optional[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    _lowerCAmelCase		(	self							):
						return ViTImageProcessor.from_pretrained('facebook/vit-msn-small'							) if is_vision_available() else None

			@slow
			def    _lowerCAmelCase		(	self							):
						torch.manual_seed(2							)
						_lowerCamelCase					:					Tuple	       =    ViTMSNForImageClassification.from_pretrained('facebook/vit-msn-small'							).to(A							)

						_lowerCamelCase					:					int	       =    self.default_image_processor
						_lowerCamelCase					:					Optional[int]	       =    prepare_img()
						_lowerCamelCase					:					int	       =    image_processor(images=A				,   return_tensors='pt'							).to(A							)

						# forward pass
						with torch.no_grad():
									_lowerCamelCase					:					str	       =    model(**A							)

						# verify the logits
						_lowerCamelCase					:					Optional[Any]	       =    torch.Size((1, 1000)							)
						self.assertEqual(outputs.logits.shape				,   A							)

						_lowerCamelCase					:					Any	       =    torch.tensor([-0.0_8_0_3, -0.4_4_5_4, -0.2_3_7_5]							).to(A							)

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





