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 json import os import fairseq import torch from fairseq data import dictionary from transformers import wav2vec2config wav2vec2ctctokenizer wav2vec2featureextractor wav2vec2forctc wav2vec2forpretraining wav2vec2processor logging from transformers models wav2vec2 modelingwav2vec2 import wav2vec2forsequenceclassification 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 selfattnlayernorm encoder layers layernorm fc1 encoder layers feedforward intermediatedense fc2 encoder layers feedforward outputdense finallayernorm encoder layers finallayernorm encoder layernorm encoder layernorm adapterlayer encoder layers adapterlayer w2vmodel layernorm featureprojection layernorm quantizer weightproj quantizer weightproj quantizer vars quantizer codevectors projectq projectq finalproj projecthid w2vencoder proj lmhead maskemb maskedspecembed poolinglayer linear projector poolinglayer projection classifier toplevelkeys lmhead quantizer weightproj quantizer codevectors projectq projecthid projector classifier def readtxtintodictfilename result with openfilename r as file for linenumber line in enumeratefile line line strip if line words line split key linenumber value words0 resultkey value return result def setrecursivelykey value fullname weighttype hfpointer for attribute in key split hfpointer getattrhfpointer attribute hfparamname none for paramkey in parammapping keys if fullname endswithparamkey hfparamname parammappingfullname split 1 weighttype param if weighttype is not none and weighttype param hfshape getattrhfpointer weighttype shape elif weighttype is not none and weighttype param shapepointer hfpointer for attribute in hfparamname split shapepointer getattrshapepointer attribute hfshape shapepointer shape let s reduce dimension value value0 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 elif weighttype param for attribute in hfparamname split hfpointer getattrhfpointer attribute hfpointer data value else hfpointer data value logger infofkey weighttype if weighttype is not none else was initialized from fullname def renamedictkey value fullname weighttype hfdict hfparamname none for paramkey in parammapping keys if fullname endswithparamkey hfparamname parammappingfullname split 1 weighttype param if weighttype is not none and weighttype param fullkey joinkey weighttype elif weighttype is not none and weighttype param fullkey joinkey hfparamname else fullkey key hfdictfullkey value if lmhead in fullkey else value0 parammapping wa linear1 weight wb linear2 weight ba linear1 bias bb linear2 bias lnw norm weight lnb norm bias def loadwav2vec2layername value hfmodelnone hfdictnone isused false for key mappedkey in mapping items mappedkey wav2vec2 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 if hfdict is not none renamedictmappedkey value name weighttype hfdict else setrecursivelymappedkey value name weighttype hfmodel return isused return isused def recursivelyloadweightsfairseqmodel hfmodel isheadless unusedweights fairseqdict fairseqmodel statedict featureextractor hfmodel wav2vec2 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 isused loadwav2vec2layername value hfmodel 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 if value shape featureextractor convlayerslayerid conv bias data shape raise valueerror 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 if value shape featureextractor convlayerslayerid conv weight data shape raise valueerror 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 if value shape featureextractor convlayerslayerid layernorm bias data shape raise valueerror ffullname has size value shape but f featureextractor convlayerslayerid 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 if value shape featureextractor convlayerslayerid layernorm weight data shape raise valueerror ffullname has size value shape but f featureextractor convlayerslayerid 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 convertwav2vec2checkpoint checkpointpath pytorchdumpfolderpath configpathnone dictpathnone isfinetunedtrue isseqclassfalse if configpath is not none config wav2vec2config frompretrainedconfigpath else config wav2vec2config if isseqclass id2label readtxtintodictdictpath config id2label id2label hfwav2vec wav2vec2forsequenceclassificationconfig featureextractor wav2vec2featureextractor featuresize1 samplingrate16000 paddingvalue0 donormalizetrue returnattentionmasktrue featureextractor savepretrainedpytorchdumpfolderpath elif isfinetuned if dictpath targetdict dictionary loaddictpath important change bos pad token id since ctc symbol is pad and not s as in fairseq config bostokenid targetdict padindex config padtokenid targetdict bosindex config eostokenid targetdict eosindex config vocabsize lentargetdict symbols vocabpath os path joinpytorchdumpfolderpath vocab json if not os path isdirpytorchdumpfolderpath logger errorpytorchdumpfolderpath should be a directory formatpytorchdumpfolderpath return os makedirspytorchdumpfolderpath existoktrue vocabdict targetdict indices fairseq has the pad and s switched vocabdictpad 0 vocabdicts 1 with openvocabpath w encodingutf8 as vocabhandle json dumpvocabdict vocabhandle tokenizer wav2vec2ctctokenizer vocabpath unktokentargetdict unkword padtokentargetdict padword bostokentargetdict bosword eostokentargetdict eosword worddelimitertoken dolowercasefalse returnattentionmask true if config featextractnorm layer else false featureextractor wav2vec2featureextractor featuresize1 samplingrate16000 paddingvalue0 donormalizetrue returnattentionmaskreturnattentionmask processor wav2vec2processorfeatureextractorfeatureextractor tokenizertokenizer processor savepretrainedpytorchdumpfolderpath hfwav2vec wav2vec2forctcconfig else hfwav2vec wav2vec2forpretrainingconfig if isfinetuned or isseqclass model fairseq checkpointutils loadmodelensembleandtask checkpointpath argoverridesdata joindictpath split 1 else taskarg argparse namespacetaskaudiopretraining task fairseq tasks setuptasktaskarg model fairseq checkpointutils loadmodelensembleandtaskcheckpointpath tasktask model model0 eval recursivelyloadweightsmodel hfwav2vec not isfinetuned hfwav2vec 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 parser addargument isseqclass actionstoretrue helpwhether the model to convert is a finetuned sequence classification model or not args parser parseargs isfinetuned not args notfinetuned and not args isseqclass convertwav2vec2checkpoint args checkpointpath args pytorchdumpfolderpath args configpath args dictpath isfinetuned args isseqclass 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 let s reduce dimension todo don t match quantizer weight_proj copy paste tweak model s weights to transformers design important change bos pad token id since ctc symbol is pad and not s as in fairseq fairseq has the pad and s switched
import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    Wav2Vec2Config,
    Wav2Vec2CTCTokenizer,
    Wav2Vec2FeatureExtractor,
    Wav2Vec2ForCTC,
    Wav2Vec2ForPreTraining,
    Wav2Vec2Processor,
    logging,
)
from transformers.models.wav2vec2.modeling_wav2vec2 import Wav2Vec2ForSequenceClassification


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_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",
    "adapter_layer": "encoder.layers.*.adapter_layer",
    "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": "lm_head",
    "mask_emb": "masked_spec_embed",
    "pooling_layer.linear": "projector",
    "pooling_layer.projection": "classifier",
}
TOP_LEVEL_KEYS = [
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
    "projector",
    "classifier",
]


