


'''simple docstring'''






from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, logging


if is_torch_available():
		import torch


__UpperCAmelCase					       =							logging.get_logger(__name__)




class 	a__       (							a__						):





							'''simple docstring'''






							lowercase__    :     Union[str, Any]					=   ["pixel_values"]







							def __init__(	self      ,					lowerCamelCase_ = True      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = PILImageResampling.BILINEAR      ,					lowerCamelCase_ = True      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = True      ,					lowerCamelCase_ = 1 / 2_55      ,					lowerCamelCase_ = True      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = None      ,					**lowerCamelCase_      ,					) ->							None:
														super().__init__(**lowerCamelCase_     )
														lowerCAmelCase__      =			size if size is not None else {'''shortest_edge''': 2_56}
														lowerCAmelCase__      =			get_size_dict(lowerCamelCase_      ,					default_to_square=lowerCamelCase_     )
														lowerCAmelCase__      =			crop_size if crop_size is not None else {'''height''': 2_24, '''width''': 2_24}
														lowerCAmelCase__      =			get_size_dict(lowerCamelCase_      ,					param_name='''crop_size'''     )
														lowerCAmelCase__      =			do_resize
														lowerCAmelCase__      =			size
														lowerCAmelCase__      =			resample
														lowerCAmelCase__      =			do_center_crop
														lowerCAmelCase__      =			crop_size
														lowerCAmelCase__      =			do_rescale
														lowerCAmelCase__      =			rescale_factor
														lowerCAmelCase__      =			do_normalize
														lowerCAmelCase__      =			image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
														lowerCAmelCase__      =			image_std if image_std is not None else IMAGENET_STANDARD_STD







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_ = PILImageResampling.BICUBIC      ,					lowerCamelCase_ = None      ,					**lowerCamelCase_      ,					) ->							np.ndarray:
														lowerCAmelCase__      =			get_size_dict(lowerCamelCase_      ,					default_to_square=lowerCamelCase_     )
														if "shortest_edge" not in size:
																					raise ValueError(F"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""     )
														lowerCAmelCase__      =			get_resize_output_image_size(lowerCamelCase_      ,					size=size['''shortest_edge''']      ,					default_to_square=lowerCamelCase_     )
														return resize(lowerCamelCase_      ,					size=lowerCamelCase_      ,					resample=lowerCamelCase_      ,					data_format=lowerCamelCase_      ,					**lowerCamelCase_     )







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_ = None      ,					**lowerCamelCase_      ,					) ->							np.ndarray:
														lowerCAmelCase__      =			get_size_dict(lowerCamelCase_     )
														if "height" not in size or "width" not in size:
																					raise ValueError(F"""The `size` parameter must contain the keys `height` and `width`. Got {size.keys()}"""     )
														return center_crop(lowerCamelCase_      ,					size=(size['''height'''], size['''width'''])      ,					data_format=lowerCamelCase_      ,					**lowerCamelCase_     )







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_ = None      ,					**lowerCamelCase_     ) ->							np.ndarray:
														return rescale(lowerCamelCase_      ,					scale=lowerCamelCase_      ,					data_format=lowerCamelCase_      ,					**lowerCamelCase_     )







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_ = None      ,					**lowerCamelCase_      ,					) ->							np.ndarray:
														return normalize(lowerCamelCase_      ,					mean=lowerCamelCase_      ,					std=lowerCamelCase_      ,					data_format=lowerCamelCase_      ,					**lowerCamelCase_     )







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = None      ,					lowerCamelCase_ = ChannelDimension.FIRST      ,					**lowerCamelCase_      ,					) ->							int:
														lowerCAmelCase__      =			do_resize if do_resize is not None else self.do_resize
														lowerCAmelCase__      =			size if size is not None else self.size
														lowerCAmelCase__      =			get_size_dict(lowerCamelCase_      ,					default_to_square=lowerCamelCase_     )
														lowerCAmelCase__      =			resample if resample is not None else self.resample
														lowerCAmelCase__      =			do_center_crop if do_center_crop is not None else self.do_center_crop
														lowerCAmelCase__      =			crop_size if crop_size is not None else self.crop_size
														lowerCAmelCase__      =			get_size_dict(lowerCamelCase_      ,					param_name='''crop_size'''     )
														lowerCAmelCase__      =			do_rescale if do_rescale is not None else self.do_rescale
														lowerCAmelCase__      =			rescale_factor if rescale_factor is not None else self.rescale_factor
														lowerCAmelCase__      =			do_normalize if do_normalize is not None else self.do_normalize
														lowerCAmelCase__      =			image_mean if image_mean is not None else self.image_mean
														lowerCAmelCase__      =			image_std if image_std is not None else self.image_std

														lowerCAmelCase__      =			make_list_of_images(lowerCamelCase_     )

														if not valid_images(lowerCamelCase_     ):
																					raise ValueError(
																					    '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
																					    '''torch.Tensor, tf.Tensor or jax.ndarray.'''     )

														if do_resize and size is None:
																					raise ValueError('''Size must be specified if do_resize is True.'''     )

														if do_center_crop and crop_size is None:
																					raise ValueError('''Crop size must be specified if do_center_crop is True.'''     )

														if do_rescale and rescale_factor is None:
																					raise ValueError('''Rescale factor must be specified if do_rescale is True.'''     )

														if do_normalize and (image_mean is None or image_std is None):
																					raise ValueError('''Image mean and std must be specified if do_normalize is True.'''     )

														# All transformations expect numpy arrays.
														lowerCAmelCase__      =			[to_numpy_array(lowerCamelCase_     ) for image in images]

														if do_resize:
																					lowerCAmelCase__      =			[self.resize(image=lowerCamelCase_      ,					size=lowerCamelCase_      ,					resample=lowerCamelCase_     ) for image in images]

														if do_center_crop:
																					lowerCAmelCase__      =			[self.center_crop(image=lowerCamelCase_      ,					size=lowerCamelCase_     ) for image in images]

														if do_rescale:
																					lowerCAmelCase__      =			[self.rescale(image=lowerCamelCase_      ,					scale=lowerCamelCase_     ) for image in images]

														if do_normalize:
																					lowerCAmelCase__      =			[self.normalize(image=lowerCamelCase_      ,					mean=lowerCamelCase_      ,					std=lowerCamelCase_     ) for image in images]

														lowerCAmelCase__      =			[to_channel_dimension_format(lowerCamelCase_      ,					lowerCamelCase_     ) for image in images]

														lowerCAmelCase__      =			{'''pixel_values''': images}
														return BatchFeature(data=lowerCamelCase_      ,					tensor_type=lowerCamelCase_     )







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_ = None     ) ->							List[Any]:
														lowerCAmelCase__      =			outputs.logits

														# Resize logits and compute semantic segmentation maps
														if target_sizes is not None:
																					if len(lowerCamelCase_     ) != len(lowerCamelCase_     ):
																												raise ValueError(
																												    '''Make sure that you pass in as many target sizes as the batch dimension of the logits'''     )

																					if is_torch_tensor(lowerCamelCase_     ):
																												lowerCAmelCase__      =			target_sizes.numpy()

																					lowerCAmelCase__      =			[]

																					for idx in range(len(lowerCamelCase_     )     ):
																												lowerCAmelCase__      =			torch.nn.functional.interpolate(
																												    logits[idx].unsqueeze(dim=0     )      ,					size=target_sizes[idx]      ,					mode='''bilinear'''      ,					align_corners=lowerCamelCase_     )
																												lowerCAmelCase__      =			resized_logits[0].argmax(dim=0     )
																												semantic_segmentation.append(lowerCamelCase_     )
														else:
																					lowerCAmelCase__      =			logits.argmax(dim=1     )
																					lowerCAmelCase__      =			[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]     )]

														return semantic_segmentation



