#!/usr/bin/env python3
"""Convert the fairseq models available in voxpopuli repo https://github.com/facebookresearch/voxpopuli

The available checkpoints should open with fairseq.
But the following error cannot be resolved with almost any version of fairseq.
https://github.com/facebookresearch/voxpopuli/issues/29

So this script manually parse the checkpoint file and reconstruct the model.

Examples

```
python convert_voxpopuli_models.py \
  --input-file wav2vec2_base_10k_ft_fr.pt \
  --output-file wav2vec2_voxpopuli_base_10k_asr_fr.pt
```
"""


def _parse_args():
    import argparse

    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument("--input-file", required=True, help="Input checkpoint file.")
    parser.add_argument("--output-file", required=False, help="Output model file.")
    return parser.parse_args()


def _removeprefix(s, prefix):
    if s.startswith(prefix):
        return s[len(prefix) :]
    return s


def _load(input_file):
    import torch
    from omegaconf import OmegaConf

    data = torch.load(input_file)
    cfg = OmegaConf.to_container(data["cfg"])
    for key in list(cfg.keys()):
        if key != "model":
            del cfg[key]
            if "w2v_args" in cfg["model"]:
                del cfg["model"]["w2v_args"][key]
    state_dict = {_removeprefix(k, "w2v_encoder."): v for k, v in data["model"].items()}
    return cfg, state_dict


def _parse_model_param(cfg, state_dict):
    key_mapping = {
        "extractor_mode": "extractor_mode",
        "conv_feature_layers": "extractor_conv_layer_config",
        "conv_bias": "extractor_conv_bias",
        "encoder_embed_dim": "encoder_embed_dim",
        "dropout_input": "encoder_projection_dropout",
        "conv_pos": "encoder_pos_conv_kernel",
        "conv_pos_groups": "encoder_pos_conv_groups",
        "encoder_layers": "encoder_num_layers",
        "encoder_attention_heads": "encoder_num_heads",
        "attention_dropout": "encoder_attention_dropout",
        "encoder_ffn_embed_dim": "encoder_ff_interm_features",
        "activation_dropout": "encoder_ff_interm_dropout",
        "dropout": "encoder_dropout",
        "layer_norm_first": "encoder_layer_norm_first",
        "layerdrop": "encoder_layer_drop",
        "encoder_layerdrop": "encoder_layer_drop",
    }
    params = {}
    src_dicts = [cfg["model"]]
    if "w2v_args" in cfg["model"]:
        src_dicts.append(cfg["model"]["w2v_args"]["model"])

    for src, tgt in key_mapping.items():
        for model_cfg in src_dicts:
            if src in model_cfg:
                params[tgt] = model_cfg[src]
                break
    if params["extractor_mode"] == "default":
        params["extractor_mode"] = "group_norm"
    # the following line is commented out to resolve lint warning; uncomment before running script
    # params["extractor_conv_layer_config"] = eval(params["extractor_conv_layer_config"])
    assert len(params) == 15
    params["aux_num_out"] = state_dict["proj.bias"].numel() if "proj.bias" in state_dict else None
    return params


def _main(args):
    import json

    import torch
    import torchaudio
    from torchaudio.models.wav2vec2.utils.import_fairseq import (
        _convert_state_dict as _convert,
    )

    cfg, state_dict = _load(args.input_file)
    params = _parse_model_param(cfg, state_dict)
    print(json.dumps(params, indent=4))
    model = torchaudio.models.wav2vec2_model(**params)
    model.load_state_dict(_convert(state_dict))
    torch.save(model.state_dict(), args.output_file)


if __name__ == "__main__":
    _main(_parse_args())
