Gemini 1.5 の JSONモードを試す

npaka

2024年5月31日 14:36

Gemini 1.5 のJSONモードを試したので、まとめました。

1. JSONモード

「JSONモード」は、LLMの出力をJSON形式に強制するためのモードです。「gemini-1.5-pro」「gemini-1.5-flash」で利用可能です。

さらに、「gemini-1.5-pro」は「制約付きデコード」をサポートします。これにより、推論時にスキーマオブジェクト (または同等のPython型) を渡すことができ、出力はそのスキーマに厳密に従います。

2. JSONモードの使い方

JSONモードの使い方は、次のとおりです。

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

# パッケージのインストール
!pip install -q -U google-generativeai

(2) 「Google AI Studio」でAPIキーを取得し、シークレットの「GOOGLE_API_KEY」に登録後、以下のセルを実行。

from google.colab import userdata
import google.generativeai as genai

# 環境変数の準備 (左端の鍵アイコンでGOOGLE_API_KEYを設定)
GOOGLE_API_KEY=userdata.get("GOOGLE_API_KEY")
genai.configure(api_key=GOOGLE_API_KEY)

(3) モデルの準備。
generation_config に {"response_mime_type": "application/json"} を指定することで、JSONモードが有効化されます。


# モデルの準備
model = genai.GenerativeModel(
    "gemini-1.5-flash",
    generation_config={"response_mime_type": "application/json"}
)

(4) プロンプトでJSON形式を指定して推論実行

# プロンプトでJSON出力を指示
prompt = """次のJSONスキーマを使用して、
人気のあるクッキーレシピを日本語で5個リストアップしてください。

Recipe = {'recipe_name': str}
Return: list[Recipe]"""

# 推論実行
raw_response = model.generate_content(prompt)
print(raw_response.text)

[
{"recipe_name": "チョコレートチップクッキー"},
{"recipe_name": "オレオクッキー"},
{"recipe_name": "スノーボールクッキー"},
{"recipe_name": "バタークッキー"},
{"recipe_name": "ジンジャーマンクッキー"}
]

(5) 文字列をJSONにパース。

import json

# 文字列をJSONにパース
response = json.loads(raw_response.text)
print(response)

[
{'recipe_name': 'チョコレートチップクッキー'},
{'recipe_name': 'オレオクッキー'},
{'recipe_name': 'スノーボールクッキー'},
{'recipe_name': 'バタークッキー'},
{'recipe_name': 'ジンジャーマンクッキー'}
]

3. 制約付きデコードによるJSON出力

制約付きデコードでJSON出力する手順は、次のとおりです。

(1) スキーマオブジェクトの定義

import typing_extensions as typing

# スキーマオブジェクトの定義
class Recipe(typing.TypedDict):
    recipe_name: str

(2) モデルの準備
「gemini-1.5-pro」のみ利用できます。

# モデルの準備
model = genai.GenerativeModel(
    model_name="models/gemini-1.5-pro"
)

(3) スキーマオブジェクトでJSON形式を指定して推論実行。
generation_config() の response_mime_type に "application/json"、response_schema にスキーマオブジェクトを指定します。

# 制約付きデコードでJSON出力を指示
result = model.generate_content(
    "人気のあるクッキーレシピを日本語で5個リストアップしてください。",
    generation_config=genai.GenerationConfig(
        response_mime_type="application/json",
        response_schema = list[Recipe]
    ),
    request_options={"timeout": 600},  # タイムアウト
)
print(result.text)

[
{"recipe_name": "サクサク美味しい！基本のクッキーレシピ"},
{"recipe_name": "簡単混ぜるだけ！HMで作るチョコチップクッキー"},
{"recipe_name": "ほろほろ食感！スノーボールクッキー"},
{"recipe_name": "可愛い！型抜きクッキー"},
{"recipe_name": "抹茶の風味豊か！抹茶クッキー"}
]