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 wav2vec2 checkpoint import argparse import os from functools import reduce import fairseq import torch from datasets import loaddataset from transformers import wav2vec2processor logging from transformers models data2vec configurationdata2vecaudio import data2vecaudioconfig copied from https github compytorchfairseqblobmainexamplesdata2vecmodelsdata2vecaudio py from transformers models data2vec data2vecaudio import data2vecaudiomodel as dummy noqa f401 from transformers models data2vec modelingdata2vecaudio import data2vecaudioforctc data2vecaudiomodel logging setverbosityinfo logger logging getloggername mapping postextractproj featureprojection projection models 0 layernorm featureprojection layernorm selfattn kproj encoder layers attention kproj selfattn vproj encoder layers attention vproj selfattn qproj encoder layers attention qproj selfattn outproj encoder layers attention outproj 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 w2vencoder proj lmhead maskemb maskedspecembed toplevelkeys lmhead 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 if hfshape value shape raise valueerror 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 isheadless unusedweights fairseqdict fairseqmodel statedict if not isheadless featureextractor hfmodel data2vecaudio featureextractor posconvembedding hfmodel data2vecaudio encoder posconvembed else featureextractor hfmodel featureextractor posconvembedding hfmodel encoder posconvembed for name value in fairseqdict items isused false if convlayers in name loadconvlayer name value featureextractor unusedweights isused true elif posconv in name loadposconvlayer name value posconvembedding unusedweights isused true else for key mappedkey in mapping items if not isheadless mappedkey data2vecaudio mappedkey if mappedkey not in toplevelkeys else mappedkey 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 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 accessbystringmodule path names path split return reducegetattr names module def setweightsfullname module fsqvalue hfweightpath hfweight accessbystringmodule hfweightpath hfvalue hfweight data if fsqvalue shape hfvalue shape raise valueerrorffullname has size fsqvalue shape but hfvalue shape was found hfweight data fsqvalue logger infoffullname was correctly initialized from hfweightpath def loadconvlayerfullname value featureextractor unusedweights name fullname splitconvlayers 1 items name split layerid intitems0 typeid intitems1 weighttype name split 1 if typeid 0 layertype conv elif typeid 2 layertype layernorm else unusedweights appendfullname return setweightsfullname featureextractor value fconvlayers layerid layertype weighttype def loadposconvlayerfullname value posconvembeddings unusedweights name fullname splitposconv 1 items name split layerid intitems0 typeid intitems1 weighttype name split 1 if typeid 0 unusedweights appendfullname return else layertype conv setweightsfullname posconvembeddings value flayers layerid layertype weighttype torch nograd def convertwav2vec2checkpoint checkpointpath pytorchdumpfolderpath configpathnone dictpathnone isfinetunedtrue if configpath is not none config data2vecaudioconfig frompretrainedconfigpath else config data2vecaudioconfig if not isfinetuned modify finalproj layer name hfwav2vec data2vecaudiomodelconfig data2veccheckpointdir os path dirnamecheckpointpath statedict torch loadcheckpointpath statedictmodelfinalproj weight statedictmodel popfinalproj 0 weight statedictmodelfinalproj bias statedictmodel popfinalproj 0 bias convertedckpt os path joindata2veccheckpointdir converted pt torch savestatedict convertedckpt else hfwav2vec data2vecaudioforctcconfig convertedckpt checkpointpath def loaddata2vecpath model fairseq checkpointutils loadmodelensembleandtaskpath return model0 eval model loaddata2vecconvertedckpt recursivelyloadweightsmodel hfwav2vec not isfinetuned processor wav2vec2processor frompretrainedfacebookwav2vec2largelv60 ds loaddatasetpatrickvonplatenlibrispeechasrdummy clean splitvalidation inputaudio xarray for x in ds 4audio inputs processorinputaudio returntensorspt paddingtrue inputvalues inputs inputvalues attentionmask inputs attentionmask inputvalues inputs inputvalues 1 attentionmask inputs attentionmask 1 hfwav2vec eval model eval if isfinetuned theiroutput modelsourceinputvalues paddingmask1 attentionmask maskfalse featuresonlytrue encoderout transpose0 1 ouroutput hfwav2vecinputvalues attentionmaskattentionmasklogits predids torch argmaxouroutput dim1 outputstring processor batchdecodepredids printfexpected output ds 4 text pred outputstring else theiroutput modelsourceinputvalues paddingmask1 attentionmask maskfalse featuresonlytrue layerresults 10 transpose0 1 ouroutput hfwav2vecinputvalues attentionmaskattentionmasklasthiddenstate printouroutput shape theiroutput shape maxabsolutediff torch maxtorch absouroutput theiroutput item printfmaxabsolutediff maxabsolutediff 1e7 success torch allcloseouroutput theiroutput atol1e3 printdo both models output the same tensors if success else if not success raise exceptionsomething went wrong hfwav2vec savepretrainedpytorchdumpfolderpath if isfinetuned processor savepretrainedpytorchdumpfolderpath else processor featureextractor 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 addargumentdictpath defaultnone typestr helppath to dict of finetuned model parser addargumentconfigpath defaultnone typestr helppath to hf config json of model to convert parser addargument notfinetuned actionstoretrue helpwhether the model to convert is a finetuned model or not args parser parseargs convertwav2vec2checkpoint args checkpointpath args pytorchdumpfolderpath args configpath args dictpath not args notfinetuned 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 wav2vec2 checkpoint copied from https github com pytorch fairseq blob main examples data2vec models data2vec_audio py noqa f401 todo don t match quantizer weight_proj copy paste tweak model s weights to transformers design modify final_proj layer name input_values inputs input_values 1 attention_mask inputs attention_mask 1 1e 7
import argparse
import os
from functools import reduce

