





from __future__ import annotations

import inspect
import unittest

from transformers import MobileViTConfig
from transformers.file_utils import is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import numpy as np
						import tensorflow as tf

						from transformers import TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation, TFMobileViTModel
						from transformers.models.mobilevit.modeling_tf_mobilevit import TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import MobileViTImageProcessor
class SCREAMING_SNAKE_CASE_    (       snake_case_       ):






						def UpperCAmelCase_					(   self   :   Union[str, Any]						)				->     int:







							"""simple docstring"""


							snake_case_		:      Optional[Any]           =						self.config_class(**self.inputs_dict						)
							self.parent.assertTrue(hasattr(_A		, 'hidden_sizes'						)						)
							self.parent.assertTrue(hasattr(_A		, 'neck_hidden_sizes'						)						)
							self.parent.assertTrue(hasattr(_A		, 'num_attention_heads'						)						)
class SCREAMING_SNAKE_CASE_    :






						def __init__(   self   :   Any		, _A   :   str		, _A   :   Optional[int]=13		, _A   :   Optional[int]=32		, _A   :   Optional[Any]=2		, _A   :   Union[str, Any]=3		, _A   :   Dict=640		, _A   :   int=4		, _A   :   Dict="silu"		, _A   :   str=3		, _A   :   Tuple=32		, _A   :   List[Any]=0.1		, _A   :   Optional[Any]=0.1		, _A   :   Tuple=0.1		, _A   :   Tuple=0.0_2		, _A   :   List[Any]=True		, _A   :   Union[str, Any]=True		, _A   :   int=10		, _A   :   List[Any]=None		, )				->     Optional[int]:







							"""simple docstring"""


							snake_case_		:      Dict           =						parent
							snake_case_		:      int           =						batch_size
							snake_case_		:      Tuple           =						image_size
							snake_case_		:      int           =						patch_size
							snake_case_		:      str           =						num_channels
							snake_case_		:      Any           =						last_hidden_size
							snake_case_		:      List[str]           =						num_attention_heads
							snake_case_		:      List[str]           =						hidden_act
							snake_case_		:      List[Any]           =						conv_kernel_size
							snake_case_		:      Any           =						output_stride
							snake_case_		:      Tuple           =						hidden_dropout_prob
							snake_case_		:      Union[str, Any]           =						attention_probs_dropout_prob
							snake_case_		:      Dict           =						classifier_dropout_prob
							snake_case_		:      Any           =						use_labels
							snake_case_		:      Optional[Any]           =						is_training
							snake_case_		:      List[Any]           =						num_labels
							snake_case_		:      Optional[Any]           =						initializer_range
							snake_case_		:      Union[str, Any]           =						scope






						def UpperCAmelCase_					(   self   :   List[Any]						)				->     Optional[Any]:







							"""simple docstring"""


							snake_case_		:      Union[str, Any]           =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

							snake_case_		:      Tuple           =						None
							snake_case_		:      Tuple           =						None
							if self.use_labels:
								snake_case_		:      Dict           =						ids_tensor([self.batch_size]		, self.num_labels						)
								snake_case_		:      int           =						ids_tensor([self.batch_size, self.image_size, self.image_size]		, self.num_labels						)

							snake_case_		:      List[str]           =						self.get_config()

							return config, pixel_values, labels, pixel_labels






						def UpperCAmelCase_					(   self   :   Optional[Any]						)				->     Tuple:







							"""simple docstring"""


							return MobileViTConfig(
							    image_size=self.image_size		, patch_size=self.patch_size		, num_channels=self.num_channels		, num_attention_heads=self.num_attention_heads		, hidden_act=self.hidden_act		, conv_kernel_size=self.conv_kernel_size		, output_stride=self.output_stride		, hidden_dropout_prob=self.hidden_dropout_prob		, attention_probs_dropout_prob=self.attention_probs_dropout_prob		, classifier_dropout_prob=self.classifier_dropout_prob		, initializer_range=self.initializer_range		, )






						def UpperCAmelCase_					(   self   :   int		, _A   :   List[Any]		, _A   :   Optional[int]		, _A   :   Any		, _A   :   Optional[int]						)				->     Any:







							"""simple docstring"""


							snake_case_		:      Any           =						TFMobileViTModel(config=_A						)
							snake_case_		:      int           =						model(_A		, training=_A						)
							snake_case_		:      int           =						self.image_size // self.output_stride
							self.parent.assertEqual(
							    result.last_hidden_state.shape		, (self.batch_size, self.last_hidden_size, expected_height, expected_width)						)






						def UpperCAmelCase_					(   self   :   List[str]		, _A   :   int		, _A   :   int		, _A   :   int		, _A   :   int						)				->     Dict:







							"""simple docstring"""


							snake_case_		:      Union[str, Any]           =						self.num_labels
							snake_case_		:      int           =						TFMobileViTForImageClassification(_A						)
							snake_case_		:      List[Any]           =						model(_A		, labels=_A		, training=_A						)
							self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.num_labels)						)






						def UpperCAmelCase_					(   self   :   Tuple		, _A   :   int		, _A   :   Dict		, _A   :   Optional[int]		, _A   :   List[str]						)				->     List[str]:







							"""simple docstring"""


							snake_case_		:      List[Any]           =						self.num_labels
							snake_case_		:      List[str]           =						TFMobileViTForSemanticSegmentation(_A						)
							snake_case_		:      Union[str, Any]           =						self.image_size // self.output_stride

							snake_case_		:      List[Any]           =						model(_A		, training=_A						)
							self.parent.assertEqual(
							    result.logits.shape		, (self.batch_size, self.num_labels, expected_height, expected_width)						)

							snake_case_		:      List[Any]           =						model(_A		, labels=_A		, training=_A						)
							self.parent.assertEqual(
							    result.logits.shape		, (self.batch_size, self.num_labels, expected_height, expected_width)						)






						def UpperCAmelCase_					(   self   :   Optional[Any]						)				->     str:







							"""simple docstring"""


							snake_case_		:      Any           =						self.prepare_config_and_inputs()
							snake_case_    ,snake_case_    ,snake_case_    ,snake_case_		:      Optional[int]           =						config_and_inputs
							snake_case_		:      Optional[Any]           =						{'pixel_values': pixel_values}
							return config, inputs_dict