'''simple docstring'''






import tempfile
import unittest

import numpy as np
from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_flax


if is_flax_available():
		import os

		from flax.core.frozen_dict import unfreeze
		from flax.traverse_util import flatten_dict

		from transformers import FlaxBertModel

		__UpperCAmelCase					       =							'''0.12'''  # assumed parallelism: 8




@require_flax
@is_staging_test
class 	a__       (							unittest.TestCase						):





							'''simple docstring'''






							@classmethod
							def 						__SCREAMING_SNAKE_CASE							(	cls     ) ->							Any:
														lowerCAmelCase__      =			TOKEN
														HfFolder.save_token(lowerCamelCase_     )







							@classmethod
							def 						__SCREAMING_SNAKE_CASE							(	cls     ) ->							Tuple:
														try:
																					delete_repo(token=cls._token      ,					repo_id='''test-model-flax'''     )
														except HTTPError:
																					pass

														try:
																					delete_repo(token=cls._token      ,					repo_id='''valid_org/test-model-flax-org'''     )
														except HTTPError:
																					pass







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[Any]:
														lowerCAmelCase__      =			BertConfig(
														    vocab_size=99      ,					hidden_size=32      ,					num_hidden_layers=5      ,					num_attention_heads=4      ,					intermediate_size=37     )
														lowerCAmelCase__      =			FlaxBertModel(lowerCamelCase_     )
														model.push_to_hub('''test-model-flax'''      ,					use_auth_token=self._token     )

														lowerCAmelCase__      =			FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax"""     )

														lowerCAmelCase__      =			flatten_dict(unfreeze(model.params     )     )
														lowerCAmelCase__      =			flatten_dict(unfreeze(new_model.params     )     )

														for key in base_params.keys():
																					lowerCAmelCase__      =			(base_params[key] - new_params[key]).sum().item()
																					self.assertLessEqual(lowerCamelCase_      ,					1e-3      ,					msg=F"""{key} not identical"""     )

														# Reset repo
														delete_repo(token=self._token      ,					repo_id='''test-model-flax'''     )

														# Push to hub via save_pretrained
														with tempfile.TemporaryDirectory() as tmp_dir:
																					model.save_pretrained(lowerCamelCase_      ,					repo_id='''test-model-flax'''      ,					push_to_hub=lowerCamelCase_      ,					use_auth_token=self._token     )

														lowerCAmelCase__      =			FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax"""     )

														lowerCAmelCase__      =			flatten_dict(unfreeze(model.params     )     )
														lowerCAmelCase__      =			flatten_dict(unfreeze(new_model.params     )     )

														for key in base_params.keys():
																					lowerCAmelCase__      =			(base_params[key] - new_params[key]).sum().item()
																					self.assertLessEqual(lowerCamelCase_      ,					1e-3      ,					msg=F"""{key} not identical"""     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Union[str, Any]:
														lowerCAmelCase__      =			BertConfig(
														    vocab_size=99      ,					hidden_size=32      ,					num_hidden_layers=5      ,					num_attention_heads=4      ,					intermediate_size=37     )
														lowerCAmelCase__      =			FlaxBertModel(lowerCamelCase_     )
														model.push_to_hub('''valid_org/test-model-flax-org'''      ,					use_auth_token=self._token     )

														lowerCAmelCase__      =			FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''     )

														lowerCAmelCase__      =			flatten_dict(unfreeze(model.params     )     )
														lowerCAmelCase__      =			flatten_dict(unfreeze(new_model.params     )     )

														for key in base_params.keys():
																					lowerCAmelCase__      =			(base_params[key] - new_params[key]).sum().item()
																					self.assertLessEqual(lowerCamelCase_      ,					1e-3      ,					msg=F"""{key} not identical"""     )

														# Reset repo
														delete_repo(token=self._token      ,					repo_id='''valid_org/test-model-flax-org'''     )

														# Push to hub via save_pretrained
														with tempfile.TemporaryDirectory() as tmp_dir:
																					model.save_pretrained(
																					    lowerCamelCase_      ,					repo_id='''valid_org/test-model-flax-org'''      ,					push_to_hub=lowerCamelCase_      ,					use_auth_token=self._token     )

														lowerCAmelCase__      =			FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''     )

														lowerCAmelCase__      =			flatten_dict(unfreeze(model.params     )     )
														lowerCAmelCase__      =			flatten_dict(unfreeze(new_model.params     )     )

														for key in base_params.keys():
																					lowerCAmelCase__      =			(base_params[key] - new_params[key]).sum().item()
																					self.assertLessEqual(lowerCamelCase_      ,					1e-3      ,					msg=F"""{key} not identical"""     )


