kotememo

【Windows】ローカル環境でgoogleの日本語モデルgemma-2-2b-jpn-itを実行

はじめに

ローカル環境で試してみたかったため、Hugging Faceからgoogleの日本語モデルgemma-2-2b-jpn-itモデルをローカルにダウンロードして実行しました。

CPU実行、GPU実行を試し、最後に対話形式にするためのループ処理にしています。

概要

  1. モデルのダウンロード
  2. CPUでの実行
  3. GPUでの実行
  4. 対話ループの実装

環境

  • Windows 11

  • CPU : 12th Gen Intel(R) Core(TM) i7-12700

  • GPU : NVIDIA GeForce RTX 4070 Ti

  • Python 3.10.9

  • transformers 4.45.1

  • accelerate 0.34.2

  • torch 2.4.1 (GPU実行時はtorch 2.4.0+cu124)

1. モデルのダウンロード

https://huggingface.co/google/gemma-2-2b-jpn-it

リポジトリをクローン、またはFiles and versionsからダウンロード出来ます。

ダウンロードしたものはフォルダにまとめておきます。

main.pyは後述で作成するソースです。

.
│  main.py
│
└─model
    └─google
        └─gemma-2-2b-jpn-it
                config.json
                generation_config.json
                gitattributes
                model-00001-of-00002.safetensors
                model-00002-of-00002.safetensors
                model.safetensors.index.json
                README.md
                special_tokens_map.json
                tokenizer.json
                tokenizer.model
                tokenizer_config.json

2. CPUでの実行

必要なパッケージをインストールします。

pip install transformers
pip install accelerate
pip install --upgrade torch

ここでtorchを更新しているのは、もし標準で古いtorchが入っていると実行時に以下のエラーが発生する可能性があるためです。

module 'torch' has no attribute '_dynamo'

さらに、CPUで実行する場合、modelの指定時にはdevice_mapをcpuへ変更することと、torch_dtypeは削除する必要があります。

モデルはダウンロードしたフォルダへの相対パスにします。

tokenizer = AutoTokenizer.from_pretrained("model/google/gemma-2-2b-jpn-it")
model = AutoModelForCausalLM.from_pretrained(
    "model/google/gemma-2-2b-jpn-it",
#    device_map="auto",
    device_map="cpu",
#    torch_dtype=torch.bfloat16,
)

上記部分以外はほぼreadme通りのソースです。

追加でモデルのロード時間と回答時間を計測しています。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import time

start_time_1 = time.time()

# モデルとトークナイザーのロード
tokenizer = AutoTokenizer.from_pretrained("model/google/gemma-2-2b-jpn-it")
model = AutoModelForCausalLM.from_pretrained(
    "model/google/gemma-2-2b-jpn-it",
#    device_map="auto",
    device_map="cpu",
#    torch_dtype=torch.bfloat16,
)

end_time_1 = time.time()
print("------------------------------")
print("モデルロード時間 ", end_time_1 - start_time_1)
print("------------------------------")

start_time_2 = time.time()

messages = [
    {"role": "user", "content": "マシーンラーニングについての詩を書いてください。"},
]
inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True, return_dict=True).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=256)
generated_text = tokenizer.batch_decode(outputs[:, inputs['input_ids'].shape[1]:], skip_special_tokens=True)[0]
print(generated_text.strip())

end_time_2 = time.time()
print("------------------------------")
print("回答時間 ", end_time_2 - start_time_2)
print("------------------------------")

実行時にtorchの警告文が表示されますが、動作はします。

The 'max_batch_size' argument of HybridCache is deprecated and will be removed in v4.46. Use the more precisely named 'batch_size' argument instead.
Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)

結果としては、モデルのロードは約8秒、回答時間は約70秒のようです。

D:\src\python\chat_with_model_test\main>python .\main.py
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████| 2/2 [00:07<00:00,  3.54s/it]
------------------------------
モデルロード時間  8.160688638687134
------------------------------
The 'max_batch_size' argument of HybridCache is deprecated and will be removed in v4.46. Use the more precisely named 'batch_size' argument instead.
Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)
## マシーンラーニングの詩

**1.**
データの海、深淵の広がり、
複雑なパターン、隠された知識。
機械の目、無数の計算、
未来を予測、新たな道を開く。

**2.**
学習の旅、複雑な過程、
教師あり、教師なし、様々な方法。
アルゴリズムの力、複雑なネットワーク、
未知の世界を解き明かす。

