codingutf8 2018 the google ai language team s and the huggingface inc team c 2018 nvidia corporation licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch tf 2 0 general utilities import os import re import numpy from utils import explicitenum expanddims isnumpyarray istorchtensor logging reshape squeeze tensorsize from utils import transpose as transposefunc logger logging getloggername class transposetypeexplicitenum no no simple simple conv1d conv1d conv2d conv2d def converttfweightnametoptweightname tfname startprefixtoremove tfweightshapenone namescopenone if namescope is not none if not tfname startswithnamescope raise valueerror fweight name tfname does not start with namescope namescope this is an internal error in transformers so unless you were doing something really evil please open an issue to report it tfname tfnamelennamescope tfname tfname lstrip tfname tfname replace 0 device ids tfname re sub r r1 tfname 12 is replaced by 2 can be used to duplicate or remove layers in tf2 0 vs pytorch tfname tfname replace is replaced by a level separation can be used to convert tf2 0 lists in pytorch nn moduleslist tfname re subr tfname remove empty levels at the end tfname tfname split convert from tf2 0 separators to pytorch separators some weights have a single name without such as finallogitsbias in bart if lentfname 1 tfname tfname1 remove level zero tfweightshape listtfweightshape when should we transpose the weights if tfname1 kernel and tfweightshape is not none and lentfweightshape 4 transpose transposetype conv2d elif tfname1 kernel and tfweightshape is not none and lentfweightshape 3 transpose transposetype conv1d elif bool tfname1 in kernel pointwisekernel depthwisekernel or embprojs in tfname or outprojs in tfname transpose transposetype simple else transpose transposetype no convert standard tf2 0 names in pytorch names if tfname1 kernel or tfname1 embeddings or tfname1 gamma tfname1 weight if tfname1 beta tfname1 bias the separableconv1d tf layer contains two weights that are translated to pytorch conv1d here if tfname1 pointwisekernel or tfname1 depthwisekernel tfname1 tfname1 replacekernel weight remove prefix if needed tfname jointfname if startprefixtoremove tfname tfname replacestartprefixtoremove 1 return tfname transpose def applytransposetranspose transposetype weight matchshapenone pttotftrue if transpose is transposetype conv2d conv2d weight pt numoutchannel numinchannel kernel0 kernel1 tf kernel0 kernel1 numinchannel numoutchannel axes 2 3 1 0 if pttotf else 3 2 0 1 weight transposefuncweight axesaxes elif transpose is transposetype conv1d conv1d weight pt numoutchannel numinchannel kernel tf kernel numinchannel numoutchannel weight transposefuncweight axes2 1 0 elif transpose is transposetype simple weight transposefuncweight if matchshape is none return weight if lenmatchshape lenweight shape weight squeezeweight elif lenmatchshape lenweight shape weight expanddimsweight axis0 if listmatchshape listweight shape try weight reshapeweight matchshape except assertionerror as e e args matchshape matchshape raise e return weight pytorch tf 2 0 def loadpytorchcheckpointintf2model tfmodel pytorchcheckpointpath tfinputsnone allowmissingkeysfalse outputloadinginfofalse prefixnone tftoptweightrenamenone treats a single file as a collection of shards with 1 shard loads all shards into a single state dictionary load pytorch checkpoints in a tf 2 0 model ptstatedict ptmodel statedict return loadpytorchweightsintf2model tfmodel ptstatedict tfinputstfinputs allowmissingkeysallowmissingkeys def loadpytorchweightsintf2model tfmodel ptstatedict tfinputsnone allowmissingkeysfalse outputloadinginfofalse prefixnone tftoptweightrenamenone load a pytorch statedict in a tf 2 0 model ptstatedict can be either an actual dict or a lazyloading safetensors archive created with the safeopen function import tensorflow as tf from keras import backend as k if tfinputs is none tfinputs tfmodel dummyinputs if prefix is none prefix if tfinputs with tf namescopeprefix tfmodeltfinputs trainingfalse make sure model is built convert old format to new format if needed from a pytorch statedict tfkeystoptkeys for key in ptstatedict keys newkey none if gamma in key newkey key replacegamma weight if beta in key newkey key replacebeta bias if runningvar in key newkey key replacerunningvar movingvariance if runningmean in key newkey key replacerunningmean movingmean new weightnorm from https github comhuggingfacetransformerspull24030 keycomponents key split name none if keycomponents3 2 parametrizations original0 name keycomponents2 g elif keycomponents3 2 parametrizations original1 name keycomponents2 v if name is not none keycomponents keycomponents 3 name newkey joinkeycomponents if newkey is none newkey key tfkeystoptkeysnewkey key matt all tf models store the actual model stem in a mainlayer class including the base model in pt the derived models with heads use the base model class as the stem instead and there is no mainlayer class this means that tf base classes have one extra layer in their weight names corresponding to the mainlayer class this code block compensates for that startprefixtoremove if not anys startswithtfmodel basemodelprefix for s in tfkeystoptkeys keys startprefixtoremove tfmodel basemodelprefix symbolicweights tfmodel trainableweights tfmodel nontrainableweights tfloadednumel 0 allpytorchweights settfkeystoptkeys keys missingkeys mismatchedkeys issafetensorarchive hasattrptstatedict gettensor for symbolicweight in symbolicweights swname symbolicweight name name transpose converttfweightnametoptweightname swname startprefixtoremovestartprefixtoremove tfweightshapesymbolicweight shape namescopeprefix if tftoptweightrename is not none name tftoptweightrenamename find associated numpy array in pytorch model state dict if name not in tfkeystoptkeys if allowmissingkeys missingkeys appendname continue elif tfmodel keystoignoreonloadmissing is not none ized missing keys don t have to be loaded if anyre searchpat name is not none for pat in tfmodel keystoignoreonloadmissing continue raise attributeerrorfname not found in pytorch model statedictname tfkeystoptkeysname if issafetensorarchive array ptstatedict gettensorstatedictname else array ptstatedictstatedictname try array applytransposetranspose array symbolicweight shape except tf errors invalidargumenterror as e if not ignoremismatchedsizes errormsg stre errormsg ntyou may consider adding ignoremismatchedsizestrue in the model frompretrained method raise tf errors invalidargumenterrorerrormsg else mismatchedkeys appendname array shape symbolicweight shape continue tfloadednumel tensorsizearray k setvaluesymbolicweight array del array immediately free memory to keep peak usage as low as possible allpytorchweights discardname logger infofloaded tfloadednumel parameters in the tf 2 0 model unexpectedkeys listallpytorchweights if tfmodel keystoignoreonloadmissing is not none for pat in tfmodel keystoignoreonloadmissing missingkeys k for k in missingkeys if re searchpat k is none if tfmodel keystoignoreonloadunexpected is not none for pat in tfmodel keystoignoreonloadunexpected unexpectedkeys k for k in unexpectedkeys if re searchpat k is none if lenunexpectedkeys 0 logger warning some weights of the pytorch model were not used when initializing the tf 2 0 model f tfmodel class name unexpectedkeysn this is expected if you are initializing f tfmodel class name from a pytorch model trained on another task or with another architecture e g initializing a tfbertforsequenceclassification model from a bertforpretraining model n this is f not expected if you are initializing tfmodel class name from a pytorch model that you expect to be exactly identical e g initializing a tfbertforsequenceclassification model from a bertforsequenceclassification model else logger warningfall pytorch model weights were used when initializing tfmodel class name n if lenmissingkeys 0 logger warning fsome weights or buffers of the tf 2 0 model tfmodel class name were not initialized from the f pytorch model and are newly initialized missingkeysnyou should probably train this model on a downstream task to be able to use it for predictions and inference else logger warning fall the weights of tfmodel class name were initialized from the pytorch model n if your task is similar to the task the model of the checkpoint was trained on fyou can already use tfmodel class name for predictions without further training if lenmismatchedkeys 0 mismatchedwarning n join f key found shape shape1 in the checkpoint and shape2 in the model instantiated for key shape1 shape2 in mismatchedkeys logger warning fsome weights of tfmodel class name were not initialized from the model checkpoint f are newly initialized because the shapes did not f match nmismatchedwarningnyou should probably train this model on a downstream task to be able to use it for predictions and inference if outputloadinginfo loadinginfo missingkeys missingkeys unexpectedkeys unexpectedkeys mismatchedkeys mismatchedkeys return tfmodel loadinginfo return tfmodel tf 2 0 pytorch def loadtf2checkpointinpytorchmodel ptmodel tfcheckpointpath tfinputsnone allowmissingkeysfalse outputloadinginfofalse try import tensorflow as tf noqa f401 import torch noqa f401 except importerror logger error loading a tensorflow model in pytorch requires both pytorch and tensorflow to be installed please see https pytorch org and https www tensorflow orginstall for installation instructions raise import transformers from modelingtfutils import loadtfweights logger infofloading tensorflow weights from tfcheckpointpath instantiate and load the associated tf 2 0 model tfmodelclassname tf ptmodel class name add tf at the beginning tfmodelclass getattrtransformers tfmodelclassname tfmodel tfmodelclassptmodel config if tfinputs is none tfinputs tfmodel dummyinputs if tfinputs is not none tfmodeltfinputs trainingfalse make sure model is built loadtfweightstfmodel tfcheckpointpath return loadtf2modelinpytorchmodel ptmodel tfmodel allowmissingkeysallowmissingkeys outputloadinginfooutputloadinginfo def loadtf2modelinpytorchmodelptmodel tfmodel allowmissingkeysfalse outputloadinginfofalse load tf2 0 symbolic weights in a pytorch model try import tensorflow as tf noqa f401 import torch noqa f401 except importerror logger error loading a tensorflow model in pytorch requires both pytorch and tensorflow to be installed please see https pytorch org and https www tensorflow orginstall for installation instructions raise tfstatedict tfweight name tfweight numpy for tfweight in tfweights return loadtf2statedictinpytorchmodel ptmodel tfstatedict allowmissingkeysallowmissingkeys outputloadinginfooutputloadinginfo def loadtf2statedictinpytorchmodelptmodel tfstatedict allowmissingkeysfalse outputloadinginfofalse import torch newptparamsdict currentptparamsdict dictptmodel namedparameters make sure we are able to load pytorch base models as well as derived models with heads tf models always have a prefix some of pytorch models base ones don t startprefixtoremove if not anys startswithptmodel basemodelprefix for s in currentptparamsdict keys startprefixtoremove ptmodel basemodelprefix build a map from potential pytorch weight names to tf 2 0 variables tfweightsmap for name tfweight in tfstatedict items ptname transpose converttfweightnametoptweightname name startprefixtoremovestartprefixtoremove tfweightshapetfweight shape tfweightsmapptname tfweight transpose alltfweights settfweightsmap keys loadedptweightsdataptr missingkeyspt for ptweightname ptweight in currentptparamsdict items handle pytorch shared weight not duplicated in tf 2 0 if ptweight dataptr in loadedptweightsdataptr newptparamsdictptweightname loadedptweightsdataptrptweight dataptr continue ptweightnametocheck ptweightname new weightnorm from https github comhuggingfacetransformerspull24030 keycomponents ptweightname split name none if keycomponents3 2 parametrizations original0 name keycomponents2 g elif keycomponents3 2 parametrizations original1 name keycomponents2 v if name is not none keycomponents keycomponents 3 name ptweightnametocheck joinkeycomponents find associated numpy array in pytorch model state dict if ptweightnametocheck not in tfweightsmap if allowmissingkeys missingkeyspt appendptweightname continue raise attributeerrorfptweightname not found in tf 2 0 model array transpose tfweightsmapptweightnametocheck array applytransposetranspose array ptweight shape pttotffalse if numpy isscalararray array numpy arrayarray if not istorchtensorarray and not isnumpyarrayarray array array numpy if isnumpyarrayarray convert to torch tensor array torch fromnumpyarray newptparamsdictptweightname array loadedptweightsdataptrptweight dataptr array alltfweights discardptweightname missingkeys unexpectedkeys ptmodel loadstatedictnewptparamsdict strictfalse missingkeys missingkeyspt some models may have keys that are not in the state by design removing them before needlessly warning the user if ptmodel keystoignoreonloadmissing is not none for pat in ptmodel keystoignoreonloadmissing missingkeys k for k in missingkeys if re searchpat k is none if ptmodel keystoignoreonloadunexpected is not none for pat in ptmodel keystoignoreonloadunexpected unexpectedkeys k for k in unexpectedkeys if re searchpat k is none if lenunexpectedkeys 0 logger warning some weights of the tf 2 0 model were not used when initializing the pytorch model f ptmodel class name unexpectedkeysn this is expected if you are initializing f ptmodel class name from a tf 2 0 model trained on another task or with another architecture e g initializing a bertforsequenceclassification model from a tfbertforpretraining model n this is f not expected if you are initializing ptmodel class name from a tf 2 0 model that you expect to be exactly identical e g initializing a bertforsequenceclassification model from a tfbertforsequenceclassification model else logger warningfall tf 2 0 model weights were used when initializing ptmodel class name n if lenmissingkeys 0 logger warning fsome weights of ptmodel class name were not initialized from the tf 2 0 model and are newly f initialized missingkeysnyou should probably train this model on a downstream task to be able to use it for predictions and inference else logger warning fall the weights of ptmodel class name were initialized from the tf 2 0 model n if your task is similar to the task the model of the checkpoint was trained on fyou can already use ptmodel class name for predictions without further training logger infofweights or buffers not loaded from tf 2 0 model alltfweights if outputloadinginfo loadinginfo missingkeys missingkeys unexpectedkeys unexpectedkeys return ptmodel loadinginfo return ptmodel coding utf 8 2018 the google ai language team s and the huggingface inc team c 2018 nvidia corporation licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch tf 2 0 general utilities possible convert a tf 2 0 model variable name in a pytorch model weight name conventions for tf2 0 scopes pytorch attribute names conversions 1___ 2 is replaced by 2 can be used to duplicate or remove layers in tf2 0 vs pytorch _ _ is replaced by a new level separation can be used to convert tf2 0 lists in pytorch nn moduleslist return tuple with pytorch model weight name transpose transposetype member indicating whether and how tf2 0 and pytorch weights matrices should be transposed with regards to each other device ids 1___ 2 is replaced by 2 can be used to duplicate or remove layers in tf2 0 vs pytorch _ _ is replaced by a level separation can be used to convert tf2 0 lists in pytorch nn moduleslist remove empty levels at the end convert from tf2 0 separators to pytorch separators some weights have a single name without such as final_logits_bias in bart remove level zero when should we transpose the weights convert standard tf2 0 names in pytorch names the separableconv1d tf layer contains two weights that are translated to pytorch conv1d here remove prefix if needed apply a transpose to some weight then tries to reshape the weight to the same shape as a given shape all in a framework agnostic way conv2d weight pt num_out_channel num_in_channel kernel 0 kernel 1 tf kernel 0 kernel 1 num_in_channel num_out_channel conv1d weight pt num_out_channel num_in_channel kernel tf kernel num_in_channel num_out_channel pytorch tf 2 0 load pytorch checkpoints in a tf 2 0 model noqa f401 noqa f401 noqa f401 treats a single file as a collection of shards with 1 shard loads all shards into a single state dictionary load pytorch checkpoints in a tf 2 0 model load pytorch state_dict in a tf 2 0 model noqa f401 noqa f401 load a pytorch state_dict in a tf 2 0 model pt_state_dict can be either an actual dict or a lazy loading safetensors archive created with the safe_open function make sure model is built convert old format to new format if needed from a pytorch state_dict new weight_norm from https github com huggingface transformers pull 24030 matt all tf models store the actual model stem in a mainlayer class including the base model in pt the derived models with heads use the base model class as the stem instead and there is no mainlayer class this means that tf base classes have one extra layer in their weight names corresponding to the mainlayer class this code block compensates for that find associated numpy array in pytorch model state dict ized missing keys don t have to be loaded immediately free memory to keep peak usage as low as possible tf 2 0 pytorch load tf 2 0 hdf5 checkpoint in a pytorch model we use hdf5 to easily do transfer learning see https github com tensorflow tensorflow blob ee16fcac960ae660e0e4496658a366e2f745e1f0 tensorflow python keras engine network py l1352 l1357 noqa f401 noqa f401 instantiate and load the associated tf 2 0 model add tf at the beginning make sure model is built load tf 2 0 model in a pytorch model load tf2 0 symbolic weights in a pytorch model noqa f401 noqa f401 make sure we are able to load pytorch base models as well as derived models with heads tf models always have a prefix some of pytorch models base ones don t build a map from potential pytorch weight names to tf 2 0 variables handle pytorch shared weight not duplicated in tf 2 0 new weight_norm from https github com huggingface transformers pull 24030 find associated numpy array in pytorch model state dict convert to torch tensor some models may have keys that are not in the state by design removing them before needlessly warning the user
import os
import re

