


'''simple docstring'''


import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
				from transformers.models.bert.modeling_flax_bert import (
				    FlaxBertForMaskedLM,
				    FlaxBertForMultipleChoice,
				    FlaxBertForNextSentencePrediction,
				    FlaxBertForPreTraining,
				    FlaxBertForQuestionAnswering,
				    FlaxBertForSequenceClassification,
				    FlaxBertForTokenClassification,
				    FlaxBertModel,
				)





class     lowercase__  (						unittest.TestCase					):





							'''simple docstring'''


							def __init__(    self    ,					__snake_case    ,					__snake_case=13    ,					__snake_case=7    ,					__snake_case=True    ,					__snake_case=True    ,					__snake_case=True    ,					__snake_case=True    ,					__snake_case=99    ,					__snake_case=32    ,					__snake_case=5    ,					__snake_case=4    ,					__snake_case=37    ,					__snake_case="gelu"    ,					__snake_case=0.1    ,					__snake_case=0.1    ,					__snake_case=512    ,					__snake_case=16    ,					__snake_case=2    ,					__snake_case=0.02    ,					__snake_case=4    ,					):
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      parent
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      batch_size
											_SCREAMING_SNAKE_CASE :       Dict     		=      seq_length
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      is_training
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      use_attention_mask
											_SCREAMING_SNAKE_CASE :       List[str]     		=      use_token_type_ids
											_SCREAMING_SNAKE_CASE :       Dict     		=      use_labels
											_SCREAMING_SNAKE_CASE :       Any     		=      vocab_size
											_SCREAMING_SNAKE_CASE :       List[str]     		=      hidden_size
											_SCREAMING_SNAKE_CASE :       Any     		=      num_hidden_layers
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      num_attention_heads
											_SCREAMING_SNAKE_CASE :       List[str]     		=      intermediate_size
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      hidden_act
											_SCREAMING_SNAKE_CASE :       Tuple     		=      hidden_dropout_prob
											_SCREAMING_SNAKE_CASE :       int     		=      attention_probs_dropout_prob
											_SCREAMING_SNAKE_CASE :       Tuple     		=      max_position_embeddings
											_SCREAMING_SNAKE_CASE :       List[str]     		=      type_vocab_size
											_SCREAMING_SNAKE_CASE :       Any     		=      type_sequence_label_size
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      initializer_range
											_SCREAMING_SNAKE_CASE :       str     		=      num_choices






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      ids_tensor([self.batch_size, self.seq_length]    ,					self.vocab_size  )

											_SCREAMING_SNAKE_CASE :       Dict     		=      None
											if self.use_attention_mask:
															_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      random_attention_mask([self.batch_size, self.seq_length]  )

											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      None
											if self.use_token_type_ids:
															_SCREAMING_SNAKE_CASE :       List[Any]     		=      ids_tensor([self.batch_size, self.seq_length]    ,					self.type_vocab_size  )

											_SCREAMING_SNAKE_CASE :       Tuple     		=      BertConfig(
											    vocab_size=self.vocab_size    ,					hidden_size=self.hidden_size    ,					num_hidden_layers=self.num_hidden_layers    ,					num_attention_heads=self.num_attention_heads    ,					intermediate_size=self.intermediate_size    ,					hidden_act=self.hidden_act    ,					hidden_dropout_prob=self.hidden_dropout_prob    ,					attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,					max_position_embeddings=self.max_position_embeddings    ,					type_vocab_size=self.type_vocab_size    ,					is_decoder=__snake_case    ,					initializer_range=self.initializer_range    ,					)

											return config, input_ids, token_type_ids, attention_mask






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[str]     		=      self.prepare_config_and_inputs()
											_SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE :       Optional[int]     		=      config_and_inputs
											_SCREAMING_SNAKE_CASE :       Dict     		=      {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
											return config, inputs_dict






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      self.prepare_config_and_inputs()
											_SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE :       List[Any]     		=      config_and_inputs

											_SCREAMING_SNAKE_CASE :       Any     		=      True
											_SCREAMING_SNAKE_CASE :       Tuple     		=      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
											_SCREAMING_SNAKE_CASE :       List[str]     		=      ids_tensor([self.batch_size, self.seq_length]    ,					vocab_size=2  )

											return (
											    config,
											    input_ids,
											    attention_mask,
											    encoder_hidden_states,
											    encoder_attention_mask,
											)





@require_flax
class     lowercase__  (						_snake_case     ,     unittest.TestCase					):





							'''simple docstring'''


							A_							:      List[str]      	=  True

							A_							:      Optional[int]      	=  (
							    (
							        FlaxBertModel,
							        FlaxBertForPreTraining,
							        FlaxBertForMaskedLM,
							        FlaxBertForMultipleChoice,
							        FlaxBertForQuestionAnswering,
							        FlaxBertForNextSentencePrediction,
							        FlaxBertForSequenceClassification,
							        FlaxBertForTokenClassification,
							        FlaxBertForQuestionAnswering,
							    )
							    if is_flax_available()
							    else ()
							)






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      FlaxBertModelTester(self  )






							@slow
							def UpperCAmelCase_	(    self  ):
											# Only check this for base model, not necessary for all model classes.
											# This will also help speed-up tests.
											_SCREAMING_SNAKE_CASE :       Tuple     		=      FlaxBertModel.from_pretrained("""bert-base-cased"""  )
											_SCREAMING_SNAKE_CASE :       List[str]     		=      model(np.ones((1, 1)  )  )
											self.assertIsNotNone(__snake_case  )








'''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_		:   List[Any]							   =		logging.get_logger(__name__)







def 			snake_case_   (  SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				try:
								with open(SCREAMING_SNAKE_CASE__							,      """rb"""     ) as flax_state_f:
												_SCREAMING_SNAKE_CASE :       Dict     		=      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 			snake_case_   (  SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     ):







				"""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
				_SCREAMING_SNAKE_CASE :       List[Any]     		=      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."""     )
								_SCREAMING_SNAKE_CASE :       Dict     		=      jax.tree_util.tree_map(
								    lambda SCREAMING_SNAKE_CASE__     : params.astype(np.floataa     ) if params.dtype == jnp.bfloataa else params							,      SCREAMING_SNAKE_CASE__     )

				_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      """"""

				_SCREAMING_SNAKE_CASE :       str     		=      flatten_dict(SCREAMING_SNAKE_CASE__							,      sep="""."""     )
				_SCREAMING_SNAKE_CASE :       str     		=      pt_model.state_dict()

				# keep track of unexpected & missing keys
				_SCREAMING_SNAKE_CASE :       Tuple     		=      []
				_SCREAMING_SNAKE_CASE :       int     		=      set(pt_model_dict.keys()     )

				for flax_key_tuple, flax_tensor in flax_state_dict.items():
								_SCREAMING_SNAKE_CASE :       Any     		=      flax_key_tuple.split("""."""     )

								if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
												_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      flax_key_tuple_array[:-1] + ["""weight"""]
												_SCREAMING_SNAKE_CASE :       List[str]     		=      jnp.transpose(SCREAMING_SNAKE_CASE__							,      (3, 2, 0, 1)     )
								elif flax_key_tuple_array[-1] == "kernel":
												_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      flax_key_tuple_array[:-1] + ["""weight"""]
												_SCREAMING_SNAKE_CASE :       Any     		=      flax_tensor.T
								elif flax_key_tuple_array[-1] == "scale":
												_SCREAMING_SNAKE_CASE :       Optional[int]     		=      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__     ):
																_SCREAMING_SNAKE_CASE :       Optional[int]     		=      (
																    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"""     )
																)

								_SCREAMING_SNAKE_CASE :       Tuple     		=      """.""".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
																_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      np.asarray(SCREAMING_SNAKE_CASE__     ) if not isinstance(SCREAMING_SNAKE_CASE__							,      np.ndarray     ) else flax_tensor
																_SCREAMING_SNAKE_CASE :       int     		=      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
				_SCREAMING_SNAKE_CASE :       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