**3.**
画像認識、音声認識、
予測、分類、様々なタスク。
機械の知恵、人間の夢を形作る、
新しい時代を築く力を持つ。

**4.**
倫理の問いかけ、責任の重さ、
人間と機械の共存、未来の道。
技術の進化、未知の領域へ、
新たな可能性を創造する。
------------------------------
回答時間  72.89190626144409
------------------------------

3. GPUでの実行

PyTorchを自身の環境に合わせてインストールします。

# CUDA 12.4
pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu124

ソースについてはdevice_maptorch_dtypeを元に戻します。

# モデルとトークナイザーのロード
tokenizer = AutoTokenizer.from_pretrained("model/google/gemma-2-2b-jpn-it")
model = AutoModelForCausalLM.from_pretrained(
    "model/google/gemma-2-2b-jpn-it",
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

結果としては、モデルのロードは約5秒、回答時間は約9であり、実行速度が大幅に早くなったことが確認できます。

D:\src\python\chat_with_model_test\main>python .\main.py
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████| 2/2 [00:03<00:00,  1.89s/it]
------------------------------
モデルロード時間  5.084553956985474
------------------------------
The 'max_batch_size' argument of HybridCache is deprecated and will be removed in v4.46. Use the more precisely named 'batch_size' argument instead.
Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)
## マシーンラーニングの詩

**1.**
データの海、深淵の広がり、
複雑なパターン、隠された知識。
機械学習、その力強さ、
未来を予測、その道を開く。

**2.**
ニューラルネットワーク、複雑な枝、
学習の旅、その過程は静か。
データから学び、進化する姿、
予測の精度、その力強さ。

**3.**
教師あり学習、正解を導く、
教師なし学習、未知の世界へ。
機械学習、その進化は止まらない、
未来の扉を開く、新たな時代へ。

**4.**
画像認識、音声認識、
複雑なタスク、その答えを見つける。
機械学習、その力強さ、
未来の技術、その可能性を語る。
------------------------------
回答時間  9.14939260482788
------------------------------

4. 対話ループの実装

上記ソースではプログラム内で指定した内容1つへの回答のみしか対応できません。

そこで、コマンドプロンプトで入力した内容を対話として保持しつつループする処理として実装します。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# モデルをロード
tokenizer = AutoTokenizer.from_pretrained("model/google/gemma-2-2b-jpn-it")
model = AutoModelForCausalLM.from_pretrained(
    "model/google/gemma-2-2b-jpn-it",
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

# 対話の履歴を保持するリスト
messages = []

def chat_with_model(user_input):
    # ユーザーからの入力をメッセージに追加
    messages.append({"role": "user", "content": user_input})
    
    # モデルに渡すための入力を準備
    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True, return_dict=True).to(model.device)

    # モデルによる応答生成
    outputs = model.generate(**inputs, max_new_tokens=1024)
    generated_text = tokenizer.batch_decode(outputs[:, inputs['input_ids'].shape[1]:], skip_special_tokens=True)[0]

    # システムの応答をメッセージに追加
    messages.append({"role": "assistant", "content": generated_text.strip()})

    return generated_text.strip()

# 対話の開始
while True:
    # 質問文を入力
    user_input = input("あなた: ")

    # 対話終了
    if user_input.lower() in ["終了", "exit", "quit"]:
        print("対話を終了します。")
        break

    # 対話履歴リセット
    if user_input.lower() in ["リセット", "reset"]:
        messages = []
        print("アシスタント: 対話履歴をリセットしました。")
        continue

    # 応答文を出力
    response = chat_with_model(user_input)
    print("\n")
    print("アシスタント:", response)
    print("\n")

「終了」「exit」「quit」を入力した場合はループを終了し、「リセット」「reset」を入力した場合は対話内容を削除して初期状態にします。

次のような動作になります。

あなた: 挨拶をしてください。

アシスタント: こんにちは! 😊

何かお手伝いできることはありますか?


あなた: その言葉を英語にしてください。


アシスタント: Here's the translation:

"Hello! 😊  Is there anything I can help you with?"



Let me know if you'd like to try another translation!


あなた: リセット
アシスタント: 対話履歴をリセットしました。
あなた: その言葉を英語にしてください。


アシスタント: すみません、どの言葉を英語にしたいのか教えてください! 😊


あなた: 終了
対話を終了します。

簡単にですが、文脈を覚えて英語変換していることと、対話内容をリセットできていることが確認できます。