





import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoImageProcessor, ViTImageProcessor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / '''utils'''))

from test_module.custom_image_processing import CustomImageProcessor  # noqa E402


lowercase_		:				str				=						get_tests_dir('''fixtures''')


class 			lowercase				(							unittest.TestCase     ):



						"""simple docstring"""

						def 		__UpperCAmelCase       (    self			:		List[str]					):





										'''simple docstring'''







										_snake_case :			Optional[Any]  						=   mock.Mock()
										_snake_case :			Dict  						=   5_00
										_snake_case :			Optional[int]  						=   {}
										_snake_case :			Optional[Any]  						=   HTTPError
										_snake_case :			str  						=   {}

										# Download this model to make sure it's in the cache.
										_snake_case :			List[Any]  						=   ViTImageProcessor.from_pretrained('hf-internal-testing/tiny-random-vit'					)
										# Under the mock environment we get a 500 error when trying to reach the model.
										with mock.patch('requests.Session.request' ,				return_value=lowerCamelCase_					) as mock_head:
														_snake_case :			Tuple  						=   ViTImageProcessor.from_pretrained('hf-internal-testing/tiny-random-vit'					)
														# This check we did call the fake head request
														mock_head.assert_called()







						def 		__UpperCAmelCase       (    self			:		Union[str, Any]					):





										'''simple docstring'''







										_snake_case :			Optional[Any]  						=   ViTImageProcessor.from_pretrained(
										    'https://huggingface.co/hf-internal-testing/tiny-random-vit/resolve/main/preprocessor_config.json'					)







						def 		__UpperCAmelCase       (    self			:		Optional[Any]					):





										'''simple docstring'''







										with self.assertRaises(lowerCamelCase_					):
														# config is in subfolder, the following should not work without specifying the subfolder
														_snake_case :			List[Any]  						=   AutoImageProcessor.from_pretrained('hf-internal-testing/stable-diffusion-all-variants'					)

										_snake_case :			List[str]  						=   AutoImageProcessor.from_pretrained(
										    'hf-internal-testing/stable-diffusion-all-variants' ,				subfolder='feature_extractor'					)

										self.assertIsNotNone(lowerCamelCase_					)



@is_staging_test
class 			lowercase				(							unittest.TestCase     ):



						"""simple docstring"""

						@classmethod
						def 		__UpperCAmelCase       (    cls			:		Tuple					):





										'''simple docstring'''







										_snake_case :			Tuple  						=   TOKEN
										HfFolder.save_token(lowerCamelCase_					)







						@classmethod
						def 		__UpperCAmelCase       (    cls			:		List[str]					):





										'''simple docstring'''







										try:
														delete_repo(token=cls._token ,				repo_id='test-image-processor'					)
										except HTTPError:
														pass

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

										try:
														delete_repo(token=cls._token ,				repo_id='test-dynamic-image-processor'					)
										except HTTPError:
														pass







						def 		__UpperCAmelCase       (    self			:		Dict					):





										'''simple docstring'''







										_snake_case :			Any  						=   ViTImageProcessor.from_pretrained(lowerCamelCase_					)
										image_processor.push_to_hub('test-image-processor' ,				use_auth_token=self._token					)

										_snake_case :			int  						=   ViTImageProcessor.from_pretrained(f'''{USER}/test-image-processor'''					)
										for k, v in image_processor.__dict__.items():
														self.assertEqual(lowerCamelCase_ ,				getattr(lowerCamelCase_ ,				lowerCamelCase_					)					)

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

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

										_snake_case :			Any  						=   ViTImageProcessor.from_pretrained(f'''{USER}/test-image-processor'''					)
										for k, v in image_processor.__dict__.items():
														self.assertEqual(lowerCamelCase_ ,				getattr(lowerCamelCase_ ,				lowerCamelCase_					)					)







						def 		__UpperCAmelCase       (    self			:		Union[str, Any]					):





										'''simple docstring'''







										_snake_case :			Tuple  						=   ViTImageProcessor.from_pretrained(lowerCamelCase_					)
										image_processor.push_to_hub('valid_org/test-image-processor' ,				use_auth_token=self._token					)

										_snake_case :			Any  						=   ViTImageProcessor.from_pretrained('valid_org/test-image-processor'					)
										for k, v in image_processor.__dict__.items():
														self.assertEqual(lowerCamelCase_ ,				getattr(lowerCamelCase_ ,				lowerCamelCase_					)					)

										# Reset repo
										delete_repo(token=self._token ,				repo_id='valid_org/test-image-processor'					)

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

										_snake_case :			Optional[int]  						=   ViTImageProcessor.from_pretrained('valid_org/test-image-processor-org'					)
										for k, v in image_processor.__dict__.items():
														self.assertEqual(lowerCamelCase_ ,				getattr(lowerCamelCase_ ,				lowerCamelCase_					)					)







						def 		__UpperCAmelCase       (    self			:		List[Any]					):





										'''simple docstring'''







										CustomImageProcessor.register_for_auto_class()
										_snake_case :			Dict  						=   CustomImageProcessor.from_pretrained(lowerCamelCase_					)

										image_processor.push_to_hub('test-dynamic-image-processor' ,				use_auth_token=self._token					)

										# This has added the proper auto_map field to the config
										self.assertDictEqual(
										    image_processor.auto_map ,				{'AutoImageProcessor': 'custom_image_processing.CustomImageProcessor'} ,				)

										_snake_case :			Optional[Any]  						=   AutoImageProcessor.from_pretrained(
										    f'''{USER}/test-dynamic-image-processor''' ,				trust_remote_code=lowerCamelCase_					)
										# Can't make an isinstance check because the new_image_processor is from the CustomImageProcessor class of a dynamic module
										self.assertEqual(new_image_processor.__class__.__name__ ,				'CustomImageProcessor'					)












