Google Colab で xGen を試す

npaka

2023年7月2日 09:37

「Google Colab」で「xGen」を試したので、まとめました。

1. xGen

「xGen」は、「Salesforce」が開発した、8Kの文脈をあつかえる7BのローカルLLMです。長い文脈をあつかう対話、質問応答、要約が同サイズのモデルと比べ大きく改善されているとのことです。

2. モデル一覧

「xGen」は、次の3つのモデルが提供されています。

・Salesforce/xgen-7b-8k-inst
・Salesforce/xgen-7b-8k-base
・Salesforce/xgen-7b-4k-base

3. Colabでの実行

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

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

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

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

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# トークナイザーとモデルの読み込み
tokenizer = AutoTokenizer.from_pretrained(
    "Salesforce/xgen-7b-8k-inst",
    trust_remote_code=True
)
model = AutoModelForCausalLM.from_pretrained(
    "Salesforce/xgen-7b-8k-inst",
    torch_dtype=torch.bfloat16,
).to("cuda")

(3) 推論の実行。
日本語も少しできますが、英語の方が精度が高いです。

# プロンプトの準備
article = "後藤ひとりはギターを愛する孤独な少女。 家では孤独でただ遊んでばかりの毎日だったが、ひょんなことから伊地知虹夏率いる「結束バンド」に加入することに。 人前で演奏することに不慣れな後藤は、立派なバンドマンになれるのか？"
prompt = f"### Human: 後藤ひとりが加入するバンドの名前は何ですか？\n\n{article}\n### Assistant:"

# 推論の実行
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
sample = model.generate(**inputs, max_new_tokens=128)
print(tokenizer.decode(sample[0]))

### Human: 後藤ひとりが加入するバンドの名前は何ですか？

後藤ひとりはギターを愛する孤独な少女。 家では孤独でただ遊んでばかりの毎日だったが、ひょんなことから伊地知虹夏率いる「結束バンド」に加入することに。 人前で演奏することに不慣れな後藤は、立派なバンドマンになれるのか？
### Assistant: 後藤ひとりは「結束バンド」に加入することになります。
<|endoftext|>