56 lines
2.3 KiB
Python
56 lines
2.3 KiB
Python
import torch, torch.nn as nn
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
|
|
source_dir = "/mnt/str/models/qwen2-0.5b-instruct"
|
|
target_dir = "/mnt/str/models/llama3-70b-instruct"
|
|
output_dir = "/mnt/str/temp/transplant"
|
|
|
|
# Load model and tokenizers
|
|
model = AutoModelForCausalLM.from_pretrained(source_dir, device_map = "auto")
|
|
tokenizer_source = AutoTokenizer.from_pretrained(source_dir)
|
|
tokenizer_target = AutoTokenizer.from_pretrained(target_dir)
|
|
tied = model.config.tie_word_embeddings
|
|
target_vocab_size = max(tokenizer_target.vocab.values()) + 1 # vocab_size member seems to be unreliable
|
|
|
|
# Embedding tensor
|
|
old_emb = model.model.embed_tokens.weight
|
|
new_emb = torch.empty((target_vocab_size, model.config.hidden_size),
|
|
dtype = old_emb.dtype, device = old_emb.device)
|
|
|
|
# Head tensor
|
|
old_head = model.lm_head.weight
|
|
new_head = torch.empty((target_vocab_size, model.config.hidden_size),
|
|
dtype = old_head.dtype, device = old_head.device)
|
|
|
|
# Initialize new tensors
|
|
for idx in range(target_vocab_size):
|
|
decode = tokenizer_target.decode(torch.tensor(idx, dtype = torch.long), decode_special_tokens = True)
|
|
encode = tokenizer_source.encode(decode, add_special_tokens = False, return_tensors = "pt")
|
|
new_emb[idx] = old_emb[encode.flatten()].mean(dim = 0)
|
|
new_head[idx] = old_head[encode.flatten()].mean(dim = 0)
|
|
|
|
# Replace embedding tensor
|
|
model.model.embed_tokens.weight = nn.Parameter(new_emb, requires_grad = False)
|
|
model.model.embed_tokens.num_embeddings = target_vocab_size
|
|
|
|
# Replace head tensor
|
|
model.lm_head.weight = nn.Parameter(new_head, requires_grad = False)
|
|
model.lm_head.out_features = tokenizer_target.vocab_size
|
|
|
|
# Update model
|
|
model.vocab_size = target_vocab_size
|
|
model.config.vocab_size = target_vocab_size
|
|
model.config.bos_token_id = tokenizer_target.bos_token_id
|
|
model.config.eos_token_id = tokenizer_target.eos_token_id
|
|
|
|
# Save
|
|
model.save_pretrained(output_dir, tie_word_embeddings = tied)
|
|
tokenizer_target.save_pretrained(output_dir)
|
|
|
|
# This is more reliable since save_pretrained seems to gives you a messed up model with some architectures,
|
|
# but it requires manually copying and modifying config.json etc.:
|
|
#
|
|
# import os
|
|
# from safetensors.torch import save_file
|
|
# save_file(model.state_dict(), os.path.join(args.output_dir, "model.safetensors"), metadata = {'format': 'pt'})
|