@require_tf
class SCREAMING_SNAKE_CASE_    (       snake_case_							,					snake_case_							,					unittest.TestCase       ):
						__magic_name__:      str			   =							(
						    (TFMobileViTModel, TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation)
						    if is_tf_available()
						    else ()
						)
						__magic_name__:      Any			   =							(
						    {"feature-extraction": TFMobileViTModel, "image-classification": TFMobileViTForImageClassification}
						    if is_tf_available()
						    else {}
						)

						__magic_name__:      str			   =							False
						__magic_name__:      Optional[Any]			   =							False
						__magic_name__:      Any			   =							False
						__magic_name__:      Optional[Any]			   =							False
						__magic_name__:      int			   =							False






						def UpperCAmelCase_					(   self   :   int						)				->     List[str]:







							"""simple docstring"""


							snake_case_		:      Any           =						TFMobileViTModelTester(self						)
							snake_case_		:      List[Any]           =						TFMobileViTConfigTester(self		, config_class=_A		, has_text_modality=_A						)






						def UpperCAmelCase_					(   self   :   Union[str, Any]						)				->     Optional[Any]:







							"""simple docstring"""


							self.config_tester.run_common_tests()






						@unittest.skip(reason='MobileViT does not use inputs_embeds'						)
						def UpperCAmelCase_					(   self   :   Tuple						)				->     List[str]:







							"""simple docstring"""


							pass






						@unittest.skip(reason='MobileViT does not support input and output embeddings'						)
						def UpperCAmelCase_					(   self   :   Union[str, Any]						)				->     Optional[int]:







							"""simple docstring"""


							pass






						@unittest.skip(reason='MobileViT does not output attentions'						)
						def UpperCAmelCase_					(   self   :   Tuple						)				->     Optional[int]:







							"""simple docstring"""


							pass






						def UpperCAmelCase_					(   self   :   List[Any]						)				->     Optional[int]:







							"""simple docstring"""


							snake_case_    ,snake_case_		:      Optional[int]           =						self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
								snake_case_		:      str           =						model_class(_A						)
								snake_case_		:      Union[str, Any]           =						inspect.signature(model.call						)
								# signature.parameters is an OrderedDict => so arg_names order is deterministic
								snake_case_		:      List[str]           =						[*signature.parameters.keys()]

								snake_case_		:      Optional[int]           =						['pixel_values']
								self.assertListEqual(arg_names[:1]		, _A						)






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







							"""simple docstring"""


							snake_case_		:      List[Any]           =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_A						)






						def UpperCAmelCase_					(   self   :   Optional[int]						)				->     Optional[Any]:







							"""simple docstring"""


							def check_hidden_states_output(_A   :   Optional[int]		, _A   :   int		, _A   :   Dict						):
								snake_case_		:      List[str]           =						model_class(_A						)

								snake_case_		:      List[str]           =						model(**self._prepare_for_class(_A		, _A						)						)

								snake_case_		:      Tuple           =						outputs.hidden_states

								snake_case_		:      Optional[Any]           =						5
								self.assertEqual(len(_A						)		, _A						)

								# MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
								# with the width and height being successively divided by 2.
								snake_case_		:      Union[str, Any]           =						2
								for i in range(len(_A						)						):
									self.assertListEqual(
									    list(hidden_states[i].shape[-2:]						)		, [self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]		, )
									divisor *= 2

								self.assertEqual(self.model_tester.output_stride		, divisor // 2						)

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

							for model_class in self.all_model_classes:
								snake_case_		:      Dict           =						True
								check_hidden_states_output(_A		, _A		, _A						)

								# check that output_hidden_states also work using config
								del inputs_dict["output_hidden_states"]
								snake_case_		:      Union[str, Any]           =						True

								check_hidden_states_output(_A		, _A		, _A						)






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







							"""simple docstring"""


							snake_case_		:      Tuple           =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*_A						)






						def UpperCAmelCase_					(   self   :   str						)				->     Optional[int]:







							"""simple docstring"""


							snake_case_		:      Any           =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_semantic_segmentation(*_A						)






						@unittest.skipIf(
						    not is_tf_available() or len(tf.config.list_physical_devices('GPU'						)						) == 0		, reason='TF does not support backprop for grouped convolutions on CPU.'		, )
						def UpperCAmelCase_					(   self   :   str						)				->     str:







							"""simple docstring"""


							super().test_dataset_conversion()






						def UpperCAmelCase_					(   self   :   str		, _A   :   List[str]		, _A   :   Any		, _A   :   Tuple=2E-1		, _A   :   Union[str, Any]=2E-1						)				->     int:







							"""simple docstring"""


							self.assertTrue(np.allclose(_A		, _A		, atol=_A		, rtol=_A						)						)






						@unittest.skipIf(
						    not is_tf_available() or len(tf.config.list_physical_devices('GPU'						)						) == 0		, reason='TF does not support backprop for grouped convolutions on CPU.'		, )
						@slow
						def UpperCAmelCase_					(   self   :   List[str]						)				->     List[str]:







							"""simple docstring"""


							snake_case_    ,snake_case_		:      int           =						self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
								# Since `TFMobileViTModel` cannot operate with the default `fit()` method.
								if model_class.__name__ != "TFMobileViTModel":
									snake_case_		:      Optional[int]           =						model_class(_A						)
									if getattr(_A		, 'hf_compute_loss'		, _A						):
										super().test_keras_fit()






						def UpperCAmelCase_					(   self   :   int						)				->     List[str]:







							"""simple docstring"""


							snake_case_    ,snake_case_		:      Any           =						self.model_tester.prepare_config_and_inputs_for_common()
							for model_class in self.all_model_classes:
								# set an ignore index to correctly test the masked loss used in
								# `TFMobileViTForSemanticSegmentation`.
								if model_class.__name__ != "TFMobileViTForSemanticSegmentation":
									snake_case_		:      Optional[Any]           =						5

								snake_case_		:      Tuple           =						model_class(_A						)
								if getattr(_A		, 'hf_compute_loss'		, _A						):
									# The number of elements in the loss should be the same as the number of elements in the label
									snake_case_		:      Optional[int]           =						self._prepare_for_class(inputs_dict.copy()		, _A		, return_labels=_A						)
									snake_case_		:      int           =						prepared_for_class[
									    sorted(prepared_for_class.keys() - inputs_dict.keys()		, reverse=_A						)[0]
									]
									snake_case_		:      List[Any]           =						added_label.shape.as_list()[:1]

									# Test that model correctly compute the loss with kwargs
									snake_case_		:      int           =						self._prepare_for_class(inputs_dict.copy()		, _A		, return_labels=_A						)
									snake_case_		:      Dict           =						{'input_ids', 'pixel_values', 'input_features'}
									snake_case_		:      Optional[Any]           =						possible_input_names.intersection(set(_A						)						).pop()
									snake_case_		:      Dict           =						prepared_for_class.pop(_A						)

									snake_case_		:      Optional[int]           =						model(_A		, **_A						)[0]
									self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]						)

									# Test that model correctly compute the loss when we mask some positions
									snake_case_		:      Optional[int]           =						self._prepare_for_class(inputs_dict.copy()		, _A		, return_labels=_A						)
									snake_case_		:      Union[str, Any]           =						{'input_ids', 'pixel_values', 'input_features'}
									snake_case_		:      Dict           =						possible_input_names.intersection(set(_A						)						).pop()
									snake_case_		:      List[str]           =						prepared_for_class.pop(_A						)
									if "labels" in prepared_for_class:
										snake_case_		:      List[str]           =						prepared_for_class['labels'].numpy()
										if len(labels.shape						) > 1 and labels.shape[1] != 1:
											# labels[0] = -100
											snake_case_		:      Tuple           =						tf.convert_to_tensor(_A						)
											snake_case_		:      List[str]           =						model(_A		, **_A						)[0]
											self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]						)
											self.assertTrue(not np.any(np.isnan(loss.numpy()						)						)						)

                # Test that model correctly compute the loss with a dict
									snake_case_		:      Dict           =						self._prepare_for_class(inputs_dict.copy()		, _A		, return_labels=_A						)
									snake_case_		:      int           =						model(_A						)[0]
									self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]						)

									# Test that model correctly compute the loss with a tuple
									snake_case_		:      List[str]           =						self._prepare_for_class(inputs_dict.copy()		, _A		, return_labels=_A						)

									# Get keys that were added with the _prepare_for_class function
									snake_case_		:      Union[str, Any]           =						prepared_for_class.keys() - inputs_dict.keys()
									snake_case_		:      int           =						inspect.signature(model.call						).parameters
									snake_case_		:      str           =						list(signature.keys()						)

									# Create a dictionary holding the location of the tensors in the tuple
									snake_case_		:      List[Any]           =						{0: input_name}
									for label_key in label_keys:
										snake_case_		:      Dict           =						signature_names.index(_A						)
										snake_case_		:      Tuple           =						label_key
									snake_case_		:      Union[str, Any]           =						sorted(tuple_index_mapping.items()						)
									# Initialize a list with their default values, update the values and convert to a tuple
									snake_case_		:      Tuple           =						[]

									for name in signature_names:
										if name != "kwargs":
											list_input.append(signature[name].default						)

									for index, value in sorted_tuple_index_mapping:
										snake_case_		:      List[Any]           =						prepared_for_class[value]

									snake_case_		:      List[Any]           =						tuple(_A						)

									# Send to model
									snake_case_		:      str           =						model(tuple_input[:-1]						)[0]

									self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]						)






						@slow
						def UpperCAmelCase_					(   self   :   Tuple						)				->     Optional[Any]:







							"""simple docstring"""


							for model_name in TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								snake_case_		:      List[str]           =						TFMobileViTModel.from_pretrained(_A						)
								self.assertIsNotNone(_A						)


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


