#!/usr/bin/env python3
"""Convert a Wav2Vec2/HuBERT model published by fairseq into torchaudio format

Examples

```
python convert_fairseq_models.py \
  --input-file hubert_base_ls960.pt \
  --output-file hubert_fairseq_base_ls960.pth

python convert_fairseq_models.py \
  --input-file hubert_large_ll60k.pt \
  --output-file hubert_fairseq_large_ll60k.pth

python convert_fairseq_models.py \
  --input-file hubert_large_ll60k_finetune_ls960.pt \
  --output-file hubert_fairseq_large_ll60k_asr_ls960.pth

python convert_fairseq_models.py \
  --input-file hubert_xtralarge_ll60k.pt \
  --output-file hubert_fairseq_xlarge_ll60k.pth

python convert_fairseq_models.py \
  --input-file hubert_xtralarge_ll60k_finetune_ls960.pt \
  --output-file hubert_fairseq_xlarge_ll60k_asr_ls960.pth
"""

import argparse

# Note: Avoiding the import of torch and fairseq on global scope as they are slow


def _parse_args():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument("--input-file", required=True, help="Input model file.")
    parser.add_argument("--output-file", required=False, help="Output model file.")
    parser.add_argument(
        "--dict-dir",
        help=(
            "Directory where letter vocabulary file, `dict.ltr.txt`, is found. "
            "Required when loading wav2vec2 model. "
            "https://dl.fbaipublicfiles.com/fairseq/wav2vec/dict.ltr.txt"
        ),
    )
    return parser.parse_args()


def _load_model(input_file, dict_dir):
    import fairseq

    overrides = {} if dict_dir is None else {"data": dict_dir}
    models, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task(
        [input_file],
        arg_overrides=overrides,
    )
    return models[0]


def _import_model(model):
    from torchaudio.models.wav2vec2.utils import import_fairseq_model

    if model.__class__.__name__ in ["HubertCtc", "Wav2VecCtc"]:
        model = model.w2v_encoder
    model = import_fairseq_model(model)
    return model


def _main(args):
    import torch

    model = _load_model(args.input_file, args.dict_dir)
    model = _import_model(model)
    torch.save(model.state_dict(), args.output_file)


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