wenet/whisper/convert_whisper_to_wenet_config_and_ckpt.py

# Copyright (c) 2023 Wenet Community. (authors: Xingchen Song)
#
# 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.
"""
Requirements:

```bash
pip install -U openai-whisper
```

Example:

```bash
# Converts the model from OpenAI to WeNet format:
python convert_whisper_to_wenet_config_and_ckpt.py \
    --whisper_ckpt large-v3.pt \
    --output_dir exp/whisper/large-v3
```
"""

import argparse
import copy
import os
import sys
import torch
import yaml

_cpath_ = sys.path[0]
sys.path.remove(_cpath_)
from whisper.tokenizer import get_tokenizer

sys.path.insert(0, _cpath_)


def convert_to_wenet_yaml(tokenizer, dims, wenet_yaml_path: str):
    configs = {}
    configs['input_dim'] = dims['n_mels']
    configs['output_dim'] = dims['n_vocab']
    assert dims['n_vocab'] == tokenizer.encoding.n_vocab, "{} v.s. {}".format(
        dims['n_vocab'], tokenizer.encoding.n_vocab)

    configs['encoder'] = 'transformer'
    configs['encoder_conf'] = {}
    configs['encoder_conf']['gradient_checkpointing'] = True
    configs['encoder_conf']['input_layer'] = 'conv1d2'
    configs['encoder_conf']['output_size'] = dims['n_audio_state']
    configs['encoder_conf']['attention_heads'] = dims['n_audio_head']
    configs['encoder_conf']['linear_units'] = dims['n_audio_state'] * 4
    configs['encoder_conf']['num_blocks'] = dims['n_audio_layer']
    configs['encoder_conf']['dropout_rate'] = 0.1
    configs['encoder_conf']['positional_dropout_rate'] = 0.1
    configs['encoder_conf']['attention_dropout_rate'] = 0.0
    configs['encoder_conf']['normalize_before'] = True
    configs['encoder_conf']['use_dynamic_chunk'] = False
    configs['encoder_conf']['use_dynamic_left_chunk'] = False
    configs['encoder_conf']['pos_enc_layer_type'] = "abs_pos_whisper"
    configs['encoder_conf']['static_chunk_size'] = -1
    configs['encoder_conf']['key_bias'] = False
    configs['encoder_conf']['activation_type'] = "gelu"

    configs['decoder'] = 'transformer'
    configs['decoder_conf'] = {}
    configs['decoder_conf']['tie_word_embedding'] = True
    configs['decoder_conf']['gradient_checkpointing'] = True
    configs['decoder_conf']['attention_heads'] = dims['n_text_head']
    configs['decoder_conf']['linear_units'] = dims['n_text_state'] * 4
    configs['decoder_conf']['num_blocks'] = dims['n_text_layer']
    configs['decoder_conf']['dropout_rate'] = 0.1
    configs['decoder_conf']['positional_dropout_rate'] = 0.1
    configs['decoder_conf']['self_attention_dropout_rate'] = 0.0
    configs['decoder_conf']['src_attention_dropout_rate'] = 0.0
    configs['decoder_conf']['input_layer'] = "embed_learnable_pe"
    configs['decoder_conf']['use_output_layer'] = True
    configs['decoder_conf']['normalize_before'] = True
    configs['decoder_conf']['src_attention'] = True
    configs['decoder_conf']['key_bias'] = False
    configs['decoder_conf']['activation_type'] = "gelu"

    configs['tokenizer'] = 'whisper'
    configs['tokenizer_conf'] = {}
    configs['tokenizer_conf']['is_multilingual'] = dims['n_vocab'] >= 51865
    configs['tokenizer_conf']['num_languages'] = dims['n_vocab'] - 51765 - \
        int(configs['tokenizer_conf']['is_multilingual'])
    configs['tokenizer_conf']['split_with_space'] = False
    configs['tokenizer_conf']['bpe_path'] = None
    configs['tokenizer_conf']['symbol_table_path'] = None
    configs['tokenizer_conf']['non_lang_syms_path'] = None
    configs['tokenizer_conf']['special_tokens'] = {}
    configs['tokenizer_conf']['special_tokens']['sot'] = tokenizer.sot
    configs['tokenizer_conf']['special_tokens']['eot'] = tokenizer.eot
    configs['tokenizer_conf']['special_tokens'][
        'sot_prev'] = tokenizer.sot_prev
    configs['tokenizer_conf']['special_tokens'][
        'transcribe'] = tokenizer.transcribe
    configs['tokenizer_conf']['special_tokens'][
        'translate'] = tokenizer.translate
    configs['tokenizer_conf']['special_tokens'][
        'no_timestamps'] = tokenizer.no_timestamps
    configs['tokenizer_conf']['special_tokens'][
        'no_speech'] = tokenizer.no_speech
    configs['tokenizer_conf']['special_tokens']['timestamp_begin'] = \
        tokenizer.timestamp_begin

    configs['ctc_conf'] = {}
    configs['ctc_conf']['ctc_blank_id'] = tokenizer.no_speech

    configs['cmvn'] = None
    configs['cmvn_conf'] = {}
    configs['cmvn_conf']['cmvn_file'] = None
    configs['cmvn_conf']['is_json_cmvn'] = None

    configs['model'] = "whisper"
    configs['model_conf'] = {}
    configs['model_conf']['ctc_weight'] = 0.3
    configs['model_conf']['lsm_weight'] = 0.1
    configs['model_conf']['length_normalized_loss'] = False

    configs['dataset'] = "asr"
    configs['dataset_conf'] = {}
    configs['dataset_conf']['filter_conf'] = {}
    configs['dataset_conf']['filter_conf'][
        'max_length'] = dims['n_audio_ctx'] * 2  # 1/2 subsample # noqa
    configs['dataset_conf']['filter_conf']['min_length'] = 0
    configs['dataset_conf']['filter_conf']['token_max_length'] = dims[
        'n_text_ctx']
    configs['dataset_conf']['filter_conf']['token_min_length'] = 1
    configs['dataset_conf']['resample_conf'] = {}
    configs['dataset_conf']['resample_conf']['resample_rate'] = 16000
    # NOTE: Disable speed_perturb, https://github.com/wenet-e2e/wenet/issues/2171
    configs['dataset_conf']['speed_perturb'] = False
    configs['dataset_conf']['spec_aug'] = True
    configs['dataset_conf']['spec_aug_conf'] = {}
    configs['dataset_conf']['spec_aug_conf']['num_t_mask'] = 2
    configs['dataset_conf']['spec_aug_conf']['num_f_mask'] = 2
    configs['dataset_conf']['spec_aug_conf']['max_t'] = 50
    configs['dataset_conf']['spec_aug_conf']['max_f'] = 10
    configs['dataset_conf']['spec_sub'] = True
    configs['dataset_conf']['spec_sub_conf'] = {}
    configs['dataset_conf']['spec_sub_conf']['num_t_sub'] = 3
    configs['dataset_conf']['spec_sub_conf']['max_t'] = 30
    configs['dataset_conf']['spec_trim'] = False
    configs['dataset_conf']['shuffle'] = True
    configs['dataset_conf']['shuffle_conf'] = {}
    configs['dataset_conf']['shuffle_conf']['shuffle_size'] = 1500
    configs['dataset_conf']['sort'] = True
    configs['dataset_conf']['sort_conf'] = {}
    configs['dataset_conf']['sort_conf']['sort_size'] = 500
    configs['dataset_conf']['feats_type'] = "log_mel_spectrogram"
    configs['dataset_conf']['log_mel_spectrogram_conf'] = {}
    configs['dataset_conf']['log_mel_spectrogram_conf']['n_fft'] = 400
    configs['dataset_conf']['log_mel_spectrogram_conf']['hop_length'] = 160
    configs['dataset_conf']['log_mel_spectrogram_conf']['num_mel_bins'] = dims[
        'n_mels']
    configs['dataset_conf']['log_mel_spectrogram_conf']['padding'] = 0
    configs['dataset_conf']['batch_conf'] = {}
    configs['dataset_conf']['batch_conf']['batch_type'] = 'dynamic'
    configs['dataset_conf']['batch_conf']['batch_size'] = 26
    configs['dataset_conf']['batch_conf']['max_frames_in_batch'] = 12000
    configs['dataset_conf']['language_conf'] = {}
    configs['dataset_conf']['language_conf']['limited_langs'] = ['zh']

    configs['grad_clip'] = 5
    configs['accum_grad'] = 4
    configs['max_epoch'] = 100
    configs['log_interval'] = 100

    configs['optim'] = "adam"
    configs['optim_conf'] = {}
    configs['optim_conf']['lr'] = 0.0005
    configs['scheduler'] = "warmuplr"
    configs['scheduler_conf'] = {}
    configs['scheduler_conf']['warmup_steps'] = 12000

    with open(wenet_yaml_path, '+w') as f:
        f.write(yaml.dump(configs))
        f.flush()

    print(configs)


def convert_to_wenet_state_dict(whisper_state_dict, wenet_state_dict_path):
    wenet_state_dict = {}
    unused = []
    print(
        "===================== start CKPT Conversion ========================="
    )
    for name in whisper_state_dict.keys():
        original_name = copy.deepcopy(name)
        name = name.replace("encoder.conv1", "encoder.embed.conv.0")
        name = name.replace("encoder.conv2", "encoder.embed.conv.2")
        name = name.replace("decoder.token_embedding", "decoder.embed.0")
        name = name.replace("encoder.blocks", "encoder.encoders")
        name = name.replace("decoder.blocks", "decoder.decoders")
        name = name.replace(".cross_attn.query", ".src_attn.linear_q")
        name = name.replace(".cross_attn.key", ".src_attn.linear_k")
        name = name.replace(".cross_attn.value", ".src_attn.linear_v")
        name = name.replace(".cross_attn.out", ".src_attn.linear_out")
        name = name.replace(".attn.query", ".self_attn.linear_q")
        name = name.replace(".attn.key", ".self_attn.linear_k")
        name = name.replace(".attn.value", ".self_attn.linear_v")
        name = name.replace(".attn.out", ".self_attn.linear_out")
        name = name.replace("mlp.0", "feed_forward.w_1")
        name = name.replace("mlp.2", "feed_forward.w_2")
        if "decoder" in name:
            name = name.replace("cross_attn_ln", "norm2")
            name = name.replace("mlp_ln", "norm3")
        else:
            name = name.replace("mlp_ln", "norm2")
        name = name.replace("attn_ln", "norm1")
        name = name.replace("encoder.ln_post", "encoder.after_norm")
        name = name.replace("decoder.ln", "decoder.after_norm")
        if original_name == "decoder.positional_embedding":
            whisper_state_dict[name] = whisper_state_dict[name].unsqueeze(0)
            name = "decoder.embed.1.pe"
        elif original_name == "encoder.positional_embedding":
            whisper_state_dict[name] = whisper_state_dict[name].unsqueeze(0)
            name = "encoder.embed.pos_enc.pe"
        print("name  {} ==> {}".format(original_name, name))
        print("type  {} ==> torch.float32".format(
            whisper_state_dict[original_name].dtype))
        print("shape {}\n".format(whisper_state_dict[original_name].shape))
        if (original_name == name):
            unused.append(name)
        else:
            wenet_state_dict[name] = whisper_state_dict[original_name].float()
    for name in unused:
        print("NOTE!!! drop {}".format(name))
    print("Saving fp32 ckpt to {}...".format(wenet_state_dict_path))
    torch.save(wenet_state_dict, wenet_state_dict_path)
    print(
        "DONE\n===================== End CKPT Conversion =========================\n"
    )