import fairseq
import torch
from datasets import load_dataset

from transformers import Wav2Vec2Processor, logging
from transformers.models.data2vec.configuration_data2vec_audio import Data2VecAudioConfig


from transformers.models.data2vec.data2vec_audio import Data2VecAudioModel as Dummy  
from transformers.models.data2vec.modeling_data2vec_audio import Data2VecAudioForCTC, Data2VecAudioModel


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

MAPPING = {
    "post_extract_proj": "feature_projection.projection",
    "models.0.layer_norm": "feature_projection.layer_norm",
    "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_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",
    "w2v_encoder.proj": "lm_head",
    "mask_emb": "masked_spec_embed",
}
TOP_LEVEL_KEYS = [
    "lm_head",
]


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

    if hf_shape != value.shape:
        raise ValueError(
            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, is_headless):
    unused_weights = []
    fairseq_dict = fairseq_model.state_dict()

    if not is_headless:
        feature_extractor = hf_model.data2vec_audio.feature_extractor
        pos_conv_embedding = hf_model.data2vec_audio.encoder.pos_conv_embed

    else:
        feature_extractor = hf_model.feature_extractor
        pos_conv_embedding = hf_model.encoder.pos_conv_embed

    for name, value in fairseq_dict.items():
        is_used = False
        if "conv_layers" in name:
            load_conv_layer(
                name,
                value,
                feature_extractor,
                unused_weights,
            )
            is_used = True
        elif "pos_conv" in name:
            load_pos_conv_layer(
                name,
                value,
                pos_conv_embedding,
                unused_weights,
            )
            is_used = True
        else:
            for key, mapped_key in MAPPING.items():
                if not is_headless:
                    mapped_key = "data2vec_audio." + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
                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:
                        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 access_by_string(module, path):
    names = path.split(".")
    return reduce(getattr, names, module)


def set_weights(full_name, module, fsq_value, hf_weight_path):
    hf_weight = access_by_string(module, hf_weight_path)
    hf_value = hf_weight.data

    if fsq_value.shape != hf_value.shape:
        raise ValueError(f"{full_name} has size {fsq_value.shape}, but {hf_value.shape} was found.")
    hf_weight.data = fsq_value
    logger.info(f"{full_name} was correctly initialized from {hf_weight_path}.")


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

    weight_type = name.split(".")[-1]
    if type_id == 0:
        layer_type = "conv"
    elif type_id == 2:
        layer_type = "layer_norm"
    else:
        unused_weights.append(full_name)
        return

    set_weights(full_name, feature_extractor, value, f"conv_layers.{layer_id}.{layer_type}.{weight_type}")


def load_pos_conv_layer(full_name, value, pos_conv_embeddings, unused_weights):
    name = full_name.split("pos_conv.")[-1]
    items = name.split(".")
    layer_id = int(items[0])
    type_id = int(items[1])

    weight_type = name.split(".")[-1]
    if type_id != 0:
        unused_weights.append(full_name)
        return
    else:
        layer_type = "conv"

    set_weights(full_name, pos_conv_embeddings, value, f"layers.{layer_id}.{layer_type}.{weight_type}")


@torch.no_grad()
def convert_wav2vec2_checkpoint(
    checkpoint_path, pytorch_dump_folder_path, config_path=None, dict_path=None, is_finetuned=True
):
    
    if config_path is not None:
        config = Data2VecAudioConfig.from_pretrained(config_path)
    else:
        config = Data2VecAudioConfig()

    if not is_finetuned:
        
        hf_wav2vec = Data2VecAudioModel(config)
        data2vec_checkpoint_dir = os.path.dirname(checkpoint_path)

        state_dict = torch.load(checkpoint_path)
        state_dict["model"]["final_proj.weight"] = state_dict["model"].pop("final_proj.0.weight")
        state_dict["model"]["final_proj.bias"] = state_dict["model"].pop("final_proj.0.bias")
        converted_ckpt = os.path.join(data2vec_checkpoint_dir, "converted.pt")
        torch.save(state_dict, converted_ckpt)
    else:
        hf_wav2vec = Data2VecAudioForCTC(config)
        converted_ckpt = checkpoint_path

    def load_data2vec(path):
        model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task([path])
        return model[0].eval()

    model = load_data2vec(converted_ckpt)

    recursively_load_weights(model, hf_wav2vec, not is_finetuned)

    processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-lv60")

    ds = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
    input_audio = [x["array"] for x in ds[:4]["audio"]]

    inputs = processor(input_audio, return_tensors="pt", padding=True)

    input_values = inputs.input_values
    attention_mask = inputs.attention_mask
    
    

    hf_wav2vec.eval()
    model.eval()
    if is_finetuned:
        their_output = model(source=input_values, padding_mask=(1 - attention_mask), mask=False, features_only=True)[
            "encoder_out"
        ].transpose(0, 1)
        our_output = hf_wav2vec(input_values, attention_mask=attention_mask)["logits"]

        pred_ids = torch.argmax(our_output, dim=-1)
        output_string = processor.batch_decode(pred_ids)

        print(f"Expected Output: {ds[:4]['text']}, Pred: {output_string}")
    else:
        their_output = model(source=input_values, padding_mask=(1 - attention_mask), mask=False, features_only=True)[
            "layer_results"
        ][-1][0].transpose(0, 1)
        our_output = hf_wav2vec(input_values, attention_mask=attention_mask)["last_hidden_state"]

    print(our_output.shape, their_output.shape)
    max_absolute_diff = torch.max(torch.abs(our_output - their_output)).item()
    print(f"max_absolute_diff = {max_absolute_diff}")  
    success = torch.allclose(our_output, their_output, atol=1e-3)
    print("Do both models output the same tensors?", "🔥" if success else "💩")
    if not success:
        raise Exception("Something went wRoNg")

    hf_wav2vec.save_pretrained(pytorch_dump_folder_path)

    if is_finetuned:
        processor.save_pretrained(pytorch_dump_folder_path)
    else:
        processor.feature_extractor.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("--dict_path", default=None, type=str, help="Path to dict of fine-tuned model")
    parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
    parser.add_argument(
        "--not_finetuned", action="store_true", help="Whether the model to convert is a fine-tuned model or not"
    )
    args = parser.parse_args()
    convert_wav2vec2_checkpoint(
        args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
    )