import numpy

from .utils import ExplicitEnum, expand_dims, is_numpy_array, is_torch_tensor, logging, reshape, squeeze, tensor_size
from .utils import transpose as transpose_func


logger = logging.get_logger(__name__)


class TransposeType(ExplicitEnum):
    

    NO = "no"
    SIMPLE = "simple"
    CONV1D = "conv1d"
    CONV2D = "conv2d"


def convert_tf_weight_name_to_pt_weight_name(
    tf_name, start_prefix_to_remove="", tf_weight_shape=None, name_scope=None
):
    
    if name_scope is not None:
        if not tf_name.startswith(name_scope):
            raise ValueError(
                f"Weight name {tf_name} does not start with name_scope {name_scope}. This is an internal error "
                "in Transformers, so (unless you were doing something really evil) please open an issue to report it!"
            )
        tf_name = tf_name[len(name_scope) :]
        tf_name = tf_name.lstrip("/")
    tf_name = tf_name.replace(":0", "")  
    tf_name = re.sub(
        r"/[^/]*___([^/]*)/", r"/\1/", tf_name
    )  
    tf_name = tf_name.replace(
        "_._", "/"
    )  
    tf_name = re.sub(r"//+", "/", tf_name)  
    tf_name = tf_name.split("/")  
    
    if len(tf_name) > 1:
        tf_name = tf_name[1:]  

    tf_weight_shape = list(tf_weight_shape)

    
    if tf_name[-1] == "kernel" and tf_weight_shape is not None and len(tf_weight_shape) == 4:
        transpose = TransposeType.CONV2D
    elif tf_name[-1] == "kernel" and tf_weight_shape is not None and len(tf_weight_shape) == 3:
        transpose = TransposeType.CONV1D
    elif bool(
        tf_name[-1] in ["kernel", "pointwise_kernel", "depthwise_kernel"]
        or "emb_projs" in tf_name
        or "out_projs" in tf_name
    ):
        transpose = TransposeType.SIMPLE
    else:
        transpose = TransposeType.NO

    
    if tf_name[-1] == "kernel" or tf_name[-1] == "embeddings" or tf_name[-1] == "gamma":
        tf_name[-1] = "weight"
    if tf_name[-1] == "beta":
        tf_name[-1] = "bias"

    
    if tf_name[-1] == "pointwise_kernel" or tf_name[-1] == "depthwise_kernel":
        tf_name[-1] = tf_name[-1].replace("_kernel", ".weight")

    
    tf_name = ".".join(tf_name)
    if start_prefix_to_remove:
        tf_name = tf_name.replace(start_prefix_to_remove, "", 1)

    return tf_name, transpose


