見出し画像

gemma-2-2b-jpn-it を試す

ぬこぬこ

tl;dr

  • Gemma 2 2B の日本語版が発表されたよ

  • Google Colab と macOS で試したよ

  • macOS では transformers と MLX で試したよ

  • 2B にしてはかなり日本語が流暢で、ローカル LLM の進展をすごく感じるよ

Gemma 2 2B JPN とは

Gemma-2-JPN は日本語に特化した Gemma 2 の事後学習モデル。英語と同様の性能。使用例として挙げられているものは、詩の作成や翻訳、コード生成など。モデルの訓練には TPU を用いており、JAX、ML Pathways を使用したよう。モデルの性能評価は GPT-3.5 と LLM as a Judge で比較。一般的なベンチマークというより、日本語で聞いたら日本語で答えるという RLHF をしているらしい（Gemma Developer イベントで話していた記憶より）。

試したくてうずうずしてきたでしょう？Google Colab と macOS にて動作確認しましたので、よろしければ試してみてください。

Google Colab で試す

基本的にサンプルコードをそのまま実行すると動きます。Hugging Face Token の設定だけお忘れ無く。Gist に公開しているので、追実験されたい方は Open in Colab からどうぞ。

サンプルにある「マシーンラーニングについての詩を書いてください。」の出力を下記に記載します。

## マシーンラーニングの詩

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

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

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

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

2B でこのレベルの日本語を出力できるのは本当にすごい。

macOS で試す

サンプルコードには丁寧に device を mps に変えるだけで動くとコメントアウトが残されているので、とりあえずそのまま動かしてみる。

ディレクトリを用意し、uv で環境構築、PyTorch と Hugging Face の Transformers をインストールし、実行。

mkdir playground-gemma-2-2b-jpn-it
cd playground-gemma-2-2b-jpn-it

uv init --python 3.11
touch samplecode.py
vim samplecode.py

uv add torch
uv add transformers

uv run samplecode.py

samplecode.py は下記のように指定してください。2B ですので、ユニファイドメモリが 8B もあれば動くはず。

# samplecode.py

import torch
from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="google/gemma-2-2b-jpn-it",
    model_kwargs={"torch_dtype": torch.bfloat16},
    device="mps",
)

messages = [
    {"role": "user", "content": "マシーンラーニングについての詩を書いてください。"},
]

outputs = pipe(messages, return_full_text=False, max_new_tokens=256)
assistant_response = outputs[0]["generated_text"].strip()
print(assistant_response)

手元の出力を下記に記載します。

データの海に浮かぶ星、
機械の魂、その進化の証。
学習の旅、複雑な道、
未知の世界を探索する。

ニューラルネットワーク、複雑な網、
情報から知識を紡ぎ出す。
パターンを見つけ、予測する力、
未来を解き明かす、機械の眼。

教師あり学習、教師なし学習、
様々な手法、複雑なプロセス。
データの海を深く探る、
機械の知恵、無限の力。

進化の過程、常に変化する、
新しい知識、新しい世界へ。
機械学習、未来を拓く鍵、
未知の領域へ、新たな一歩を踏み出す。

とりあえずサンプルコードはここまで。

Hugging Face Spaces にて Demo を公開しています。もしよろしければお試しください。

MLX でも試す

MLX は macOS 用のフレームワークです。gemma-2-2b-jpn-it を mlx で使えるよう変換し、mlx-community にアップロードしました。

こちらをもとに動かしてみます。uv / pip どちらのコマンドも記載しておきます。

# uv を使う場合
uv run --with mlx-lm mlx_lm.generate --model mlx-community/gemma-2-2b-jpn-it --prompt "こんにちは。" --max-tokens 128 --temp 0.0


# pip を使う場合
pip instal mlx-lm
python -m mlx_lm.generate --model mlx-community/gemma-2-2b-jpn-it --prompt "こんにちは。" --max-tokens 128 --temp 0.0

実行してみましょう。

uv run --with mlx-lm mlx_lm.generate --model mlx-community/gemma-2-2b-jpn-it --prompt "こんにちは。" --max-tokens 128 --temp 0.0
 ~/Downloads/playground-mlx  on master ?11                                   with yutahayashi@MyMacStudio  at 15:34:29
> uv run --with mlx-lm mlx_lm.generate --model mlx-community/gemma-2-2b-jpn-it --prompt "こんにちは。" --max-tokens 128 --temp 0.0


config.json: 100%|█████████████████████████████████████████████████████████████████████| 866/866 [00:00<00:00, 3.22MB/s]
model.safetensors.index.json: 100%|████████████████████████████████████████████████| 21.1k/21.1k [00:00<00:00, 35.5MB/s]
tokenizer_config.json: 100%|███████████████████████████████████████████████████████| 46.9k/46.9k [00:00<00:00, 6.87MB/s]
special_tokens_map.json: 100%|█████████████████████████████████████████████████████████| 555/555 [00:00<00:00, 1.75MB/s]
tokenizer.model: 100%|█████████████████████████████████████████████████████████████| 4.24M/4.24M [00:00<00:00, 4.45MB/s]
tokenizer.json: 100%|██████████████████████████████████████████████████████████████| 17.5M/17.5M [00:01<00:00, 12.2MB/s]
model.safetensors: 100%|███████████████████████████████████████████████████████████| 5.23G/5.23G [04:02<00:00, 21.5MB/s]
Fetching 7 files: 100%|███████████████████████████████████████████████████████████████████| 7/7 [04:03<00:00, 34.78s/it]
==========json: 100%|██████████████████████████████████████████████████████████████| 17.5M/17.5M [00:01<00:00, 12.8MB/s]
Prompt: <bos><start_of_turn>user
こんにちは。<end_of_turn>
<start_of_turn>model

こんにちは！😊

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



<end_of_turn>
==========
Prompt: 12 tokens, 41.454 tokens-per-sec
Generation: 17 tokens, 43.487 tokens-per-sec
Peak memory: 7.359 GB

メモリ 8GB に収まっていますね！吊るしの macOS でも動きそうです。よろしければプロンプトなどを変えて試してみてください。

追記

Hugging Face 公式からデモが出ていました！こちらも是非！


すべて勉強代に充てさせていただきます！アウトプットします！