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 sew checkpoint import argparse import json import os import fairseq import torch from fairseq data import dictionary register sew s fairseq modules from sewasapp import tasks noqa f401 from transformers import sewconfig sewforctc sewmodel wav2vec2ctctokenizer wav2vec2featureextractor wav2vec2processor logging logging setverbosityinfo logger logging getloggername mapping postextractproj featureprojection 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 upsample 0 encoder upsample projection encoder layernorm encoder layernorm w2vmodel layernorm layernorm w2vencoder proj lmhead maskemb maskedspecembed 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 isfinetuned unusedweights fairseqdict fairseqmodel statedict featureextractor hfmodel sew featureextractor if isfinetuned else 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 mappedkey sew mappedkey if isfinetuned and mappedkey lmhead 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 weight in name weighttype weight elif bias in name weighttype bias 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 def convertconfigmodel isfinetuned config sewconfig if isfinetuned fsconfig model w2vencoder w2vmodel cfg else fsconfig model cfg config convbias fsconfig convbias convlayers evalfsconfig convfeaturelayers config convdim x0 for x in convlayers config convkernel x1 for x in convlayers config convstride x2 for x in convlayers config featextractactivation gelu config featextractnorm layer if fsconfig extractormode layernorm else group config finaldropout 0 0 config hiddenact fsconfig activationfn name config hiddensize fsconfig encoderembeddim config initializerrange 0 02 config intermediatesize fsconfig encoderffnembeddim config layernormeps 1e5 config layerdrop fsconfig encoderlayerdrop config numattentionheads fsconfig encoderattentionheads config numconvposembeddinggroups fsconfig convposgroups config numconvposembeddings fsconfig convpos config numfeatextractlayers lenconvlayers config numhiddenlayers fsconfig encoderlayers config squeezefactor fsconfig squeezefactor take care of any params that are overridden by the wav2vecctc model if isfinetuned fsconfig model cfg config finaldropout fsconfig finaldropout config layerdrop fsconfig layerdrop config activationdropout fsconfig activationdropout config applyspecaugment fsconfig maskprob 0 or fsconfig maskchannelprob 0 config attentiondropout fsconfig attentiondropout config featprojdropout fsconfig dropoutinput config hiddendropout fsconfig dropout config maskfeaturelength fsconfig maskchannellength config maskfeatureprob fsconfig maskchannelprob config masktimelength fsconfig masklength config masktimeprob fsconfig maskprob config featureextractortype wav2vec2featureextractor config tokenizerclass wav2vec2ctctokenizer return config torch nograd def convertsewcheckpoint checkpointpath pytorchdumpfolderpath configpathnone dictpathnone isfinetunedtrue if isfinetuned model fairseq checkpointutils loadmodelensembleandtask checkpointpath argoverridesdata joindictpath split 1 else model fairseq checkpointutils loadmodelensembleandtaskcheckpointpath if configpath is not none config sewconfig frompretrainedconfigpath else config convertconfigmodel0 isfinetuned model model0 eval returnattentionmask true if config featextractnorm layer else false featureextractor wav2vec2featureextractor featuresize1 samplingrate16000 paddingvalue0 donormalizetrue returnattentionmaskreturnattentionmask if isfinetuned if dictpath targetdict dictionary loaddictpath important change bos pad token id since ctc symbol is pad and not s as in fairseq targetdict indicestargetdict bosword targetdict padindex targetdict indicestargetdict padword targetdict bosindex 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 with openvocabpath w encodingutf8 as vocabhandle json dumptargetdict indices vocabhandle tokenizer wav2vec2ctctokenizer vocabpath unktokentargetdict unkword padtokentargetdict padword bostokentargetdict bosword eostokentargetdict eosword worddelimitertoken dolowercasefalse processor wav2vec2processorfeatureextractorfeatureextractor tokenizertokenizer processor savepretrainedpytorchdumpfolderpath hfmodel sewforctcconfig else hfmodel sewmodelconfig featureextractor savepretrainedpytorchdumpfolderpath recursivelyloadweightsmodel hfmodel isfinetuned hfmodel 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 isfinetuned actionstoretrue helpwhether the model to convert is a finetuned model or not args parser parseargs convertsewcheckpoint args checkpointpath args pytorchdumpfolderpath args configpath args dictpath args isfinetuned 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 sew checkpoint register sew s fairseq modules noqa f401 take care of any params that are overridden by the wav2vecctc model 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
import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary


from sew_asapp import tasks  