def read_txt_into_dict(filename):
    result = {}
    with open(filename, "r") as file:
        for line_number, line in enumerate(file):
            line = line.strip()
            if line:
                words = line.split()
                key = line_number
                value = words[0]
                result[key] = value
    return result


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

    hf_param_name = None
    for param_key in PARAM_MAPPING.keys():
        if full_name.endswith(param_key):
            hf_param_name = PARAM_MAPPING[full_name.split(".")[-1]]
            weight_type = "param"

    if weight_type is not None and weight_type != "param":
        hf_shape = getattr(hf_pointer, weight_type).shape
    elif weight_type is not None and weight_type == "param":
        shape_pointer = hf_pointer
        for attribute in hf_param_name.split("."):
            shape_pointer = getattr(shape_pointer, attribute)
        hf_shape = shape_pointer.shape

        
        value = value[0]
    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
    elif weight_type == "param":
        for attribute in hf_param_name.split("."):
            hf_pointer = getattr(hf_pointer, attribute)
        hf_pointer.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 rename_dict(key, value, full_name, weight_type, hf_dict):
    hf_param_name = None
    for param_key in PARAM_MAPPING.keys():
        if full_name.endswith(param_key):
            hf_param_name = PARAM_MAPPING[full_name.split(".")[-1]]
            weight_type = "param"

    if weight_type is not None and weight_type != "param":
        full_key = ".".join([key, weight_type])
    elif weight_type is not None and weight_type == "param":
        full_key = ".".join([key, hf_param_name])
    else:
        full_key = key

    hf_dict[full_key] = value if "lm_head" in full_key else value[0]


PARAM_MAPPING = {
    "W_a": "linear_1.weight",
    "W_b": "linear_2.weight",
    "b_a": "linear_1.bias",
    "b_b": "linear_2.bias",
    "ln_W": "norm.weight",
    "ln_b": "norm.bias",
}