def apply_transpose(transpose: TransposeType, weight, match_shape=None, pt_to_tf=True):
    
    if transpose is TransposeType.CONV2D:
        
        
        
        axes = (2, 3, 1, 0) if pt_to_tf else (3, 2, 0, 1)
        weight = transpose_func(weight, axes=axes)
    elif transpose is TransposeType.CONV1D:
        
        
        
        weight = transpose_func(weight, axes=(2, 1, 0))
    elif transpose is TransposeType.SIMPLE:
        weight = transpose_func(weight)

    if match_shape is None:
        return weight

    if len(match_shape) < len(weight.shape):
        weight = squeeze(weight)
    elif len(match_shape) > len(weight.shape):
        weight = expand_dims(weight, axis=0)

    if list(match_shape) != list(weight.shape):
        try:
            weight = reshape(weight, match_shape)
        except AssertionError as e:
            e.args += (match_shape, match_shape)
            raise e

    return weight







def load_pytorch_checkpoint_in_tf2_model(
    tf_model,
    pytorch_checkpoint_path,
    tf_inputs=None,
    allow_missing_keys=False,
    output_loading_info=False,
    _prefix=None,
    tf_to_pt_weight_rename=None,
):
    
    try:
        import tensorflow as tf  
        import torch  
        from safetensors.torch import load_file as safe_load_file  
    except ImportError:
        logger.error(
            "Loading a PyTorch model in TensorFlow, requires both PyTorch and TensorFlow to be installed. Please see "
            "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions."
        )
        raise

    
    if isinstance(pytorch_checkpoint_path, str):
        pytorch_checkpoint_path = [pytorch_checkpoint_path]

    
    pt_state_dict = {}
    for path in pytorch_checkpoint_path:
        pt_path = os.path.abspath(path)
        logger.info(f"Loading PyTorch weights from {pt_path}")
        if pt_path.endswith(".safetensors"):
            state_dict = safe_load_file(pt_path)
        else:
            state_dict = torch.load(pt_path, map_location="cpu")

        pt_state_dict.update(state_dict)

    logger.info(f"PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()):,} parameters")

    return load_pytorch_weights_in_tf2_model(
        tf_model,
        pt_state_dict,
        tf_inputs=tf_inputs,
        allow_missing_keys=allow_missing_keys,
        output_loading_info=output_loading_info,
        _prefix=_prefix,
        tf_to_pt_weight_rename=tf_to_pt_weight_rename,
    )


