codingutf8 2021 the huggingface inc team 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 convert wavlm checkpoint import argparse import torch step 1 clone https github commicrosoftunilm step 2 git checkout to https github commicrosoftunilmcommitb94ec76c36f02fb2b0bf0dcb0b8554a2185173cd step 3 cd unilm step 4 ln s realpath wavlmmodules py create simlink import classes from unilm wavlm wavlm import wavlm as wavlmorig from unilm wavlm wavlm import wavlmconfig as wavlmconfigorig from transformers import wavlmconfig wavlmmodel logging logging setverbosityinfo logger logging getloggername mapping postextractproj featureprojection projection encoder posconv 0 encoder posconvembed conv selfattn kproj encoder layers attention kproj selfattn vproj encoder layers attention vproj selfattn qproj encoder layers attention qproj selfattn outproj encoder layers attention outproj selfattn greplinear encoder layers attention grurelposlinear selfattn relativeattentionbias encoder layers attention relattnembed selfattn grepa encoder layers attention grurelposconst selfattnlayernorm encoder layers layernorm fc1 encoder layers feedforward intermediatedense fc2 encoder layers feedforward outputdense finallayernorm encoder layers finallayernorm encoder layernorm encoder layernorm w2vmodel layernorm featureprojection layernorm quantizer weightproj quantizer weightproj quantizer vars quantizer codevectors projectq projectq finalproj projecthid w2vencoder proj ctcproj maskemb maskedspecembed toplevelkeys ctcproj quantizer weightproj quantizer codevectors projectq projecthid def setrecursivelyhfpointer key value fullname weighttype for attribute in key split hfpointer getattrhfpointer attribute if weighttype is not none hfshape getattrhfpointer weighttype shape else hfshape hfpointer shape assert hfshape value shape fshape of hf key weighttype if weighttype is not none else is hfshape but should be f value shape for fullname if weighttype weight hfpointer weight data value elif weighttype weightg hfpointer weightg data value elif weighttype weightv hfpointer weightv data value elif weighttype bias hfpointer bias data value else hfpointer data value logger infofkey weighttype if weighttype is not none else was initialized from fullname def recursivelyloadweightsfairseqmodel hfmodel unusedweights fairseqdict fairseqmodel statedict featureextractor hfmodel featureextractor for name value in fairseqdict items isused false if convlayers in name loadconvlayer name value featureextractor unusedweights hfmodel config featextractnorm group isused true else for key mappedkey in mapping items if key in name or key splitw2vmodel 1 name split 0 isused true if in mappedkey layerindex name splitkey0 split 2 mappedkey mappedkey replace layerindex if weightg in name weighttype weightg elif weightv in name weighttype weightv elif bias in name and relativeattentionbias not in name weighttype bias elif weight in name todo don t match quantizer weightproj weighttype weight else weighttype none setrecursivelyhfmodel mappedkey value name weighttype continue if not isused unusedweights appendname logger warningfunused weights unusedweights def loadconvlayerfullname value featureextractor unusedweights usegroupnorm name fullname splitconvlayers 1 items name split layerid intitems0 typeid intitems1 if typeid 0 if bias in name assert value shape featureextractor convlayerslayerid conv bias data shape ffullname has size value shape but f featureextractor convlayerslayerid conv bias data shape was found featureextractor convlayerslayerid conv bias data value logger infoffeat extract conv layer layerid was initialized from fullname elif weight in name assert value shape featureextractor convlayerslayerid conv weight data shape ffullname has size value shape but f featureextractor convlayerslayerid conv weight data shape was found featureextractor convlayerslayerid conv weight data value logger infoffeat extract conv layer layerid was initialized from fullname elif typeid 2 and not usegroupnorm or typeid 2 and layerid 0 and usegroupnorm if bias in name assert value shape featureextractor convlayerslayerid layernorm bias data shape ffullname has size value shape but featureextractorlayerid layernorm bias data shape was found featureextractor convlayerslayerid layernorm bias data value logger infoffeat extract layer norm weight of layer layerid was initialized from fullname elif weight in name assert value shape featureextractor convlayerslayerid layernorm weight data shape ffullname has size value shape but f featureextractorlayerid layernorm weight data shape was found featureextractor convlayerslayerid layernorm weight data value logger infoffeat extract layer norm weight of layer layerid was initialized from fullname else unusedweights appendfullname torch nograd def convertwavlmcheckpointcheckpointpath pytorchdumpfolderpath configpathnone load the pretrained checkpoints checkpoint torch loadcheckpointpath cfg wavlmconfigorigcheckpointcfg model wavlmorigcfg model loadstatedictcheckpointmodel model eval if configpath is not none config wavlmconfig frompretrainedconfigpath else config wavlmconfig hfwavlm wavlmmodelconfig recursivelyloadweightsmodel hfwavlm hfwavlm savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser parser addargumentpytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model parser addargumentcheckpointpath defaultnone typestr helppath to fairseq checkpoint parser addargumentconfigpath defaultnone typestr helppath to hf config json of model to convert args parser parseargs convertwavlmcheckpointargs checkpointpath args pytorchdumpfolderpath args configpath coding utf 8 2021 the huggingface inc team 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 convert wavlm checkpoint step 1 clone https github com microsoft unilm step 2 git checkout to https github com microsoft unilm commit b94ec76c36f02fb2b0bf0dcb0b8554a2185173cd step 3 cd unilm step 4 ln s realpath wavlm modules py create simlink import classes todo don t match quantizer weight_proj load the pre trained checkpoints
import argparse

import torch






from unilm.wavlm.WavLM import WavLM as WavLMOrig
from unilm.wavlm.WavLM import WavLMConfig as WavLMConfigOrig

