





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 lowerCAmelCase__(__snake_case	,__snake_case	,__snake_case	,__snake_case=None							)    ->				Tuple:







	'''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 lowerCAmelCase__(__snake_case	,__snake_case							)    ->				List[Any]:







	'''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 lowerCAmelCase__(__snake_case	,__snake_case							)    ->				int:







	'''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 lowerCAmelCase__(__snake_case	,__snake_case							)    ->				str:







	'''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  __A   (   lowerCAmelCase			):

			'''simple docstring'''







			def __init__(   self	,  __lowerCAmelCase	,  __lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__        =		dataset
				lowerCamelCase__        =		prefix





			def __getitem__(   self	,  __lowerCAmelCase							):



				'''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  __A   (   lowerCAmelCase			):

			'''simple docstring'''







			def __init__(   self	,  __lowerCAmelCase = None	,  __lowerCAmelCase = None	,  __lowerCAmelCase = None	,  __lowerCAmelCase=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(__lowerCAmelCase	,  '''index.json'''							)							) as f:
						lowerCamelCase__        =		json.load(__lowerCAmelCase							)
				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	,  __lowerCAmelCase							):



				'''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=__lowerCAmelCase							) as f:
						lowerCamelCase__        =		f.get_tensor(weight_info.get('''weight_name'''	,  __lowerCAmelCase							)							)

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

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





			def __iter__(   self							):



				'''simple docstring'''







				return iter(self.all_keys							)





			def __len__(   self							):



				'''simple docstring'''







				return len(self.all_keys							)