def load_pytorch_model_in_tf2_model(tf_model, pt_model, tf_inputs=None, allow_missing_keys=False):
    
    pt_state_dict = pt_model.state_dict()

    return load_pytorch_weights_in_tf2_model(
        tf_model, pt_state_dict, tf_inputs=tf_inputs, allow_missing_keys=allow_missing_keys
    )


def load_pytorch_weights_in_tf2_model(
    tf_model,
    pt_state_dict,
    tf_inputs=None,
    allow_missing_keys=False,
    output_loading_info=False,
    _prefix=None,
    tf_to_pt_weight_rename=None,
):
    
    try:
        import tensorflow as tf  
        import torch  
    except ImportError:
        logger.error(
            "Loading a PyTorch model in TensorFlow, requires both PyTorch and TensorFlow to be installed. Please see "
            "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions."
        )
        raise

    pt_state_dict = {k: v.numpy() for k, v in pt_state_dict.items()}
    return load_pytorch_state_dict_in_tf2_model(
        tf_model,
        pt_state_dict,
        tf_inputs=tf_inputs,
        allow_missing_keys=allow_missing_keys,
        output_loading_info=output_loading_info,
        _prefix=_prefix,
        tf_to_pt_weight_rename=tf_to_pt_weight_rename,
    )


