tokyotech-llm/Swallow-7b-plus-hfをGoogleColobで動かす

AI

初めに

Swallow-7bよりも精度が高いモデルが出たみたいなので触ってみます

huggingface.co

まどマギプロンプトの場合は、以下のようになります

環境

準備

ライブラリの追加

量子化をするため必要なものも追加します。

!pip install torch
!pip install sentencepiece
!pip install protobuf
!pip install transformers bitsandbytes accelerate

モデルのロード

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "tokyotech-llm/Swallow-7b-plus-hf"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name ,load_in_4bit=True, low_cpu_mem_usage=True, device_map="auto")

実行

サンプルプロンプト

推論

prompt = "東京工業大学の主なキャンパスは、"
input_ids = tokenizer.encode(
    prompt,
    add_special_tokens=False,
    return_tensors="pt"
)
tokens = model.generate(
    input_ids.to(device=model.device),
    max_new_tokens=128,
    temperature=0.99,
    top_p=0.95,
    do_sample=True,
)

out = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(out)

結果

 東京工業大学の主なキャンパスは、 大岡山キャンパスと田町キャンパスの2か所です。ƒ 大岡山キャンパスは最寄駅の「大岡山駅」のすぐ近くにあり、学生食堂や銀行ATM、書店など、 大学施設が充実しています。一方、田町キャンパスはJR山手線・京浜東北線「田町駅」から徒歩8分のところにあります。2011年に新設された田町キャンパスには、研究室や情報処理室のほかに、教室、 セミナー室などがあります。
東京工業大学の評判や口コミ

まどマギプロンプト

推論

prompt = "まどマギで一番可愛いキャラは、"
input_ids = tokenizer.encode(
    prompt,
    add_special_tokens=False,
    return_tensors="pt"
)
tokens = model.generate(
    input_ids.to(device=model.device),
    max_new_tokens=128,
    temperature=0.99,
    top_p=0.95,
    do_sample=True,
)

out = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(out)

結果

まどマギで一番可愛いキャラは、満場一致で
 surely
13: 風吹けば名無し@\(^o^)/ 2017/03/13(月) 14:59:52.35 ID:1B3NVnk6r.net
>>12
ほむらやで
ほむらやで
4: 風吹けば名無し@\(^o^)/ 2017/03/13(月) 14:56:40.16 ID:1B3NVnk6

使用リソース

4bit量子化で実行したときは、以下の通りです