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 hubertconfig hubertforsequenceclassification wav2vec2featureextractor logging logging setverbosityinfo logger logging getloggername supportedmodels utterancelevel torch nograd def converts3prlcheckpointbasemodelname configpath checkpointpath modeldumppath checkpoint torch loadcheckpointpath maplocationcpu if checkpointconfigdownstreamexpertmodelrcselect not in supportedmodels raise notimplementederrorfthe supported s3prl models are supportedmodels downstreamdict checkpointdownstream hfcongfig hubertconfig frompretrainedconfigpath hfmodel hubertforsequenceclassification frompretrainedbasemodelname confighfcongfig hffeatureextractor wav2vec2featureextractor frompretrained basemodelname returnattentionmasktrue donormalizefalse if hfcongfig useweightedlayersum hfmodel layerweights data checkpointfeaturizerweights hfmodel projector weight data downstreamdictprojector weight hfmodel projector bias data downstreamdictprojector bias hfmodel classifier weight data downstreamdictmodel postnet linear weight hfmodel classifier bias data downstreamdictmodel postnet linear bias 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 HubertConfig, HubertForSequenceClassification, Wav2Vec2FeatureExtractor, logging


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

SUPPORTED_MODELS = ["UtteranceLevel"]


@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")
    if checkpoint["Config"]["downstream_expert"]["modelrc"]["select"] not in SUPPORTED_MODELS:
        raise NotImplementedError(f"The supported s3prl models are {SUPPORTED_MODELS}")

    downstream_dict = checkpoint["Downstream"]

    hf_congfig = HubertConfig.from_pretrained(config_path)
    hf_model = HubertForSequenceClassification.from_pretrained(base_model_name, config=hf_congfig)
    hf_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
        base_model_name, return_attention_mask=True, do_normalize=False
    )

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

    hf_model.projector.weight.data = downstream_dict["projector.weight"]
    hf_model.projector.bias.data = downstream_dict["projector.bias"]
    hf_model.classifier.weight.data = downstream_dict["model.post_net.linear.weight"]
    hf_model.classifier.bias.data = downstream_dict["model.post_net.linear.bias"]

    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)