apply_chat_vector.py

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)

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  # マージ比率