def load_wav2vec2_layer(name, value, hf_model=None, hf_dict=None):
    is_used = False
    for key, mapped_key in MAPPING.items():
        mapped_key = "wav2vec2." + 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
            if hf_dict is not None:
                rename_dict(mapped_key, value, name, weight_type, hf_dict)
            else:
                set_recursively(mapped_key, value, name, weight_type, hf_model)
            return is_used
    return is_used


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

    feature_extractor = hf_model.wav2vec2.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:
            is_used = load_wav2vec2_layer(name, value, hf_model)
        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:
            if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
                raise ValueError(
                    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:
            if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
                raise ValueError(
                    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:
            if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
                raise ValueError(
                    f"{full_name} has size {value.shape}, but"
                    f" {feature_extractor.conv_layers[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:
            if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
                raise ValueError(
                    f"{full_name} has size {value.shape}, but"
                    f" {feature_extractor.conv_layers[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_wav2vec2_checkpoint(
    checkpoint_path, pytorch_dump_folder_path, config_path=None, dict_path=None, is_finetuned=True, is_seq_class=False
):
    
    if config_path is not None:
        config = Wav2Vec2Config.from_pretrained(config_path)
    else:
        config = Wav2Vec2Config()

    if is_seq_class:
        id2label = read_txt_into_dict(dict_path)
        config.id2label = id2label
        hf_wav2vec = Wav2Vec2ForSequenceClassification(config)
        feature_extractor = Wav2Vec2FeatureExtractor(
            feature_size=1,
            sampling_rate=16000,
            padding_value=0,
            do_normalize=True,
            return_attention_mask=True,
        )
        feature_extractor.save_pretrained(pytorch_dump_folder_path)

    elif is_finetuned:
        if dict_path:
            target_dict = Dictionary.load(dict_path)

            
            
            config.bos_token_id = target_dict.pad_index
            config.pad_token_id = target_dict.bos_index
            config.eos_token_id = target_dict.eos_index
            config.vocab_size = len(target_dict.symbols)
            vocab_path = os.path.join(pytorch_dump_folder_path, "vocab.json")
            if not os.path.isdir(pytorch_dump_folder_path):
                logger.error("--pytorch_dump_folder_path ({}) should be a directory".format(pytorch_dump_folder_path))
                return
            os.makedirs(pytorch_dump_folder_path, exist_ok=True)
            vocab_dict = target_dict.indices

            
            vocab_dict["<pad>"] = 0
            vocab_dict["<s>"] = 1
            with open(vocab_path, "w", encoding="utf-8") as vocab_handle:
                json.dump(vocab_dict, vocab_handle)
            tokenizer = Wav2Vec2CTCTokenizer(
                vocab_path,
                unk_token=target_dict.unk_word,
                pad_token=target_dict.pad_word,
                bos_token=target_dict.bos_word,
                eos_token=target_dict.eos_word,
                word_delimiter_token="|",
                do_lower_case=False,
            )
            return_attention_mask = True if config.feat_extract_norm == "layer" else False
            feature_extractor = Wav2Vec2FeatureExtractor(
                feature_size=1,
                sampling_rate=16000,
                padding_value=0,
                do_normalize=True,
                return_attention_mask=return_attention_mask,
            )
            processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)
            processor.save_pretrained(pytorch_dump_folder_path)

        hf_wav2vec = Wav2Vec2ForCTC(config)
    else:
        hf_wav2vec = Wav2Vec2ForPreTraining(config)

    if is_finetuned or is_seq_class:
        model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task(
            [checkpoint_path], arg_overrides={"data": "/".join(dict_path.split("/")[:-1])}
        )
    else:
        task_arg = argparse.Namespace(task="audio_pretraining")
        task = fairseq.tasks.setup_task(task_arg)

        model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path], task=task)

    model = model[0].eval()

    recursively_load_weights(model, hf_wav2vec, not is_finetuned)

    hf_wav2vec.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"
    )
    parser.add_argument(
        "--is_seq_class",
        action="store_true",
        help="Whether the model to convert is a fine-tuned sequence classification model or not",
    )
    args = parser.parse_args()

    is_finetuned = not args.not_finetuned and not args.is_seq_class
    convert_wav2vec2_checkpoint(
        args.checkpoint_path,
        args.pytorch_dump_folder_path,
        args.config_path,
        args.dict_path,
        is_finetuned,
        args.is_seq_class,
    )