def lowerCAmelCase__(__snake_case	,__snake_case							)    ->				str:







	'''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










import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
			import torch

if is_vision_available():
			from PIL import Image

			from transformers import ViltImageProcessor






class  __A   (   unittest.TestCase			):

			'''simple docstring'''







			def __init__(   self	,  __lowerCAmelCase	,  __lowerCAmelCase=7	,  __lowerCAmelCase=3	,  __lowerCAmelCase=1_8	,  __lowerCAmelCase=3_0	,  __lowerCAmelCase=4_0_0	,  __lowerCAmelCase=True	,  __lowerCAmelCase=None	,  __lowerCAmelCase=2	,  __lowerCAmelCase=True	,  __lowerCAmelCase=[0.5, 0.5, 0.5]	,  __lowerCAmelCase=[0.5, 0.5, 0.5]	,  ):



				'''simple docstring'''







				lowerCamelCase__        =		size if size is not None else {'''shortest_edge''': 3_0}
				lowerCamelCase__        =		parent
				lowerCamelCase__        =		batch_size
				lowerCamelCase__        =		num_channels
				lowerCamelCase__        =		image_size
				lowerCamelCase__        =		min_resolution
				lowerCamelCase__        =		max_resolution
				lowerCamelCase__        =		do_resize
				lowerCamelCase__        =		size
				lowerCamelCase__        =		size_divisor
				lowerCamelCase__        =		do_normalize
				lowerCamelCase__        =		image_mean
				lowerCamelCase__        =		image_std





			def 		__lowerCamelCase    (   self							):



				'''simple docstring'''







				return {
				    "image_mean": self.image_mean,
				    "image_std": self.image_std,
				    "do_normalize": self.do_normalize,
				    "do_resize": self.do_resize,
				    "size": self.size,
				    "size_divisor": self.size_divisor,
				}





			def 		__lowerCamelCase    (   self	,  __lowerCAmelCase	,  __lowerCAmelCase=False							):



				'''simple docstring'''







				if not batched:
					lowerCamelCase__        =		self.size['''shortest_edge''']
					lowerCamelCase__        =		image_inputs[0]
					if isinstance(__lowerCAmelCase	,  Image.Image							):
						lowerCamelCase__    ,     lowerCamelCase__        =		image.size
					else:
						lowerCamelCase__    ,     lowerCamelCase__        =		image.shape[1], image.shape[2]
					lowerCamelCase__        =		size / min(__lowerCAmelCase	,  __lowerCAmelCase							)
					if h < w:
						lowerCamelCase__    ,     lowerCamelCase__        =		size, scale * w
					else:
						lowerCamelCase__    ,     lowerCamelCase__        =		scale * h, size

					lowerCamelCase__        =		int((1_3_3_3 / 8_0_0) * size							)
					if max(__lowerCAmelCase	,  __lowerCAmelCase							) > max_size:
						lowerCamelCase__        =		max_size / max(__lowerCAmelCase	,  __lowerCAmelCase							)
						lowerCamelCase__        =		newh * scale
						lowerCamelCase__        =		neww * scale

					lowerCamelCase__    ,     lowerCamelCase__        =		int(newh + 0.5							), int(neww + 0.5							)
					lowerCamelCase__    ,     lowerCamelCase__        =		(
					    newh // self.size_divisor * self.size_divisor,
					    neww // self.size_divisor * self.size_divisor,
					)

				else:
					lowerCamelCase__        =		[]
					for image in image_inputs:
						lowerCamelCase__    ,     lowerCamelCase__        =		self.get_expected_values([image]							)
						expected_values.append((expected_height, expected_width)							)
					lowerCamelCase__        =		max(__lowerCAmelCase	,  key=lambda __lowerCAmelCase							: item[0]							)[0]
					lowerCamelCase__        =		max(__lowerCAmelCase	,  key=lambda __lowerCAmelCase							: item[1]							)[1]

				return expected_height, expected_width






@require_torch
@require_vision
class  __A   (   lowerCAmelCase      ,    unittest.TestCase			):

			'''simple docstring'''







			lowerCAmelCase_												=						ViltImageProcessor if is_vision_available() else None





			def 		__lowerCamelCase    (   self							):



				'''simple docstring'''







				lowerCamelCase__        =		ViltImageProcessingTester(self							)





			@property
			def 		__lowerCamelCase    (   self							):



				'''simple docstring'''







				return self.image_processor_tester.prepare_image_processor_dict()





			def 		__lowerCamelCase    (   self							):



				'''simple docstring'''







				lowerCamelCase__        =		self.image_processing_class(**self.image_processor_dict							)
				self.assertTrue(hasattr(__lowerCAmelCase	,  '''image_mean'''							)							)
				self.assertTrue(hasattr(__lowerCAmelCase	,  '''image_std'''							)							)
				self.assertTrue(hasattr(__lowerCAmelCase	,  '''do_normalize'''							)							)
				self.assertTrue(hasattr(__lowerCAmelCase	,  '''do_resize'''							)							)
				self.assertTrue(hasattr(__lowerCAmelCase	,  '''size'''							)							)
				self.assertTrue(hasattr(__lowerCAmelCase	,  '''size_divisor'''							)							)





			def 		__lowerCamelCase    (   self							):



				'''simple docstring'''







				lowerCamelCase__        =		self.image_processing_class.from_dict(self.image_processor_dict							)
				self.assertEqual(image_processor.size	,  {'''shortest_edge''': 3_0}							)

				lowerCamelCase__        =		self.image_processing_class.from_dict(self.image_processor_dict	,  size=4_2							)
				self.assertEqual(image_processor.size	,  {'''shortest_edge''': 4_2}							)





			def 		__lowerCamelCase    (   self							):



				'''simple docstring'''







				pass





			def 		__lowerCamelCase    (   self							):



				'''simple docstring'''







				lowerCamelCase__        =		self.image_processing_class(**self.image_processor_dict							)
				# create random PIL images
				lowerCamelCase__        =		prepare_image_inputs(self.image_processor_tester	,  equal_resolution=__lowerCAmelCase							)
				for image in image_inputs:
					self.assertIsInstance(__lowerCAmelCase	,  Image.Image							)

				# Test not batched input
				lowerCamelCase__        =		image_processing(image_inputs[0]	,  return_tensors='''pt'''							).pixel_values

				lowerCamelCase__    ,     lowerCamelCase__        =		self.image_processor_tester.get_expected_values(__lowerCAmelCase							)
				self.assertEqual(
				    encoded_images.shape	,  (1, self.image_processor_tester.num_channels, expected_height, expected_width)	,  )

				# Test batched
				lowerCamelCase__        =		image_processing(__lowerCAmelCase	,  return_tensors='''pt'''							).pixel_values

				lowerCamelCase__    ,     lowerCamelCase__        =		self.image_processor_tester.get_expected_values(__lowerCAmelCase	,  batched=__lowerCAmelCase							)
				self.assertEqual(
				    encoded_images.shape	,  (
				        self.image_processor_tester.batch_size,
				        self.image_processor_tester.num_channels,
				        expected_height,
				        expected_width,
				    )	,  )





			def 		__lowerCamelCase    (   self							):



				'''simple docstring'''







				lowerCamelCase__        =		self.image_processing_class(**self.image_processor_dict							)
				# create random numpy tensors
				lowerCamelCase__        =		prepare_image_inputs(self.image_processor_tester	,  equal_resolution=__lowerCAmelCase	,  numpify=__lowerCAmelCase							)
				for image in image_inputs:
					self.assertIsInstance(__lowerCAmelCase	,  np.ndarray							)

				# Test not batched input
				lowerCamelCase__        =		image_processing(image_inputs[0]	,  return_tensors='''pt'''							).pixel_values

				lowerCamelCase__    ,     lowerCamelCase__        =		self.image_processor_tester.get_expected_values(__lowerCAmelCase							)
				self.assertEqual(
				    encoded_images.shape	,  (1, self.image_processor_tester.num_channels, expected_height, expected_width)	,  )

				# Test batched
				lowerCamelCase__        =		image_processing(__lowerCAmelCase	,  return_tensors='''pt'''							).pixel_values

				lowerCamelCase__    ,     lowerCamelCase__        =		self.image_processor_tester.get_expected_values(__lowerCAmelCase	,  batched=__lowerCAmelCase							)
				self.assertEqual(
				    encoded_images.shape	,  (
				        self.image_processor_tester.batch_size,
				        self.image_processor_tester.num_channels,
				        expected_height,
				        expected_width,
				    )	,  )





			def 		__lowerCamelCase    (   self							):



				'''simple docstring'''







				lowerCamelCase__        =		self.image_processing_class(**self.image_processor_dict							)
				# create random PyTorch tensors
				lowerCamelCase__        =		prepare_image_inputs(self.image_processor_tester	,  equal_resolution=__lowerCAmelCase	,  torchify=__lowerCAmelCase							)
				for image in image_inputs:
					self.assertIsInstance(__lowerCAmelCase	,  torch.Tensor							)

				# Test not batched input
				lowerCamelCase__        =		image_processing(image_inputs[0]	,  return_tensors='''pt'''							).pixel_values

				lowerCamelCase__    ,     lowerCamelCase__        =		self.image_processor_tester.get_expected_values(__lowerCAmelCase							)
				self.assertEqual(
				    encoded_images.shape	,  (1, self.image_processor_tester.num_channels, expected_height, expected_width)	,  )

				# Test batched
				lowerCamelCase__        =		image_processing(__lowerCAmelCase	,  return_tensors='''pt'''							).pixel_values

				lowerCamelCase__    ,     lowerCamelCase__        =		self.image_processor_tester.get_expected_values(__lowerCAmelCase	,  batched=__lowerCAmelCase							)
				self.assertEqual(
				    encoded_images.shape	,  (
				        self.image_processor_tester.batch_size,
				        self.image_processor_tester.num_channels,
				        expected_height,
				        expected_width,
				    )	,  )



