



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tensorflow_text_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowercase_						=       {
    """configuration_bert""": ["""BERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """BertConfig""", """BertOnnxConfig"""],
    """tokenization_bert""": ["""BasicTokenizer""", """BertTokenizer""", """WordpieceTokenizer"""],
}

try:
							if not is_tokenizers_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowercase_						=       ["""BertTokenizerFast"""]

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowercase_						=       [
							    """BERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
							    """BertForMaskedLM""",
							    """BertForMultipleChoice""",
							    """BertForNextSentencePrediction""",
							    """BertForPreTraining""",
							    """BertForQuestionAnswering""",
							    """BertForSequenceClassification""",
							    """BertForTokenClassification""",
							    """BertLayer""",
							    """BertLMHeadModel""",
							    """BertModel""",
							    """BertPreTrainedModel""",
							    """load_tf_weights_in_bert""",
							]

try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowercase_						=       [
							    """TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
							    """TFBertEmbeddings""",
							    """TFBertForMaskedLM""",
							    """TFBertForMultipleChoice""",
							    """TFBertForNextSentencePrediction""",
							    """TFBertForPreTraining""",
							    """TFBertForQuestionAnswering""",
							    """TFBertForSequenceClassification""",
							    """TFBertForTokenClassification""",
							    """TFBertLMHeadModel""",
							    """TFBertMainLayer""",
							    """TFBertModel""",
							    """TFBertPreTrainedModel""",
							]
try:
							if not is_tensorflow_text_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowercase_						=       ["""TFBertTokenizer"""]

try:
							if not is_flax_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowercase_						=       [
							    """FlaxBertForCausalLM""",
							    """FlaxBertForMaskedLM""",
							    """FlaxBertForMultipleChoice""",
							    """FlaxBertForNextSentencePrediction""",
							    """FlaxBertForPreTraining""",
							    """FlaxBertForQuestionAnswering""",
							    """FlaxBertForSequenceClassification""",
							    """FlaxBertForTokenClassification""",
							    """FlaxBertModel""",
							    """FlaxBertPreTrainedModel""",
							]

if TYPE_CHECKING:
							from .configuration_bert import BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, BertConfig, BertOnnxConfig
							from .tokenization_bert import BasicTokenizer, BertTokenizer, WordpieceTokenizer

							try:
														if not is_tokenizers_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_bert_fast import BertTokenizerFast

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_bert import (
														    BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
														    BertForMaskedLM,
														    BertForMultipleChoice,
														    BertForNextSentencePrediction,
														    BertForPreTraining,
														    BertForQuestionAnswering,
														    BertForSequenceClassification,
														    BertForTokenClassification,
														    BertLayer,
														    BertLMHeadModel,
														    BertModel,
														    BertPreTrainedModel,
														    load_tf_weights_in_bert,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tf_bert import (
														    TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
														    TFBertEmbeddings,
														    TFBertForMaskedLM,
														    TFBertForMultipleChoice,
														    TFBertForNextSentencePrediction,
														    TFBertForPreTraining,
														    TFBertForQuestionAnswering,
														    TFBertForSequenceClassification,
														    TFBertForTokenClassification,
														    TFBertLMHeadModel,
														    TFBertMainLayer,
														    TFBertModel,
														    TFBertPreTrainedModel,
														)

							try:
														if not is_tensorflow_text_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_bert_tf import TFBertTokenizer

							try:
														if not is_flax_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_flax_bert import (
														    FlaxBertForCausalLM,
														    FlaxBertForMaskedLM,
														    FlaxBertForMultipleChoice,
														    FlaxBertForNextSentencePrediction,
														    FlaxBertForPreTraining,
														    FlaxBertForQuestionAnswering,
														    FlaxBertForSequenceClassification,
														    FlaxBertForTokenClassification,
														    FlaxBertModel,
														    FlaxBertPreTrainedModel,
														)

else:
							import sys

							lowercase_						=       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)








import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


lowercase_						=       logging.get_logger(__name__)



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_					):
			lowercase__   					=			r"\w+[.]\d+"
			lowercase__   					=			re.findall(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)
			for pat in pats:
						lowercase__   					=			key.replace(SCREAMING_SNAKE_CASE_  ,					"_".join(pat.split("."					)					)					)
			return key



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					):
			lowercase__   					=			pt_tuple_key[:-1] + ("scale",)
			if (
			    any("norm" in str_ for str_ in pt_tuple_key					)
			    and (pt_tuple_key[-1] == "bias")
			    and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
			    and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
			):
						lowercase__   					=			pt_tuple_key[:-1] + ("scale",)
						return renamed_pt_tuple_key, pt_tensor
			elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
						lowercase__   					=			pt_tuple_key[:-1] + ("scale",)
						return renamed_pt_tuple_key, pt_tensor

			# embedding
			if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
						lowercase__   					=			pt_tuple_key[:-1] + ("embedding",)
						return renamed_pt_tuple_key, pt_tensor

			# conv layer
			lowercase__   					=			pt_tuple_key[:-1] + ("kernel",)
			if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
						lowercase__   					=			pt_tensor.transpose(2  ,					3  ,					1  ,					0					)
						return renamed_pt_tuple_key, pt_tensor

			# linear layer
			lowercase__   					=			pt_tuple_key[:-1] + ("kernel",)
			if pt_tuple_key[-1] == "weight":
						lowercase__   					=			pt_tensor.T
						return renamed_pt_tuple_key, pt_tensor

			# old PyTorch layer norm weight
			lowercase__   					=			pt_tuple_key[:-1] + ("weight",)
			if pt_tuple_key[-1] == "gamma":
						return renamed_pt_tuple_key, pt_tensor

			# old PyTorch layer norm bias
			lowercase__   					=			pt_tuple_key[:-1] + ("bias",)
			if pt_tuple_key[-1] == "beta":
						return renamed_pt_tuple_key, pt_tensor

			return pt_tuple_key, pt_tensor




def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_=42					):
			# Step 1: Convert pytorch tensor to numpy
			lowercase__   					=			{k: v.numpy() for k, v in pt_state_dict.items()}

			# Step 2: Since the model is stateless, get random Flax params
			lowercase__   					=			flax_model.init_weights(PRNGKey(SCREAMING_SNAKE_CASE_					)					)

			lowercase__   					=			flatten_dict(SCREAMING_SNAKE_CASE_					)
			lowercase__   					=			{}

			# Need to change some parameters name to match Flax names
			for pt_key, pt_tensor in pt_state_dict.items():
						lowercase__   					=			rename_key(SCREAMING_SNAKE_CASE_					)
						lowercase__   					=			tuple(renamed_pt_key.split("."					)					)

						# Correctly rename weight parameters
						lowercase__					,    lowercase__   					=			rename_key_and_reshape_tensor(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)

						if flax_key in random_flax_state_dict:
									if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
												raise ValueError(
												    f'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
												    f'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''					)

        # also add unexpected weight so that warning is thrown
						lowercase__   					=			jnp.asarray(SCREAMING_SNAKE_CASE_					)

			return unflatten_dict(SCREAMING_SNAKE_CASE_					)