def       _snake_case     (  A				,					A )		->					Optional[int]:
							lowerCAmelCase__      =			True
							lowerCAmelCase__      =			flatten_dict(modela.params )
							lowerCAmelCase__      =			flatten_dict(modela.params )
							for key in flat_params_a.keys():
														if np.sum(np.abs(flat_params_a[key] - flat_params_a[key] ) ) > 1E-4:
																					lowerCAmelCase__      =			False

							return models_are_equal






@require_flax
class 	a__       (							unittest.TestCase						):





							'''simple docstring'''






							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
														lowerCAmelCase__      =			BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''     )
														lowerCAmelCase__      =			FlaxBertModel(lowerCamelCase_     )

														lowerCAmelCase__      =			'''bert'''
														with tempfile.TemporaryDirectory() as tmp_dir:
																					model.save_pretrained(os.path.join(lowerCamelCase_      ,					lowerCamelCase_     )     )

																					with self.assertRaises(lowerCamelCase_     ):
																												lowerCAmelCase__      =			FlaxBertModel.from_pretrained(lowerCamelCase_     )

																					lowerCAmelCase__      =			FlaxBertModel.from_pretrained(lowerCamelCase_      ,					subfolder=lowerCamelCase_     )

														self.assertTrue(check_models_equal(lowerCamelCase_      ,					lowerCamelCase_     )     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[Any]:
														lowerCAmelCase__      =			BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''     )
														lowerCAmelCase__      =			FlaxBertModel(lowerCamelCase_     )

														lowerCAmelCase__      =			'''bert'''
														with tempfile.TemporaryDirectory() as tmp_dir:
																					model.save_pretrained(os.path.join(lowerCamelCase_      ,					lowerCamelCase_     )      ,					max_shard_size='''10KB'''     )

																					with self.assertRaises(lowerCamelCase_     ):
																												lowerCAmelCase__      =			FlaxBertModel.from_pretrained(lowerCamelCase_     )

																					lowerCAmelCase__      =			FlaxBertModel.from_pretrained(lowerCamelCase_      ,					subfolder=lowerCamelCase_     )

														self.assertTrue(check_models_equal(lowerCamelCase_      ,					lowerCamelCase_     )     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
														lowerCAmelCase__      =			'''bert'''
														lowerCAmelCase__      =			'''hf-internal-testing/tiny-random-bert-subfolder'''

														with self.assertRaises(lowerCamelCase_     ):
																					lowerCAmelCase__      =			FlaxBertModel.from_pretrained(lowerCamelCase_     )

														lowerCAmelCase__      =			FlaxBertModel.from_pretrained(lowerCamelCase_      ,					subfolder=lowerCamelCase_     )

														self.assertIsNotNone(lowerCamelCase_     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[Any]:
														lowerCAmelCase__      =			'''bert'''
														lowerCAmelCase__      =			'''hf-internal-testing/tiny-random-bert-sharded-subfolder'''
														with self.assertRaises(lowerCamelCase_     ):
																					lowerCAmelCase__      =			FlaxBertModel.from_pretrained(lowerCamelCase_     )

														lowerCAmelCase__      =			FlaxBertModel.from_pretrained(lowerCamelCase_      ,					subfolder=lowerCamelCase_     )

														self.assertIsNotNone(lowerCamelCase_     )