def load_pytorch_state_dict_in_tf2_model(
    tf_model,
    pt_state_dict,
    tf_inputs=None,
    allow_missing_keys=False,
    output_loading_info=False,
    _prefix=None,
    tf_to_pt_weight_rename=None,
    ignore_mismatched_sizes=False,
):
    
    import tensorflow as tf
    from keras import backend as K

    if tf_inputs is None:
        tf_inputs = tf_model.dummy_inputs

    if _prefix is None:
        _prefix = ""
    if tf_inputs:
        with tf.name_scope(_prefix):
            tf_model(tf_inputs, training=False)  
    
    tf_keys_to_pt_keys = {}
    for key in pt_state_dict.keys():
        new_key = None
        if "gamma" in key:
            new_key = key.replace("gamma", "weight")
        if "beta" in key:
            new_key = key.replace("beta", "bias")
        if "running_var" in key:
            new_key = key.replace("running_var", "moving_variance")
        if "running_mean" in key:
            new_key = key.replace("running_mean", "moving_mean")

        
        key_components = key.split(".")
        name = None
        if key_components[-3::2] == ["parametrizations", "original0"]:
            name = key_components[-2] + "_g"
        elif key_components[-3::2] == ["parametrizations", "original1"]:
            name = key_components[-2] + "_v"
        if name is not None:
            key_components = key_components[:-3] + [name]
            new_key = ".".join(key_components)

        if new_key is None:
            new_key = key
        tf_keys_to_pt_keys[new_key] = key

    
    
    
    
    start_prefix_to_remove = ""
    if not any(s.startswith(tf_model.base_model_prefix) for s in tf_keys_to_pt_keys.keys()):
        start_prefix_to_remove = tf_model.base_model_prefix + "."

    symbolic_weights = tf_model.trainable_weights + tf_model.non_trainable_weights
    tf_loaded_numel = 0
    all_pytorch_weights = set(tf_keys_to_pt_keys.keys())
    missing_keys = []
    mismatched_keys = []
    is_safetensor_archive = hasattr(pt_state_dict, "get_tensor")
    for symbolic_weight in symbolic_weights:
        sw_name = symbolic_weight.name
        name, transpose = convert_tf_weight_name_to_pt_weight_name(
            sw_name,
            start_prefix_to_remove=start_prefix_to_remove,
            tf_weight_shape=symbolic_weight.shape,
            name_scope=_prefix,
        )
        if tf_to_pt_weight_rename is not None:
            name = tf_to_pt_weight_rename(name)

        
        if name not in tf_keys_to_pt_keys:
            if allow_missing_keys:
                missing_keys.append(name)
                continue
            elif tf_model._keys_to_ignore_on_load_missing is not None:
                
                if any(re.search(pat, name) is not None for pat in tf_model._keys_to_ignore_on_load_missing):
                    continue
            raise AttributeError(f"{name} not found in PyTorch model")
        state_dict_name = tf_keys_to_pt_keys[name]
        if is_safetensor_archive:
            array = pt_state_dict.get_tensor(state_dict_name)
        else:
            array = pt_state_dict[state_dict_name]
        try:
            array = apply_transpose(transpose, array, symbolic_weight.shape)
        except tf.errors.InvalidArgumentError as e:
            if not ignore_mismatched_sizes:
                error_msg = str(e)
                error_msg += (
                    "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method."
                )
                raise tf.errors.InvalidArgumentError(error_msg)
            else:
                mismatched_keys.append((name, array.shape, symbolic_weight.shape))
                continue

        tf_loaded_numel += tensor_size(array)

        K.set_value(symbolic_weight, array)
        del array  
        all_pytorch_weights.discard(name)

    logger.info(f"Loaded {tf_loaded_numel:,} parameters in the TF 2.0 model.")

    unexpected_keys = list(all_pytorch_weights)

    if tf_model._keys_to_ignore_on_load_missing is not None:
        for pat in tf_model._keys_to_ignore_on_load_missing:
            missing_keys = [k for k in missing_keys if re.search(pat, k) is None]
    if tf_model._keys_to_ignore_on_load_unexpected is not None:
        for pat in tf_model._keys_to_ignore_on_load_unexpected:
            unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]

    if len(unexpected_keys) > 0:
        logger.warning(
            "Some weights of the PyTorch model were not used when initializing the TF 2.0 model"
            f" {tf_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"
            f" {tf_model.__class__.__name__} from a PyTorch model trained on another task or with another architecture"
            " (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n- This IS"
            f" NOT expected if you are initializing {tf_model.__class__.__name__} from a PyTorch model that you expect"
            " to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a"
            " BertForSequenceClassification model)."
        )
    else:
        logger.warning(f"All PyTorch model weights were used when initializing {tf_model.__class__.__name__}.\n")
    if len(missing_keys) > 0:
        logger.warning(
            f"Some weights or buffers of the TF 2.0 model {tf_model.__class__.__name__} were not initialized from the"
            f" PyTorch model and are newly initialized: {missing_keys}\nYou should probably TRAIN this model on a"
            " down-stream task to be able to use it for predictions and inference."
        )
    else:
        logger.warning(
            f"All the weights of {tf_model.__class__.__name__} were initialized from the PyTorch model.\n"
            "If your task is similar to the task the model of the checkpoint was trained on, "
            f"you can already use {tf_model.__class__.__name__} for predictions without further training."
        )

    if len(mismatched_keys) > 0:
        mismatched_warning = "\n".join(
            [
                f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"
                for key, shape1, shape2 in mismatched_keys
            ]
        )
        logger.warning(
            f"Some weights of {tf_model.__class__.__name__} were not initialized from the model checkpoint"
            f" are newly initialized because the shapes did not"
            f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able"
            " to use it for predictions and inference."
        )

    if output_loading_info:
        loading_info = {
            "missing_keys": missing_keys,
            "unexpected_keys": unexpected_keys,
            "mismatched_keys": mismatched_keys,
        }
        return tf_model, loading_info

    return tf_model