from transformers import (
    SEWConfig,
    SEWForCTC,
    SEWModel,
    Wav2Vec2CTCTokenizer,
    Wav2Vec2FeatureExtractor,
    Wav2Vec2Processor,
    logging,
)


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

MAPPING = {
    "post_extract_proj": "feature_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.upsample.0": "encoder.upsample.projection",
    "encoder.layer_norm": "encoder.layer_norm",
    "w2v_model.layer_norm": "layer_norm",
    "w2v_encoder.proj": "lm_head",
    "mask_emb": "masked_spec_embed",
}


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, is_finetuned):
    unused_weights = []
    fairseq_dict = fairseq_model.state_dict()

    feature_extractor = hf_model.sew.feature_extractor if is_finetuned else 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():
                mapped_key = "sew." + mapped_key if (is_finetuned and mapped_key != "lm_head") 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 "weight" in name:
                        weight_type = "weight"
                    elif "bias" in name:
                        weight_type = "bias"
                    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)


def convert_config(model, is_finetuned):
    config = SEWConfig()
    if is_finetuned:
        fs_config = model.w2v_encoder.w2v_model.cfg
    else:
        fs_config = model.cfg

    config.conv_bias = fs_config.conv_bias
    conv_layers = eval(fs_config.conv_feature_layers)
    config.conv_dim = [x[0] for x in conv_layers]
    config.conv_kernel = [x[1] for x in conv_layers]
    config.conv_stride = [x[2] for x in conv_layers]
    config.feat_extract_activation = "gelu"
    config.feat_extract_norm = "layer" if fs_config.extractor_mode == "layer_norm" else "group"
    config.final_dropout = 0.0
    config.hidden_act = fs_config.activation_fn.name
    config.hidden_size = fs_config.encoder_embed_dim
    config.initializer_range = 0.02
    config.intermediate_size = fs_config.encoder_ffn_embed_dim
    config.layer_norm_eps = 1e-5
    config.layerdrop = fs_config.encoder_layerdrop
    config.num_attention_heads = fs_config.encoder_attention_heads
    config.num_conv_pos_embedding_groups = fs_config.conv_pos_groups
    config.num_conv_pos_embeddings = fs_config.conv_pos
    config.num_feat_extract_layers = len(conv_layers)
    config.num_hidden_layers = fs_config.encoder_layers
    config.squeeze_factor = fs_config.squeeze_factor

    
    if is_finetuned:
        fs_config = model.cfg
        config.final_dropout = fs_config.final_dropout
        config.layerdrop = fs_config.layerdrop
    config.activation_dropout = fs_config.activation_dropout
    config.apply_spec_augment = fs_config.mask_prob > 0 or fs_config.mask_channel_prob > 0
    config.attention_dropout = fs_config.attention_dropout
    config.feat_proj_dropout = fs_config.dropout_input
    config.hidden_dropout = fs_config.dropout
    config.mask_feature_length = fs_config.mask_channel_length
    config.mask_feature_prob = fs_config.mask_channel_prob
    config.mask_time_length = fs_config.mask_length
    config.mask_time_prob = fs_config.mask_prob

    config.feature_extractor_type = "Wav2Vec2FeatureExtractor"
    config.tokenizer_class = "Wav2Vec2CTCTokenizer"

    return config


@torch.no_grad()
def convert_sew_checkpoint(
    checkpoint_path, pytorch_dump_folder_path, config_path=None, dict_path=None, is_finetuned=True
):
    

    if is_finetuned:
        model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task(
            [checkpoint_path], arg_overrides={"data": "/".join(dict_path.split("/")[:-1])}
        )
    else:
        model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path])

    if config_path is not None:
        config = SEWConfig.from_pretrained(config_path)
    else:
        config = convert_config(model[0], is_finetuned)
    model = model[0].eval()

    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,
    )

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

            
            
            target_dict.indices[target_dict.bos_word] = target_dict.pad_index
            target_dict.indices[target_dict.pad_word] = target_dict.bos_index
            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)
            with open(vocab_path, "w", encoding="utf-8") as vocab_handle:
                json.dump(target_dict.indices, 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,
            )
            processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)
            processor.save_pretrained(pytorch_dump_folder_path)

        hf_model = SEWForCTC(config)
    else:
        hf_model = SEWModel(config)
        feature_extractor.save_pretrained(pytorch_dump_folder_path)

    recursively_load_weights(model, hf_model, is_finetuned)

    hf_model.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(
        "--is_finetuned", action="store_true", help="Whether the model to convert is a fine-tuned model or not"
    )
    args = parser.parse_args()
    convert_sew_checkpoint(
        args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, args.is_finetuned
    )