from transformers import WavLMConfig, WavLMModel, logging


logging.set_verbosity_info()
logger = logging.get_logger(__name__)

MAPPING = {
    "post_extract_proj": "feature_projection.projection",
    "encoder.pos_conv.0": "encoder.pos_conv_embed.conv",
    "self_attn.k_proj": "encoder.layers.*.attention.k_proj",
    "self_attn.v_proj": "encoder.layers.*.attention.v_proj",
    "self_attn.q_proj": "encoder.layers.*.attention.q_proj",
    "self_attn.out_proj": "encoder.layers.*.attention.out_proj",
    "self_attn.grep_linear": "encoder.layers.*.attention.gru_rel_pos_linear",
    "self_attn.relative_attention_bias": "encoder.layers.*.attention.rel_attn_embed",
    "self_attn.grep_a": "encoder.layers.*.attention.gru_rel_pos_const",
    "self_attn_layer_norm": "encoder.layers.*.layer_norm",
    "fc1": "encoder.layers.*.feed_forward.intermediate_dense",
    "fc2": "encoder.layers.*.feed_forward.output_dense",
    "final_layer_norm": "encoder.layers.*.final_layer_norm",
    "encoder.layer_norm": "encoder.layer_norm",
    "w2v_model.layer_norm": "feature_projection.layer_norm",
    "quantizer.weight_proj": "quantizer.weight_proj",
    "quantizer.vars": "quantizer.codevectors",
    "project_q": "project_q",
    "final_proj": "project_hid",
    "w2v_encoder.proj": "ctc_proj",
    "mask_emb": "masked_spec_embed",
}
TOP_LEVEL_KEYS = [
    "ctc_proj",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]


def set_recursively(hf_pointer, key, value, full_name, weight_type):
    for attribute in key.split("."):
        hf_pointer = getattr(hf_pointer, attribute)

    if weight_type is not None:
        hf_shape = getattr(hf_pointer, weight_type).shape
    else:
        hf_shape = hf_pointer.shape

    assert hf_shape == value.shape, (
        f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be"
        f" {value.shape} for {full_name}"
    )

    if weight_type == "weight":
        hf_pointer.weight.data = value
    elif weight_type == "weight_g":
        hf_pointer.weight_g.data = value
    elif weight_type == "weight_v":
        hf_pointer.weight_v.data = value
    elif weight_type == "bias":
        hf_pointer.bias.data = value
    else:
        hf_pointer.data = value

    logger.info(f"{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}.")


def recursively_load_weights(fairseq_model, hf_model):
    unused_weights = []
    fairseq_dict = fairseq_model.state_dict()

    feature_extractor = hf_model.feature_extractor

    for name, value in fairseq_dict.items():
        is_used = False
        if "conv_layers" in name:
            load_conv_layer(
                name,
                value,
                feature_extractor,
                unused_weights,
                hf_model.config.feat_extract_norm == "group",
            )
            is_used = True
        else:
            for key, mapped_key in MAPPING.items():
                if key in name or key.split("w2v_model.")[-1] == name.split(".")[0]:
                    is_used = True
                    if "*" in mapped_key:
                        layer_index = name.split(key)[0].split(".")[-2]
                        mapped_key = mapped_key.replace("*", layer_index)
                    if "weight_g" in name:
                        weight_type = "weight_g"
                    elif "weight_v" in name:
                        weight_type = "weight_v"
                    elif "bias" in name and "relative_attention_bias" not in name:
                        weight_type = "bias"
                    elif "weight" in name:
                        
                        weight_type = "weight"
                    else:
                        weight_type = None

                    set_recursively(hf_model, mapped_key, value, name, weight_type)
                continue
        if not is_used:
            unused_weights.append(name)

    logger.warning(f"Unused weights: {unused_weights}")


def load_conv_layer(full_name, value, feature_extractor, unused_weights, use_group_norm):
    name = full_name.split("conv_layers.")[-1]
    items = name.split(".")
    layer_id = int(items[0])
    type_id = int(items[1])

    if type_id == 0:
        if "bias" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
                f"{full_name} has size {value.shape}, but"
                f" {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found."
            )
            feature_extractor.conv_layers[layer_id].conv.bias.data = value
            logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}.")
        elif "weight" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
                f"{full_name} has size {value.shape}, but"
                f" {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found."
            )
            feature_extractor.conv_layers[layer_id].conv.weight.data = value
            logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}.")
    elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
        if "bias" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
                f"{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was"
                " found."
            )
            feature_extractor.conv_layers[layer_id].layer_norm.bias.data = value
            logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.")
        elif "weight" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
                f"{full_name} has size {value.shape}, but"
                f" {feature_extractor[layer_id].layer_norm.weight.data.shape} was found."
            )
            feature_extractor.conv_layers[layer_id].layer_norm.weight.data = value
            logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.")
    else:
        unused_weights.append(full_name)


@torch.no_grad()
def convert_wavlm_checkpoint(checkpoint_path, pytorch_dump_folder_path, config_path=None):
    
    checkpoint = torch.load(checkpoint_path)
    cfg = WavLMConfigOrig(checkpoint["cfg"])
    model = WavLMOrig(cfg)
    model.load_state_dict(checkpoint["model"])
    model.eval()

    if config_path is not None:
        config = WavLMConfig.from_pretrained(config_path)
    else:
        config = WavLMConfig()

    hf_wavlm = WavLMModel(config)

    recursively_load_weights(model, hf_wavlm)

    hf_wavlm.save_pretrained(pytorch_dump_folder_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
    parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to fairseq checkpoint")
    parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
    args = parser.parse_args()
    convert_wavlm_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)