def load_tf2_checkpoint_in_pytorch_model(
    pt_model, tf_checkpoint_path, tf_inputs=None, allow_missing_keys=False, output_loading_info=False
):
    
    try:
        import tensorflow as tf  
        import torch  
    except ImportError:
        logger.error(
            "Loading a TensorFlow model in PyTorch, requires both PyTorch and TensorFlow to be installed. Please see "
            "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions."
        )
        raise

    import transformers

    from .modeling_tf_utils import load_tf_weights

    logger.info(f"Loading TensorFlow weights from {tf_checkpoint_path}")

    
    tf_model_class_name = "TF" + pt_model.__class__.__name__  
    tf_model_class = getattr(transformers, tf_model_class_name)
    tf_model = tf_model_class(pt_model.config)

    if tf_inputs is None:
        tf_inputs = tf_model.dummy_inputs

    if tf_inputs is not None:
        tf_model(tf_inputs, training=False)  

    load_tf_weights(tf_model, tf_checkpoint_path)

    return load_tf2_model_in_pytorch_model(
        pt_model, tf_model, allow_missing_keys=allow_missing_keys, output_loading_info=output_loading_info
    )


def load_tf2_model_in_pytorch_model(pt_model, tf_model, allow_missing_keys=False, output_loading_info=False):
    
    weights = tf_model.weights

    return load_tf2_weights_in_pytorch_model(
        pt_model, weights, allow_missing_keys=allow_missing_keys, output_loading_info=output_loading_info
    )