@require_tf
class SCREAMING_SNAKE_CASE_    (       unittest.TestCase       ):






						@slow
						def UpperCAmelCase_					(   self   :   Optional[Any]						)				->     Tuple:







							"""simple docstring"""


							snake_case_		:      int           =						TFMobileViTForImageClassification.from_pretrained('apple/mobilevit-xx-small'						)

							snake_case_		:      int           =						MobileViTImageProcessor.from_pretrained('apple/mobilevit-xx-small'						)
							snake_case_		:      str           =						prepare_img()
							snake_case_		:      List[str]           =						image_processor(images=_A		, return_tensors='tf'						)

							# forward pass
							snake_case_		:      Tuple           =						model(**_A		, training=_A						)

							# verify the logits
							snake_case_		:      Tuple           =						tf.TensorShape((1, 1000)						)
							self.assertEqual(outputs.logits.shape		, _A						)

							snake_case_		:      Optional[int]           =						tf.constant([-1.9_3_6_4, -1.2_3_2_7, -0.4_6_5_3]						)

							tf.debugging.assert_near(outputs.logits[0, :3]		, _A		, atol=1E-4		, rtol=1E-04						)






						@slow
						def UpperCAmelCase_					(   self   :   List[Any]						)				->     Optional[int]:







							"""simple docstring"""


							snake_case_		:      str           =						TFMobileViTForSemanticSegmentation.from_pretrained('apple/deeplabv3-mobilevit-xx-small'						)

							snake_case_		:      Any           =						MobileViTImageProcessor.from_pretrained('apple/deeplabv3-mobilevit-xx-small'						)

							snake_case_		:      List[Any]           =						prepare_img()
							snake_case_		:      Optional[Any]           =						image_processor(images=_A		, return_tensors='tf'						)

							# forward pass
							snake_case_		:      Optional[int]           =						model(inputs.pixel_values		, training=_A						)
							snake_case_		:      Dict           =						outputs.logits

							# verify the logits
							snake_case_		:      Tuple           =						tf.TensorShape((1, 21, 32, 32)						)
							self.assertEqual(logits.shape		, _A						)

							snake_case_		:      List[str]           =						tf.constant(
							    [
							        [[6.9_7_1_3, 6.9_7_8_6, 7.2_4_2_2], [7.2_8_9_3, 7.2_8_2_5, 7.4_4_4_6], [7.6_5_8_0, 7.8_7_9_7, 7.9_4_2_0]],
							        [[-1_0.6_8_6_9, -1_0.3_2_5_0, -1_0.3_4_7_1], [-1_0.4_2_2_8, -9.9_8_6_8, -9.7_1_3_2], [-1_1.0_4_0_5, -1_1.0_2_2_1, -1_0.7_3_1_8]],
							        [[-3.3_0_8_9, -2.8_5_3_9, -2.6_7_4_0], [-3.2_7_0_6, -2.5_6_2_1, -2.5_1_0_8], [-3.2_5_3_4, -2.6_6_1_5, -2.6_6_5_1]],
							    ]						)

							tf.debugging.assert_near(logits[0, :3, :3, :3]		, _A		, rtol=1E-4		, atol=1E-4						)






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 						_UpperCAmelCase							( snake_case      ,		snake_case      ,		snake_case      ,		snake_case=None     ):





					"""simple docstring"""

					_lowerCAmelCase         =       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         =       weight.view(torch.intaa     )
										_lowerCAmelCase         =       """bfloat16"""
					_lowerCAmelCase         =       weight.cpu().numpy()
					_lowerCAmelCase         =       os.path.join(snake_case      ,		F'{weight_name}.dat'     )
					if index is not None:
										if dtype is None:
															_lowerCAmelCase         =       str(array.dtype     )
										_lowerCAmelCase         =       {"""dtype""": dtype, """shape""": list(array.shape     )}
					if array.ndim == 0:
										_lowerCAmelCase         =       array[None]
					_lowerCAmelCase         =       np.memmap(snake_case      ,		dtype=array.dtype      ,		mode="""w+"""      ,		shape=array.shape     )
					_lowerCAmelCase         =       array[:]
					file_array.flush()
					return index


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

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

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

					_lowerCAmelCase         =       np.memmap(snake_case      ,		dtype=snake_case      ,		shape=snake_case      ,		mode="""r"""     )

					if len(weight_info["""shape"""]     ) == 0:
										_lowerCAmelCase         =       weight[0]
					_lowerCAmelCase         =       torch.tensor(snake_case     )
					if weight_info["dtype"] == "bfloat16":
										_lowerCAmelCase         =       weight.view(torch.bfloataa     )

					return weight


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					if index is None or len(snake_case     ) == 0:
										# Nothing to save
										return

					_lowerCAmelCase         =       os.path.join(snake_case      ,		"""index.json"""     )
					if os.path.isfile(snake_case     ):
										with open(snake_case      ,		"""r"""      ,		encoding="""utf-8"""     ) as f:
															_lowerCAmelCase         =       json.load(snake_case     )
					else:
										_lowerCAmelCase         =       {}
					current_index.update(snake_case     )

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


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					os.makedirs(snake_case      ,		exist_ok=snake_case     )
					_lowerCAmelCase         =       {}
					for name, parameter in state_dict.items():
										_lowerCAmelCase         =       offload_weight(snake_case      ,		snake_case      ,		snake_case      ,		index=snake_case     )

					# Update index
					save_offload_index(snake_case      ,		snake_case     )





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def __init__(				self		,							_snake_case		,							_snake_case						):



									"""simple docstring"""





									_lowerCAmelCase         =       dataset
									_lowerCAmelCase         =       prefix




				def __getitem__(				self		,							_snake_case						):



									"""simple docstring"""





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




				def __iter__(				self						):



									"""simple docstring"""





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




				def __len__(				self						):



									"""simple docstring"""





									return len(self.dataset						)





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def __init__(				self		,							_snake_case = None		,							_snake_case = None		,							_snake_case = None		,							_snake_case=None		,							):



									"""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         =       {} if state_dict is None else state_dict
									_lowerCAmelCase         =       save_folder
									if index is None and save_folder is not None:
														with open(os.path.join(_snake_case		,							"""index.json"""						)						) as f:
																			_lowerCAmelCase         =       json.load(_snake_case						)
									_lowerCAmelCase         =       {} if index is None else index
									_lowerCAmelCase         =       list(self.state_dict.keys()						)
									self.all_keys.extend([key for key in self.index if key not in self.all_keys]						)
									_lowerCAmelCase         =       device




				def __getitem__(				self		,							_snake_case						):



									"""simple docstring"""





									if key in self.state_dict:
														return self.state_dict[key]
									_lowerCAmelCase         =       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         =       """cpu""" if self.device is None else self.device
														with safe_open(weight_info["""safetensors_file"""]		,							framework="""pt"""		,							device=_snake_case						) as f:
																			_lowerCAmelCase         =       f.get_tensor(weight_info.get("""weight_name"""		,							_snake_case						)						)

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

									_lowerCAmelCase         =       os.path.join(self.save_folder		,							F'{key}.dat'						)
									return load_offloaded_weight(_snake_case		,							_snake_case						)




				def __iter__(				self						):



									"""simple docstring"""





									return iter(self.all_keys						)




				def __len__(				self						):



									"""simple docstring"""





									return len(self.all_keys						)


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					_lowerCAmelCase         =       {}
					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



