| |
|
|
| from copy import deepcopy |
| import torch |
|
|
| WHISPER_MAPPING = { |
| "layers": "blocks", |
| "fc1": "mlp.0", |
| "fc2": "mlp.2", |
| "final_layer_norm": "mlp_ln", |
| "layers": "blocks", |
| ".self_attn.q_proj": ".attn.query", |
| ".self_attn.k_proj": ".attn.key", |
| ".self_attn.v_proj": ".attn.value", |
| ".self_attn_layer_norm": ".attn_ln", |
| ".self_attn.out_proj": ".attn.out", |
| ".encoder_attn.q_proj": ".cross_attn.query", |
| ".encoder_attn.k_proj": ".cross_attn.key", |
| ".encoder_attn.v_proj": ".cross_attn.value", |
| ".encoder_attn_layer_norm": ".cross_attn_ln", |
| ".encoder_attn.out_proj": ".cross_attn.out", |
| "decoder.layer_norm.": "decoder.ln.", |
| "encoder.layer_norm.": "encoder.ln_post.", |
| "embed_tokens": "token_embedding", |
| "encoder.embed_positions.weight": "encoder.positional_embedding", |
| "decoder.embed_positions.weight": "decoder.positional_embedding", |
| "layer_norm": "ln_post", |
| } |
|
|
|
|
| def rename_keys(s_dict): |
| keys = list(s_dict.keys()) |
| for key in keys: |
| new_key = key |
| for k, v in WHISPER_MAPPING.items(): |
| if k in key: |
| new_key = new_key.replace(k, v) |
|
|
| print(f"{key} -> {new_key}") |
|
|
| s_dict[new_key] = s_dict.pop(key) |
| return s_dict |
|
|
|
|
| def convert_hf_whisper(hf_model_name_or_path: str, whisper_state_path: str): |
| from transformers import WhisperForConditionalGeneration |
| transformer_model = WhisperForConditionalGeneration.from_pretrained(hf_model_name_or_path) |
| config = transformer_model.config |
|
|
| |
| dims = { |
| 'n_mels': config.num_mel_bins, |
| 'n_vocab': config.vocab_size, |
| 'n_audio_ctx': config.max_source_positions, |
| 'n_audio_state': config.d_model, |
| 'n_audio_head': config.encoder_attention_heads, |
| 'n_audio_layer': config.encoder_layers, |
| 'n_text_ctx': config.max_target_positions, |
| 'n_text_state': config.d_model, |
| 'n_text_head': config.decoder_attention_heads, |
| 'n_text_layer': config.decoder_layers |
| } |
|
|
| state_dict = deepcopy(transformer_model.model.state_dict()) |
| state_dict = rename_keys(state_dict) |
|
|
| torch.save({"dims": dims, "model_state_dict": state_dict}, whisper_state_path) |