def load_tf2_weights_in_pytorch_model(pt_model, tf_weights, allow_missing_keys=False, output_loading_info=False):
    
    try:
        import tensorflow as tf  
        import torch  
    except ImportError:
        logger.error(
            "Loading a TensorFlow model in PyTorch, requires both PyTorch and TensorFlow to be installed. Please see "
            "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions."
        )
        raise

    tf_state_dict = {tf_weight.name: tf_weight.numpy() for tf_weight in tf_weights}
    return load_tf2_state_dict_in_pytorch_model(
        pt_model, tf_state_dict, allow_missing_keys=allow_missing_keys, output_loading_info=output_loading_info
    )


def load_tf2_state_dict_in_pytorch_model(pt_model, tf_state_dict, allow_missing_keys=False, output_loading_info=False):
    import torch

    new_pt_params_dict = {}
    current_pt_params_dict = dict(pt_model.named_parameters())

    
    
    start_prefix_to_remove = ""
    if not any(s.startswith(pt_model.base_model_prefix) for s in current_pt_params_dict.keys()):
        start_prefix_to_remove = pt_model.base_model_prefix + "."

    
    tf_weights_map = {}
    for name, tf_weight in tf_state_dict.items():
        pt_name, transpose = convert_tf_weight_name_to_pt_weight_name(
            name, start_prefix_to_remove=start_prefix_to_remove, tf_weight_shape=tf_weight.shape
        )
        tf_weights_map[pt_name] = (tf_weight, transpose)

    all_tf_weights = set(tf_weights_map.keys())
    loaded_pt_weights_data_ptr = {}
    missing_keys_pt = []
    for pt_weight_name, pt_weight in current_pt_params_dict.items():
        
        if pt_weight.data_ptr() in loaded_pt_weights_data_ptr:
            new_pt_params_dict[pt_weight_name] = loaded_pt_weights_data_ptr[pt_weight.data_ptr()]
            continue

        pt_weight_name_to_check = pt_weight_name
        
        key_components = pt_weight_name.split(".")
        name = None
        if key_components[-3::2] == ["parametrizations", "original0"]:
            name = key_components[-2] + "_g"
        elif key_components[-3::2] == ["parametrizations", "original1"]:
            name = key_components[-2] + "_v"
        if name is not None:
            key_components = key_components[:-3] + [name]
            pt_weight_name_to_check = ".".join(key_components)

        
        if pt_weight_name_to_check not in tf_weights_map:
            if allow_missing_keys:
                missing_keys_pt.append(pt_weight_name)
                continue

            raise AttributeError(f"{pt_weight_name} not found in TF 2.0 model")

        array, transpose = tf_weights_map[pt_weight_name_to_check]

        array = apply_transpose(transpose, array, pt_weight.shape, pt_to_tf=False)

        if numpy.isscalar(array):
            array = numpy.array(array)
        if not is_torch_tensor(array) and not is_numpy_array(array):
            array = array.numpy()
        if is_numpy_array(array):
            
            array = torch.from_numpy(array)

        new_pt_params_dict[pt_weight_name] = array
        loaded_pt_weights_data_ptr[pt_weight.data_ptr()] = array
        all_tf_weights.discard(pt_weight_name)

    missing_keys, unexpected_keys = pt_model.load_state_dict(new_pt_params_dict, strict=False)
    missing_keys += missing_keys_pt

    
    
    if pt_model._keys_to_ignore_on_load_missing is not None:
        for pat in pt_model._keys_to_ignore_on_load_missing:
            missing_keys = [k for k in missing_keys if re.search(pat, k) is None]

    if pt_model._keys_to_ignore_on_load_unexpected is not None:
        for pat in pt_model._keys_to_ignore_on_load_unexpected:
            unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]

    if len(unexpected_keys) > 0:
        logger.warning(
            "Some weights of the TF 2.0 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 TF 2.0 model trained on another task or with another architecture"
            " (e.g. initializing a BertForSequenceClassification model from a TFBertForPreTraining model).\n- This IS"
            f" NOT expected if you are initializing {pt_model.__class__.__name__} from a TF 2.0 model that you expect"
            " to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"
            " TFBertForSequenceClassification model)."
        )
    else:
        logger.warning(f"All TF 2.0 model weights were used when initializing {pt_model.__class__.__name__}.\n")
    if len(missing_keys) > 0:
        logger.warning(
            f"Some weights of {pt_model.__class__.__name__} were not initialized from the TF 2.0 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."
        )
    else:
        logger.warning(
            f"All the weights of {pt_model.__class__.__name__} were initialized from the TF 2.0 model.\n"
            "If your task is similar to the task the model of the checkpoint was trained on, "
            f"you can already use {pt_model.__class__.__name__} for predictions without further training."
        )

    logger.info(f"Weights or buffers not loaded from TF 2.0 model: {all_tf_weights}")

    if output_loading_info:
        loading_info = {"missing_keys": missing_keys, "unexpected_keys": unexpected_keys}
        return pt_model, loading_info

    return pt_model