Skip to content

Instantly share code, notes, and snippets.

@mutaguchi
Created April 17, 2024 10:40
Show Gist options
  • Save mutaguchi/23c5cfcf872fb951af2c9dc67b5a52d2 to your computer and use it in GitHub Desktop.
Save mutaguchi/23c5cfcf872fb951af2c9dc67b5a52d2 to your computer and use it in GitHub Desktop.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
models = {
"stablelm": "stabilityai/japanese-stablelm-base-gamma-7b",
"chatntq": "NTQAI/chatntq-ja-7b-v1.0",
"mistral": "mistralai/Mistral-7B-v0.1",
"starling": "Nexusflow/Starling-LM-7B-beta",
"antler": "Elizezen/Antler-7B",
}
def load_model(model_name):
return {
"name": model_name,
"model": AutoModelForCausalLM.from_pretrained(models[model_name], torch_dtype=torch.bfloat16, device_map="cpu")
}
# ChatVectorを適用する対象モデル
cp_model = load_model("antler")
# ChatVector1を抽出するベースモデルとinstructionモデル
base_model1 = load_model("stablelm")
inst_model1 = load_model("chatntq")
ratio1 = 0.5 # マージ比率
# ChatVector2を抽出するベースモデルとinstructionモデル(Noneでも可)
base_model2 = load_model("mistral")
inst_model2 = load_model("starling")
ratio2 = 0.5 # マージ比率
skip_layers = ["model.embed_tokens.weight", "lm_head.weight"]
for k, v in cp_model["model"].state_dict().items():
if (k in skip_layers) or ("layernorm" in k):
continue
chat_vector1 = inst_model1["model"].state_dict()[
k] - base_model1["model"].state_dict()[k]
chat_vector2 = None
if base_model2 is not None and inst_model2 is not None:
chat_vector2 = inst_model2["model"].state_dict(
)[k] - base_model2["model"].state_dict()[k]
new_v = v + ratio1 * chat_vector1.to(v.device)
if chat_vector2 is not None:
new_v += ratio2 * chat_vector2.to(v.device)
v.copy_(new_v)
# 保存
path = f"./{cp_model['name']}-{inst_model1['name']}_{ratio1}-{inst_model2['name']}_{ratio2}-7B"
cp_model["model"].save_pretrained(path)
AutoTokenizer.from_pretrained(
models[cp_model['name']]).save_pretrained(path)
@mutaguchi
Copy link
Author

Mistral-7B派生の日本語モデルに、日本語Chat Vectorと英語Chat Vectorを0.5:0.5の比率でマージします。

レシピ:
Antler + 0.5 * (chatntq - StableLM) + 0.5 * (starling - mistral)

7Bモデルを5個用いるので、75GBくらいのファイルダウンロードが発生します。
モデル(.safetensors)はカレントディレクトリに生成されます。トークナイザーは、Antlerのものを自動で保存します。
モデル生成後は適宜、GGUFなどで量子化して利用して下さい。

参考:Chat Vectorを使って日本語LLMをチャットモデルに改造する #Python - Qiita

@mutaguchi
Copy link
Author

Antlerをベースにしたモデルは創造的な文章を書きますが、設定や指示に従ったり、チャット的な通常の会話は苦手です。

その場合、以下のようなchatntqをベースとしたレシピにすると、創造性は減少するものの、設定や指示が入りやすく、チャット的な会話が成立しやすいモデルとなります。

chatntq + 0.5 * (starling - mistral) + 1.0 * (Antler - StableLM)

コードは以下の通り。

# ChatVectorを適用する対象モデル
cp_model = load_model("chatntq")

# ChatVector1を抽出するベースモデルとinstructionモデル
base_model1 = load_model("mistral")
inst_model1 = load_model("starling")
ratio1 = 0.5  # マージ比率

# ChatVector2を抽出するベースモデルとinstructionモデル(Noneでも可)
base_model2 = load_model("stablelm")
inst_model2 = load_model("antler")
ratio2 = 1.0  # マージ比率

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment