





"""simple docstring"""





import json
import os
import sys
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

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

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    PROCESSOR_MAPPING,
    TOKENIZER_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    AutoProcessor,
    AutoTokenizer,
    BertTokenizer,
    ProcessorMixin,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
)
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test
from transformers.tokenization_utils import TOKENIZER_CONFIG_FILE
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_tokenizers_available


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

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402
from test_module.custom_processing import CustomProcessor  # noqa E402
from test_module.custom_tokenization import CustomTokenizer  # noqa E402


_UpperCamelCase							=      get_tests_dir('fixtures/dummy_feature_extractor_config.json')
_UpperCamelCase							=      get_tests_dir('fixtures/vocab.json')
_UpperCamelCase							=      get_tests_dir('fixtures')
class 		SCREAMING_SNAKE_CASE_							(    unittest.TestCase   ):
				"""simple docstring"""
				__snake_case   :       List[str]	 =		["""[UNK]""", """[CLS]""", """[SEP]""", """[PAD]""", """[MASK]""", """bla""", """blou"""]




				def 					__lowercase  (     self  :Dict			):
											__lowerCamelCase      :      str         =0




				def 					__lowercase  (     self  :List[Any]			):
											__lowerCamelCase      :      Tuple         =AutoProcessor.from_pretrained('''facebook/wav2vec2-base-960h'''			)
											self.assertIsInstance(__lowercase						,    __lowercase			)




				def 					__lowercase  (     self  :Tuple			):
											with tempfile.TemporaryDirectory() as tmpdirname:
																		__lowerCamelCase      :      str         =WavaVecaConfig()
																		__lowerCamelCase      :      Dict         =AutoProcessor.from_pretrained('''facebook/wav2vec2-base-960h'''			)

																		# save in new folder
																		model_config.save_pretrained(__lowercase			)
																		processor.save_pretrained(__lowercase			)

																		__lowerCamelCase      :      Dict         =AutoProcessor.from_pretrained(__lowercase			)

											self.assertIsInstance(__lowercase						,    __lowercase			)




				def 					__lowercase  (     self  :Union[str, Any]			):
											with tempfile.TemporaryDirectory() as tmpdirname:
																		# copy relevant files
																		copyfile(__lowercase						,    os.path.join(__lowercase						,    __lowercase			)			)
																		copyfile(__lowercase						,    os.path.join(__lowercase						,    '''vocab.json'''			)			)

																		__lowerCamelCase      :      List[str]         =AutoProcessor.from_pretrained(__lowercase			)

											self.assertIsInstance(__lowercase						,    __lowercase			)




				def 					__lowercase  (     self  :List[str]			):
											with tempfile.TemporaryDirectory() as tmpdirname:
																		__lowerCamelCase      :      Dict         =WavaVecaFeatureExtractor()
																		__lowerCamelCase      :      Union[str, Any]         =AutoTokenizer.from_pretrained('''facebook/wav2vec2-base-960h'''			)

																		__lowerCamelCase      :      List[Any]         =WavaVecaProcessor(__lowercase						,    __lowercase			)

																		# save in new folder
																		processor.save_pretrained(__lowercase			)

																		# drop `processor_class` in tokenizer
																		with open(os.path.join(__lowercase						,    __lowercase			)						,    '''r'''			) as f:
																									__lowerCamelCase      :      int         =json.load(__lowercase			)
																									config_dict.pop('''processor_class'''			)

																		with open(os.path.join(__lowercase						,    __lowercase			)						,    '''w'''			) as f:
																									f.write(json.dumps(__lowercase			)			)

																		__lowerCamelCase      :      Union[str, Any]         =AutoProcessor.from_pretrained(__lowercase			)

											self.assertIsInstance(__lowercase						,    __lowercase			)




				def 					__lowercase  (     self  :Optional[int]			):
											with tempfile.TemporaryDirectory() as tmpdirname:
																		__lowerCamelCase      :      List[Any]         =WavaVecaFeatureExtractor()
																		__lowerCamelCase      :      Tuple         =AutoTokenizer.from_pretrained('''facebook/wav2vec2-base-960h'''			)

																		__lowerCamelCase      :      int         =WavaVecaProcessor(__lowercase						,    __lowercase			)

																		# save in new folder
																		processor.save_pretrained(__lowercase			)

																		# drop `processor_class` in feature extractor
																		with open(os.path.join(__lowercase						,    __lowercase			)						,    '''r'''			) as f:
																									__lowerCamelCase      :      int         =json.load(__lowercase			)
																									config_dict.pop('''processor_class'''			)

																		with open(os.path.join(__lowercase						,    __lowercase			)						,    '''w'''			) as f:
																									f.write(json.dumps(__lowercase			)			)

																		__lowerCamelCase      :      List[Any]         =AutoProcessor.from_pretrained(__lowercase			)

											self.assertIsInstance(__lowercase						,    __lowercase			)




				def 					__lowercase  (     self  :List[Any]			):
											with tempfile.TemporaryDirectory() as tmpdirname:
																		__lowerCamelCase      :      int         =WavaVecaConfig(processor_class='''Wav2Vec2Processor'''			)
																		model_config.save_pretrained(__lowercase			)
																		# copy relevant files
																		copyfile(__lowercase						,    os.path.join(__lowercase						,    '''vocab.json'''			)			)
																		# create emtpy sample processor
																		with open(os.path.join(__lowercase						,    __lowercase			)						,    '''w'''			) as f:
																									f.write('''{}'''			)

																		__lowerCamelCase      :      Union[str, Any]         =AutoProcessor.from_pretrained(__lowercase			)

											self.assertIsInstance(__lowercase						,    __lowercase			)




				def 					__lowercase  (     self  :List[str]			):
											# If remote code is not set, we will time out when asking whether to load the model.
											with self.assertRaises(__lowercase			):
																		__lowerCamelCase      :      str         =AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor'''			)
											# If remote code is disabled, we can't load this config.
											with self.assertRaises(__lowercase			):
																		__lowerCamelCase      :      str         =AutoProcessor.from_pretrained(
																		    '''hf-internal-testing/test_dynamic_processor'''						,    trust_remote_code=__lowercase			)

											__lowerCamelCase      :      str         =AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor'''						,    trust_remote_code=__lowercase			)
											self.assertTrue(processor.special_attribute_present			)
											self.assertEqual(processor.__class__.__name__						,    '''NewProcessor'''			)

											__lowerCamelCase      :      List[Any]         =processor.feature_extractor
											self.assertTrue(feature_extractor.special_attribute_present			)
											self.assertEqual(feature_extractor.__class__.__name__						,    '''NewFeatureExtractor'''			)

											__lowerCamelCase      :      Optional[int]         =processor.tokenizer
											self.assertTrue(tokenizer.special_attribute_present			)
											if is_tokenizers_available():
																		self.assertEqual(tokenizer.__class__.__name__						,    '''NewTokenizerFast'''			)

																		# Test we can also load the slow version
																		__lowerCamelCase      :      Optional[int]         =AutoProcessor.from_pretrained(
																		    '''hf-internal-testing/test_dynamic_processor'''						,    trust_remote_code=__lowercase						,    use_fast=__lowercase			)
																		__lowerCamelCase      :      Dict         =new_processor.tokenizer
																		self.assertTrue(new_tokenizer.special_attribute_present			)
																		self.assertEqual(new_tokenizer.__class__.__name__						,    '''NewTokenizer'''			)
											else:
																		self.assertEqual(tokenizer.__class__.__name__						,    '''NewTokenizer'''			)




				def 					__lowercase  (     self  :Union[str, Any]			):
											try:
																		AutoConfig.register('''custom'''						,    __lowercase			)
																		AutoFeatureExtractor.register(__lowercase						,    __lowercase			)
																		AutoTokenizer.register(__lowercase						,    slow_tokenizer_class=__lowercase			)
																		AutoProcessor.register(__lowercase						,    __lowercase			)
																		# Trying to register something existing in the Transformers library will raise an error
																		with self.assertRaises(__lowercase			):
																									AutoProcessor.register(__lowercase						,    __lowercase			)

																		# Now that the config is registered, it can be used as any other config with the auto-API
																		__lowerCamelCase      :      Optional[Any]         =CustomFeatureExtractor.from_pretrained(__lowercase			)

																		with tempfile.TemporaryDirectory() as tmp_dir:
																									__lowerCamelCase      :      Optional[Any]         =os.path.join(__lowercase						,    '''vocab.txt'''			)
																									with open(__lowercase						,    '''w'''						,    encoding='''utf-8'''			) as vocab_writer:
																																vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens]			)			)
																									__lowerCamelCase      :      str         =CustomTokenizer(__lowercase			)

																		__lowerCamelCase      :      int         =CustomProcessor(__lowercase						,    __lowercase			)

																		with tempfile.TemporaryDirectory() as tmp_dir:
																									processor.save_pretrained(__lowercase			)
																									__lowerCamelCase      :      Optional[Any]         =AutoProcessor.from_pretrained(__lowercase			)
																									self.assertIsInstance(__lowercase						,    __lowercase			)

											finally:
																		if "custom" in CONFIG_MAPPING._extra_content:
																									del CONFIG_MAPPING._extra_content["custom"]
																		if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
																									del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]
																		if CustomConfig in TOKENIZER_MAPPING._extra_content:
																									del TOKENIZER_MAPPING._extra_content[CustomConfig]
																		if CustomConfig in PROCESSOR_MAPPING._extra_content:
																									del PROCESSOR_MAPPING._extra_content[CustomConfig]




				def 					__lowercase  (     self  :List[str]			):
											class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
															"""simple docstring"""
															__snake_case   :       Dict	 =		False

											class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
															"""simple docstring"""
															__snake_case   :       Optional[int]	 =		False

											class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
															"""simple docstring"""
															__snake_case   :       List[str]	 =		"""AutoFeatureExtractor"""
															__snake_case   :       Optional[int]	 =		"""AutoTokenizer"""
															__snake_case   :       Union[str, Any]	 =		False

											try:
																		AutoConfig.register('''custom'''						,    __lowercase			)
																		AutoFeatureExtractor.register(__lowercase						,    __lowercase			)
																		AutoTokenizer.register(__lowercase						,    slow_tokenizer_class=__lowercase			)
																		AutoProcessor.register(__lowercase						,    __lowercase			)
																		# If remote code is not set, the default is to use local classes.
																		__lowerCamelCase      :      str         =AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor'''			)
																		self.assertEqual(processor.__class__.__name__						,    '''NewProcessor'''			)
																		self.assertFalse(processor.special_attribute_present			)
																		self.assertFalse(processor.feature_extractor.special_attribute_present			)
																		self.assertFalse(processor.tokenizer.special_attribute_present			)

																		# If remote code is disabled, we load the local ones.
																		__lowerCamelCase      :      str         =AutoProcessor.from_pretrained(
																		    '''hf-internal-testing/test_dynamic_processor'''						,    trust_remote_code=__lowercase			)
																		self.assertEqual(processor.__class__.__name__						,    '''NewProcessor'''			)
																		self.assertFalse(processor.special_attribute_present			)
																		self.assertFalse(processor.feature_extractor.special_attribute_present			)
																		self.assertFalse(processor.tokenizer.special_attribute_present			)

																		# If remote is enabled, we load from the Hub.
																		__lowerCamelCase      :      str         =AutoProcessor.from_pretrained(
																		    '''hf-internal-testing/test_dynamic_processor'''						,    trust_remote_code=__lowercase			)
																		self.assertEqual(processor.__class__.__name__						,    '''NewProcessor'''			)
																		self.assertTrue(processor.special_attribute_present			)
																		self.assertTrue(processor.feature_extractor.special_attribute_present			)
																		self.assertTrue(processor.tokenizer.special_attribute_present			)

											finally:
																		if "custom" in CONFIG_MAPPING._extra_content:
																									del CONFIG_MAPPING._extra_content["custom"]
																		if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
																									del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]
																		if CustomConfig in TOKENIZER_MAPPING._extra_content:
																									del TOKENIZER_MAPPING._extra_content[CustomConfig]
																		if CustomConfig in PROCESSOR_MAPPING._extra_content:
																									del PROCESSOR_MAPPING._extra_content[CustomConfig]




				def 					__lowercase  (     self  :Dict			):
											__lowerCamelCase      :      Tuple         =AutoProcessor.from_pretrained('''hf-internal-testing/tiny-random-bert'''			)
											self.assertEqual(processor.__class__.__name__						,    '''BertTokenizerFast'''			)




				def 					__lowercase  (     self  :Tuple			):
											__lowerCamelCase      :      List[Any]         =AutoProcessor.from_pretrained('''hf-internal-testing/tiny-random-convnext'''			)
											self.assertEqual(processor.__class__.__name__						,    '''ConvNextImageProcessor'''			)

@is_staging_test
class 		SCREAMING_SNAKE_CASE_							(    unittest.TestCase   ):
				"""simple docstring"""
				__snake_case   :       Dict	 =		["""[UNK]""", """[CLS]""", """[SEP]""", """[PAD]""", """[MASK]""", """bla""", """blou"""]




				@classmethod
				def 					__lowercase  (     cls  :Dict			):
											__lowerCamelCase      :      Optional[int]         =TOKEN
											HfFolder.save_token(__lowercase			)




				@classmethod
				def 					__lowercase  (     cls  :int			):
											try:
																		delete_repo(token=cls._token						,    repo_id='''test-processor'''			)
											except HTTPError:
																		pass

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

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




				def 					__lowercase  (     self  :Optional[Any]			):
											__lowerCamelCase      :      Optional[int]         =WavaVecaProcessor.from_pretrained(__lowercase			)
											with tempfile.TemporaryDirectory() as tmp_dir:
																		processor.save_pretrained(
																		    os.path.join(__lowercase						,    '''test-processor'''			)						,    push_to_hub=__lowercase						,    use_auth_token=self._token			)

																		__lowerCamelCase      :      Any         =WavaVecaProcessor.from_pretrained(f'{USER}/test-processor'			)
																		for k, v in processor.feature_extractor.__dict__.items():
																									self.assertEqual(__lowercase						,    getattr(new_processor.feature_extractor						,    __lowercase			)			)
																		self.assertDictEqual(new_processor.tokenizer.get_vocab()						,    processor.tokenizer.get_vocab()			)




				def 					__lowercase  (     self  :Union[str, Any]			):
											__lowerCamelCase      :      Tuple         =WavaVecaProcessor.from_pretrained(__lowercase			)

											with tempfile.TemporaryDirectory() as tmp_dir:
																		processor.save_pretrained(
																		    os.path.join(__lowercase						,    '''test-processor-org'''			)						,    push_to_hub=__lowercase						,    use_auth_token=self._token						,    organization='''valid_org'''						,    )

																		__lowerCamelCase      :      List[str]         =WavaVecaProcessor.from_pretrained('''valid_org/test-processor-org'''			)
																		for k, v in processor.feature_extractor.__dict__.items():
																									self.assertEqual(__lowercase						,    getattr(new_processor.feature_extractor						,    __lowercase			)			)
																		self.assertDictEqual(new_processor.tokenizer.get_vocab()						,    processor.tokenizer.get_vocab()			)




				def 					__lowercase  (     self  :List[Any]			):
											CustomFeatureExtractor.register_for_auto_class()
											CustomTokenizer.register_for_auto_class()
											CustomProcessor.register_for_auto_class()

											__lowerCamelCase      :      List[Any]         =CustomFeatureExtractor.from_pretrained(__lowercase			)

											with tempfile.TemporaryDirectory() as tmp_dir:
																		__lowerCamelCase      :      int         =os.path.join(__lowercase						,    '''vocab.txt'''			)
																		with open(__lowercase						,    '''w'''						,    encoding='''utf-8'''			) as vocab_writer:
																									vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens]			)			)
																		__lowerCamelCase      :      List[Any]         =CustomTokenizer(__lowercase			)

											__lowerCamelCase      :      Optional[int]         =CustomProcessor(__lowercase						,    __lowercase			)

											with tempfile.TemporaryDirectory() as tmp_dir:
																		create_repo(f'{USER}/test-dynamic-processor'						,    token=self._token			)
																		__lowerCamelCase      :      int         =Repository(__lowercase						,    clone_from=f'{USER}/test-dynamic-processor'						,    token=self._token			)
																		processor.save_pretrained(__lowercase			)

																		# This has added the proper auto_map field to the feature extractor config
																		self.assertDictEqual(
																		    processor.feature_extractor.auto_map						,    {
																		        '''AutoFeatureExtractor''': '''custom_feature_extraction.CustomFeatureExtractor''',
																		        '''AutoProcessor''': '''custom_processing.CustomProcessor''',
																		    }						,    )

																		# This has added the proper auto_map field to the tokenizer config
																		with open(os.path.join(__lowercase						,    '''tokenizer_config.json'''			)			) as f:
																									__lowerCamelCase      :      Optional[int]         =json.load(__lowercase			)
																		self.assertDictEqual(
																		    tokenizer_config['''auto_map''']						,    {
																		        '''AutoTokenizer''': ['''custom_tokenization.CustomTokenizer''', None],
																		        '''AutoProcessor''': '''custom_processing.CustomProcessor''',
																		    }						,    )

																		# The code has been copied from fixtures
																		self.assertTrue(os.path.isfile(os.path.join(__lowercase						,    '''custom_feature_extraction.py'''			)			)			)
																		self.assertTrue(os.path.isfile(os.path.join(__lowercase						,    '''custom_tokenization.py'''			)			)			)
																		self.assertTrue(os.path.isfile(os.path.join(__lowercase						,    '''custom_processing.py'''			)			)			)

																		repo.push_to_hub()

											__lowerCamelCase      :      Dict         =AutoProcessor.from_pretrained(f'{USER}/test-dynamic-processor'						,    trust_remote_code=__lowercase			)
											# Can't make an isinstance check because the new_processor is from the CustomProcessor class of a dynamic module
											self.assertEqual(new_processor.__class__.__name__						,    '''CustomProcessor'''			)











"""simple docstring"""





from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


_UpperCamelCase							=      logging.get_logger(__name__)



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							str    ,						SCREAMING_SNAKE_CASE       :							str   ):


							'''simple docstring'''







							try:
														with open(SCREAMING_SNAKE_CASE    ,						'''rb'''   ) as flax_state_f:
																					__lowerCamelCase      :      int         =from_bytes(SCREAMING_SNAKE_CASE    ,						flax_state_f.read()   )
							except UnpicklingError as e:
														try:
																					with open(SCREAMING_SNAKE_CASE   ) as f:
																												if f.read().startswith('''version'''   ):
																																			raise OSError(
																																			    '''You seem to have cloned a repository without having git-lfs installed. Please'''
																																			    ''' install git-lfs and run `git lfs install` followed by `git lfs pull` in the'''
																																			    ''' folder you cloned.'''   )
																												else:
																																			raise ValueError from e
														except (UnicodeDecodeError, ValueError):
																					raise EnvironmentError(F'Unable to convert {model_file} to Flax deserializable object. '   )

							return load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							int    ,						SCREAMING_SNAKE_CASE       :							Union[str, Any]   ):


							'''simple docstring'''







							try:
														import torch  # noqa: F401
							except ImportError:
														logger.error(
														    '''Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see'''
														    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
														    ''' instructions.'''   )
														raise

							# check if we have bf16 weights
							__lowerCamelCase      :      str         =flatten_dict(jax.tree_util.tree_map(lambda SCREAMING_SNAKE_CASE   : x.dtype == jnp.bfloataa    ,						SCREAMING_SNAKE_CASE   )   ).values()
							if any(SCREAMING_SNAKE_CASE   ):
														# convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

														# and bf16 is not fully supported in PT yet.
														logger.warning(
														    '''Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` '''
														    '''before loading those in PyTorch model.'''   )
														__lowerCamelCase      :      List[str]         =jax.tree_util.tree_map(
														    lambda SCREAMING_SNAKE_CASE   : params.astype(np.floataa   ) if params.dtype == jnp.bfloataa else params    ,						SCREAMING_SNAKE_CASE   )

							__lowerCamelCase      :      int         =''''''

							__lowerCamelCase      :      List[str]         =flatten_dict(SCREAMING_SNAKE_CASE    ,						sep='''.'''   )
							__lowerCamelCase      :      List[str]         =pt_model.state_dict()

							# keep track of unexpected & missing keys
							__lowerCamelCase      :      Union[str, Any]         =[]
							__lowerCamelCase      :      int         =set(pt_model_dict.keys()   )

							for flax_key_tuple, flax_tensor in flax_state_dict.items():
														__lowerCamelCase      :      Tuple         =flax_key_tuple.split('''.'''   )

														if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
																					__lowerCamelCase      :      Optional[int]         =flax_key_tuple_array[:-1] + ['''weight''']
																					__lowerCamelCase      :      Union[str, Any]         =jnp.transpose(SCREAMING_SNAKE_CASE    ,						(3, 2, 0, 1)   )
														elif flax_key_tuple_array[-1] == "kernel":
																					__lowerCamelCase      :      Optional[Any]         =flax_key_tuple_array[:-1] + ['''weight''']
																					__lowerCamelCase      :      Dict         =flax_tensor.T
														elif flax_key_tuple_array[-1] == "scale":
																					__lowerCamelCase      :      str         =flax_key_tuple_array[:-1] + ['''weight''']

														if "time_embedding" not in flax_key_tuple_array:
																					for i, flax_key_tuple_string in enumerate(SCREAMING_SNAKE_CASE   ):
																												__lowerCamelCase      :      Union[str, Any]         =(
																												    flax_key_tuple_string.replace('''_0'''    ,						'''.0'''   )
																												    .replace('''_1'''    ,						'''.1'''   )
																												    .replace('''_2'''    ,						'''.2'''   )
																												    .replace('''_3'''    ,						'''.3'''   )
																												    .replace('''_4'''    ,						'''.4'''   )
																												    .replace('''_5'''    ,						'''.5'''   )
																												    .replace('''_6'''    ,						'''.6'''   )
																												    .replace('''_7'''    ,						'''.7'''   )
																												    .replace('''_8'''    ,						'''.8'''   )
																												    .replace('''_9'''    ,						'''.9'''   )
																												)

														__lowerCamelCase      :      Optional[int]         ='''.'''.join(SCREAMING_SNAKE_CASE   )

														if flax_key in pt_model_dict:
																					if flax_tensor.shape != pt_model_dict[flax_key].shape:
																												raise ValueError(
																												    F'Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '
																												    F'to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'   )
																					else:
																												# add weight to pytorch dict
																												__lowerCamelCase      :      Tuple         =np.asarray(SCREAMING_SNAKE_CASE   ) if not isinstance(SCREAMING_SNAKE_CASE    ,						np.ndarray   ) else flax_tensor
																												__lowerCamelCase      :      Any         =torch.from_numpy(SCREAMING_SNAKE_CASE   )
																												# remove from missing keys
																												missing_keys.remove(SCREAMING_SNAKE_CASE   )
														else:
																					# weight is not expected by PyTorch model
																					unexpected_keys.append(SCREAMING_SNAKE_CASE   )

							pt_model.load_state_dict(SCREAMING_SNAKE_CASE   )

							# re-transform missing_keys to list
							__lowerCamelCase      :      Optional[Any]         =list(SCREAMING_SNAKE_CASE   )

							if len(SCREAMING_SNAKE_CASE   ) > 0:
														logger.warning(
														    '''Some weights of the Flax model were not used when initializing the PyTorch model'''
														    F' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'
														    F' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'
														    ''' (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This'''
														    F' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'
														    ''' to be exactly identical (e.g. initializing a BertForSequenceClassification model from a'''
														    ''' FlaxBertForSequenceClassification model).'''   )
							if len(SCREAMING_SNAKE_CASE   ) > 0:
														logger.warning(
														    F'Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'
														    F' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'
														    ''' use it for predictions and inference.'''   )

							return pt_model




