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 hubert checkpoint import argparse import torch from transformers import unispeechsatconfig unispeechsatforaudioframeclassification unispeechsatforsequenceclassification unispeechsatforxvector wav2vec2featureextractor logging logging setverbosityinfo logger logging getloggername def convertclassificationbasemodelname hfconfig downstreamdict model unispeechsatforsequenceclassification frompretrainedbasemodelname confighfconfig model projector weight data downstreamdictprojector weight model projector bias data downstreamdictprojector bias model classifier weight data downstreamdictmodel postnet linear weight model classifier bias data downstreamdictmodel postnet linear bias return model def convertdiarizationbasemodelname hfconfig downstreamdict model unispeechsatforaudioframeclassification frompretrainedbasemodelname confighfconfig model classifier weight data downstreamdictmodel linear weight model classifier bias data downstreamdictmodel linear bias return model def convertxvectorbasemodelname hfconfig downstreamdict model unispeechsatforxvector frompretrainedbasemodelname confighfconfig model projector weight data downstreamdictconnector weight model projector bias data downstreamdictconnector bias for i kernelsize in enumeratehfconfig tdnnkernel model tdnni kernel weight data downstreamdict fmodel framelevelfeatureextractor module i kernel weight model tdnni kernel bias data downstreamdictfmodel framelevelfeatureextractor module i kernel bias model featureextractor weight data downstreamdictmodel utterancelevelfeatureextractor linear1 weight model featureextractor bias data downstreamdictmodel utterancelevelfeatureextractor linear1 bias model classifier weight data downstreamdictmodel utterancelevelfeatureextractor linear2 weight model classifier bias data downstreamdictmodel utterancelevelfeatureextractor linear2 bias model objective weight data downstreamdictobjective w return model torch nograd def converts3prlcheckpointbasemodelname configpath checkpointpath modeldumppath checkpoint torch loadcheckpointpath maplocationcpu downstreamdict checkpointdownstream hfconfig unispeechsatconfig frompretrainedconfigpath hffeatureextractor wav2vec2featureextractor frompretrained basemodelname returnattentionmasktrue donormalizefalse arch hfconfig architectures0 if arch endswithforsequenceclassification hfmodel convertclassificationbasemodelname hfconfig downstreamdict elif arch endswithforaudioframeclassification hfmodel convertdiarizationbasemodelname hfconfig downstreamdict elif arch endswithforxvector hfmodel convertxvectorbasemodelname hfconfig downstreamdict else raise notimplementederrorfs3prl weights conversion is not supported for arch if hfconfig useweightedlayersum hfmodel layerweights data checkpointfeaturizerweights hffeatureextractor savepretrainedmodeldumppath hfmodel savepretrainedmodeldumppath if name main parser argparse argumentparser parser addargument basemodelname defaultnone typestr helpname of the huggingface pretrained base model parser addargumentconfigpath defaultnone typestr helppath to the huggingface classifier config parser addargumentcheckpointpath defaultnone typestr helppath to the s3prl checkpoint parser addargumentmodeldumppath defaultnone typestr helppath to the final converted model args parser parseargs converts3prlcheckpointargs basemodelname args configpath args checkpointpath args modeldumppath 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 hubert checkpoint copy paste tweak model s weights to transformers design
import argparse

import torch

from transformers import (
    UniSpeechSatConfig,
    UniSpeechSatForAudioFrameClassification,
    UniSpeechSatForSequenceClassification,
    UniSpeechSatForXVector,
    Wav2Vec2FeatureExtractor,
    logging,
)


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


def convert_classification(base_model_name, hf_config, downstream_dict):
    model = UniSpeechSatForSequenceClassification.from_pretrained(base_model_name, config=hf_config)
    model.projector.weight.data = downstream_dict["projector.weight"]
    model.projector.bias.data = downstream_dict["projector.bias"]
    model.classifier.weight.data = downstream_dict["model.post_net.linear.weight"]
    model.classifier.bias.data = downstream_dict["model.post_net.linear.bias"]
    return model


def convert_diarization(base_model_name, hf_config, downstream_dict):
    model = UniSpeechSatForAudioFrameClassification.from_pretrained(base_model_name, config=hf_config)
    model.classifier.weight.data = downstream_dict["model.linear.weight"]
    model.classifier.bias.data = downstream_dict["model.linear.bias"]
    return model


def convert_xvector(base_model_name, hf_config, downstream_dict):
    model = UniSpeechSatForXVector.from_pretrained(base_model_name, config=hf_config)
    model.projector.weight.data = downstream_dict["connector.weight"]
    model.projector.bias.data = downstream_dict["connector.bias"]
    for i, kernel_size in enumerate(hf_config.tdnn_kernel):
        model.tdnn[i].kernel.weight.data = downstream_dict[
            f"model.framelevel_feature_extractor.module.{i}.kernel.weight"
        ]
        model.tdnn[i].kernel.bias.data = downstream_dict[f"model.framelevel_feature_extractor.module.{i}.kernel.bias"]

    model.feature_extractor.weight.data = downstream_dict["model.utterancelevel_feature_extractor.linear1.weight"]
    model.feature_extractor.bias.data = downstream_dict["model.utterancelevel_feature_extractor.linear1.bias"]
    model.classifier.weight.data = downstream_dict["model.utterancelevel_feature_extractor.linear2.weight"]
    model.classifier.bias.data = downstream_dict["model.utterancelevel_feature_extractor.linear2.bias"]
    model.objective.weight.data = downstream_dict["objective.W"]
    return model


@torch.no_grad()
def convert_s3prl_checkpoint(base_model_name, config_path, checkpoint_path, model_dump_path):
    
    checkpoint = torch.load(checkpoint_path, map_location="cpu")

    downstream_dict = checkpoint["Downstream"]

    hf_config = UniSpeechSatConfig.from_pretrained(config_path)
    hf_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
        base_model_name, return_attention_mask=True, do_normalize=False
    )

    arch = hf_config.architectures[0]
    if arch.endswith("ForSequenceClassification"):
        hf_model = convert_classification(base_model_name, hf_config, downstream_dict)
    elif arch.endswith("ForAudioFrameClassification"):
        hf_model = convert_diarization(base_model_name, hf_config, downstream_dict)
    elif arch.endswith("ForXVector"):
        hf_model = convert_xvector(base_model_name, hf_config, downstream_dict)
    else:
        raise NotImplementedError(f"S3PRL weights conversion is not supported for {arch}")

    if hf_config.use_weighted_layer_sum:
        hf_model.layer_weights.data = checkpoint["Featurizer"]["weights"]

    hf_feature_extractor.save_pretrained(model_dump_path)
    hf_model.save_pretrained(model_dump_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--base_model_name", default=None, type=str, help="Name of the huggingface pretrained base model."
    )
    parser.add_argument("--config_path", default=None, type=str, help="Path to the huggingface classifier config.")
    parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to the s3prl checkpoint.")
    parser.add_argument("--model_dump_path", default=None, type=str, help="Path to the final converted model.")
    args = parser.parse_args()
    convert_s3prl_checkpoint(args.base_model_name, args.config_path, args.checkpoint_path, args.model_dump_path)