AKU-d_ms-0.5B-chat-v0.1

Overview

スクラッチで学習したMistralアーキテクチャの0.5BモデルをさらにSFTとDPOで会話タスクにファインチューニングしたモデルです。

開発の経緯や開発時のメモ書きは下記のURLを参照ください。 Zenn_記事

Usage

Requirement

pip install transformers sentencepiece protobuf

Code

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "YukiTomita-CC/AKU-d_ms-0.5B-chat-v0.1"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32, device_map="auto").to("cuda")
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)

messages = [
    {"role": "user", "content": "好きな食べ物は何ですか?"}
]

input_ids = tokenizer.apply_chat_template(
    messages,
    return_tensors="pt",
).to(model.device)
outputs = model.generate(input_ids, max_new_tokens=128, do_sample=True, top_p=0.9, top_k=50, num_return_sequences=5)

for output in outputs:
    print(tokenizer.decode(output[input_ids.shape[-1]:], skip_special_tokens=True))

Notice

• messagesは5ターン以下(len(messages)<=10)を推奨します。
  • 理由としては最大5ターンでしか学習していないためで、それ以上続く場合は最新の5ターンをスライスしてください。
• 学習時にデータを正規化しているため、!と?の半角全角にかなり敏感です。どちらも半角を推奨します。
• system promptには対応していません。また、userとassistantが交互となることを前提としています。