import torch
from transformers import AutoModelForSpeechSeq2Seq
import re
from tqdm import tqdm
from collections import OrderedDict
import json
with open('convert_hf2openai.json', 'r') as f:
reverse_translation = json.load(f)
reverse_translation = OrderedDict(reverse_translation)
def save_model(model, save_path):
def reverse_translate(current_param):
for pattern, repl in reverse_translation.items():
if re.match(pattern, current_param):
return re.sub(pattern, repl, current_param)
config = model.config
model_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,
}
original_model_state_dict = model.state_dict()
new_state_dict = {}
for key, value in tqdm(original_model_state_dict.items()):
key = key.replace("model.", "")
new_key = reverse_translate(key)
if new_key is not None:
new_state_dict[new_key] = value
pytorch_model = {"dims": model_dims, "model_state_dict": new_state_dict}
torch.save(pytorch_model, save_path)
model_id = "Oriserve/Whisper-Hindi2Hinglish-Swift"
model = AutoModelForSpeechSeq2Seq.from_pretrained(
model_id,
low_cpu_mem_usage=True,
use_safetensors=True
)
model_save_path = "Whisper-Hindi2Hinglish-Swift.pt"
save_model(model,model_save_path)