def convert_to_wenet_units(tokenizer, units_txt_path):
    """ NOTE(xcsong):
        The "units.txt" file is solely for adapting to the training API of Wenet
        and for quickly checking the corresponding text of an ID when necessary.
        It does not play any role in the tokenization process,
        which is carried out by the tokenizer of openai-whisper.
    """
    n_vocab = tokenizer.encoding.n_vocab
    with open(units_txt_path, "+w") as f:
        for i in range(n_vocab):
            unit = str(tokenizer.encoding.decode_single_token_bytes(i))
            if len(unit) == 0:
                unit = str(i)
                print("can not decode id {}, convert to str({})".format(i, i))
            unit = unit.replace(" ", "<space>")
            f.write("{} {}\n".format(unit, i))
            f.flush()


def get_args():
    parser = argparse.ArgumentParser(description='load and parse whisper')
    # yapf: disable
    parser.add_argument(
        '--whisper_ckpt',
        required=True,
        help='https://openaipublic.azureedge.net/main/whisper/models/e5b1a55b89c1367dacf97e3e19bfd829a01529dbfdeefa8caeb59b3f1b81dadb/large-v3.pt'  # noqa
    )
    # yapf: enable
    parser.add_argument('--output_dir',
                        default='.',
                        help='output file in wenet\'s style: ' +
                        'units.txt, train.yaml, model.pt')
    args = parser.parse_args()
    return args


def main():
    args = get_args()
    checkpoint = torch.load(args.whisper_ckpt, map_location="cpu")
    multilingual = checkpoint["dims"]['n_vocab'] >= 51865
    num_languages = checkpoint["dims"]['n_vocab'] - 51765 - int(multilingual)
    tokenizer = get_tokenizer(multilingual=multilingual,
                              num_languages=num_languages)

    convert_to_wenet_state_dict(
        checkpoint["model_state_dict"],
        os.path.join(args.output_dir, 'wenet_whisper.pt'))
    convert_to_wenet_units(tokenizer, os.path.join(args.output_dir,
                                                   'units.txt'))
    convert_to_wenet_yaml(tokenizer, checkpoint["dims"],
                          os.path.join(args.output_dir, 'train.yaml'))


if __name__ == "__main__":

    main()