import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
lowercase_		:				Optional[int]				=						object()

# For specifying empty leaf dict `{}`
lowercase_		:				List[Any]				=						object()

def       A__(		__lowerCAmelCase	,			__lowerCAmelCase ):
				_snake_case :			Optional[Any]  						=   tuple((re.compile(x + '$' ) for x in qs) )
				for i in range(len(__lowerCAmelCase ) - len(__lowerCAmelCase ) + 1 ):
								_snake_case :			Tuple  						=   [x.match(__lowerCAmelCase ) for x, y in zip(__lowerCAmelCase	,			ks[i:] )]
								if matches and all(__lowerCAmelCase ):
												return True
				return False

def       A__(		__lowerCAmelCase ):
				def replace(__lowerCAmelCase	,			__lowerCAmelCase ):
								for rule, replacement in rules:
												if _match(__lowerCAmelCase	,			__lowerCAmelCase ):
																return replacement
								return val

				return replace

def       A__(		):
				return [
				    # embeddings
				    (("transformer", "wpe", "embedding"), P('mp'	,			__lowerCAmelCase )),
				    (("transformer", "wte", "embedding"), P('mp'	,			__lowerCAmelCase )),
				    # atention
				    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(__lowerCAmelCase	,			'mp' )),
				    (("attention", "out_proj", "kernel"), P('mp'	,			__lowerCAmelCase )),
				    (("attention", "out_proj", "bias"), None),
				    # mlp
				    (("mlp", "c_fc", "kernel"), P(__lowerCAmelCase	,			'mp' )),
				    (("mlp", "c_fc", "bias"), P('mp' )),
				    (("mlp", "c_proj", "kernel"), P('mp'	,			__lowerCAmelCase )),
				    (("mlp", "c_proj", "bias"), None),
				    # layer norms
				    ((r"ln_\d+", "bias"), None),
				    ((r"\d+", r"ln_\d+", "scale"), None),
				    (("ln_f", "bias"), None),
				    (("ln_f", "scale"), None),
				]

def       A__(		__lowerCAmelCase ):
				_snake_case :			Optional[Any]  						=   _get_partition_rules()
				_snake_case :			Optional[Any]  						=   _replacement_rules(__lowerCAmelCase )
				_snake_case :			str  						=   {k: _unmatched for k in flatten_dict(__lowerCAmelCase )}
				_snake_case :			str  						=   {k: replace(__lowerCAmelCase	,			__lowerCAmelCase ) for k, v in initd.items()}
				assert _unmatched not in result.values(), "Incomplete partition spec."
				return freeze(unflatten_dict(__lowerCAmelCase ) )






