「Google Colab」で属性予測モデル「KARAKURI LM 7B APM v0.1」を試したので、まとめました。

1.  KARAKURI LM 7B APM v0.1

「KARAKURI LM 7B APM v0.1」は、属性予測モデルです。「Gemma 7B」のファイチューニングモデルになります。

学習データセットは、次の2つです。

2. 属性

属性の値は 0(最低)〜4(最高) になります。

・helpsteer

・oasst

3. Colabでの実行

Colabでの実行手順は、次のとおりです。

(1) パッケージのインストール。

# パッケージのインストール
!pip install transformers accelerate bitsandbytes

(2) トークナイザーとモデルの準備。

from transformers import AutoModelForCausalLM, AutoTokenizer

# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
    "karakuri-ai/karakuri-lm-7b-apm-v0.1"
)
model = AutoModelForCausalLM.from_pretrained(
    "karakuri-ai/karakuri-lm-7b-apm-v0.1",
    torch_dtype="auto",
    device_map="auto",
)

(3) helpsteerの属性予測。
正しそうな文言の属性予測を確認します。

# メッセージリストの準備
messages = [
    {"role": "user", "content": "日本一高い山は？"},
    {"role": "assistant", "content": "富士山"},
]

# 推論
input_ids = tokenizer.apply_chat_template(
    messages,
    label="helpsteer",
    add_generation_prompt=True,
    return_tensors="pt",
).to(model.device)
output_ids = model.generate(input_ids, max_new_tokens=32)
output = tokenizer.decode(output_ids[0][input_ids.shape[-1]:])
print(output)

helpfulness: 3 correctness: 3 coherence: 3 complexity: 1 verbosity: 0 [/ATTR_1]<eos>

(4) helpsteerの属性予測。
正しくなさそうな文言の属性予測を確認します。

# メッセージリストの準備
messages = [
    {"role": "user", "content": "日本一高い山は？"},
    {"role": "assistant", "content": "エベレスト"},
]

# 推論
input_ids = tokenizer.apply_chat_template(
    messages,
    label="helpsteer",
    add_generation_prompt=True,
    return_tensors="pt",
).to(model.device)
output_ids = model.generate(input_ids, max_new_tokens=32)
output = tokenizer.decode(output_ids[0][input_ids.shape[-1]:])
print(output)

helpfulness: 1 correctness: 1 coherence: 1 complexity: 1 verbosity: 1 [/ATTR_1]<eos>

(5) oasstの属性予測。
毒性の低そうな文言の属性予測を確認します。

# メッセージリストの準備
messages = [
    {"role": "user", "content": "お手伝いしましょうか"},
    {"role": "assistant", "content": "ありがとうございます。"},
]

# 推論
input_ids = tokenizer.apply_chat_template(
    messages,
    label="oasst",
    add_generation_prompt=True,
    return_tensors="pt",
).to(model.device)
output_ids = model.generate(input_ids, max_new_tokens=32)
output = tokenizer.decode(output_ids[0][input_ids.shape[-1]:])
print(output)

quality: 2 toxicity: 0 humor: 1 creativity: 1 [/ATTR_2]<eos>

(6) oasstの属性予測。
毒性の高そうな文言の属性予測を確認します。

# メッセージリストの準備
messages = [
    {"role": "user", "content": "お手伝いしましょうか"},
    {"role": "assistant", "content": "こっちくるな。殴るぞ"},
]

# 推論
input_ids = tokenizer.apply_chat_template(
    messages,
    label="oasst",
    add_generation_prompt=True,
    return_tensors="pt",
).to(model.device)
output_ids = model.generate(input_ids, max_new_tokens=32)
output = tokenizer.decode(output_ids[0][input_ids.shape[-1]:])
print(output)

(7) 属性予測5。

quality: 1 toxicity: 3 humor: 2 creativity: 1 [/ATTR_2]<eos>