Vertex AIのgeminiを動かすgradioのサンプルコード

みなさんGemini使ってますか？

入力トークンが100万トークンを入れられるGemini-1.5-flashやGemini-1.5-pro、しかもGPT-4に比べて安価に使えるということで最近Googleの渾身のLLM・Geminiを重宝しています。GPT-3.5をファインチューニングするよりも入力トークンにデータを乗っけてしまった方が精度が良くなったという経験もあり、入力トークンが大きいモデルを最近は使っています。

GeminiをVertex AIで使う場合は、ユーザーの入力を学習データには利用しない旨が書かれているので、安心して利用できます。

gradioとは

機械学習モデルのアプリを公開・シェアするのに便利なwebインターフェースを作成します。デモアプリを作るのに便利です。ChatInterfaceというものがあり、数行のコードで、チャット画面（Retry/Undo/Clearボタン付き）を立ち上げることができます。

今回紹介するプログラムの特徴

ストリーミング対応・会話履歴ありというプログラムになっています。ストリーミング対応なので、長い出力を待たなくて良いし、会話履歴を保持しているため、「続きを生成」することも可能です。手元でパッとGeminiのテキスト生成を試したい時に便利です。

コードを動かす前に実行するコマンド

GCPのプロジェクトを先に作成して、課金方法を設定しておいてください。
Vertex AIはAPIの認証にGoogle Cloudサービスアカウントを使用するので、以下のコマンドを実行して認証を通します。

pip install gradio
pip install --upgrade google-cloud-aiplatform
gcloud auth application-default login

コードの紹介

認証が終わったら以下のコードをコピペして実行してください。

import gradio as gr
import vertexai
from vertexai.generative_models import GenerativeModel, Part, Content
import vertexai.preview.generative_models as generative_models

# モデルを定義する関数
def create_vertex_ai_client(model_name):
    vertexai.init(project="<プロジェクトIDを入力>", location="asia-northeast1")
    system_prompt = "あなたはアシスタントAIです。"
    generation_config = {
        "max_output_tokens": 8192,
        "temperature": 1.0,
        "top_p": 0.95,
        "top_k": 20,
    }

    safety_settings = {
        generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_NONE,
        generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_NONE,
        generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_NONE,
        generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_NONE,
        generative_models.HarmCategory.HARM_CATEGORY_UNSPECIFIED: generative_models.HarmBlockThreshold.BLOCK_NONE,
    }

    model = GenerativeModel(
        model_name=model_name,
        system_instruction=[system_prompt],
        generation_config=generation_config,
        safety_settings=safety_settings,
    )

    return model

# チャットの応答を取得する関数、会話の履歴も送信する
def get_chat_response_vertex_ai(model, input, history):
    gemini_history = []
    for row in history:
        input_from_user = Part.from_text(row[0])
        input_from_model = Part.from_text(row[1])

        gemini_history.append(Content(role="user", parts=[input_from_user]))
        gemini_history.append(Content(role="model", parts=[input_from_model]))
    
    chat = model.start_chat(history=gemini_history)
    try:
        responses = chat.send_message(input ,stream=True)
        for response in responses:
            yield response.text
    except Exception as e:
        print(f"Error: {e}")
        yield ""

# チャットの応答をストリーミングする関数
def stream_response(input, history, model_name):
    if input is not None:
        partial_message = ""
        model = create_vertex_ai_client(model_name)
        for response in get_chat_response_vertex_ai(model, input, history):
            partial_message += response
            yield partial_message 

css ="""
.contain { display: flex; flex-direction: column; }
.gradio-container { height: 100vh !important; }
#component-0 { height: 100%; }
#chatbot { flex-grow: 1; overflow: auto;}
"""

# チャットインターフェースを作成
with gr.Blocks(css=css) as demo:
    chatbot = gr.ChatInterface(
        stream_response,
        additional_inputs=[
            gr.Radio(["gemini-1.5-flash-preview-0514", "gemini-1.5-pro-preview-0514"], value="gemini-1.5-pro-preview-0514" ,label="model_name", info="モデルを選択"),
        ],
    )
    
demo.queue().launch(debug=True)

実行結果

このようにモデルを選択するラジオボタンと、チャットインターフェースが表示されます。
会話履歴を送信するのでしりとりも可能です。

Geminiの微妙なところ

実際に長文の生成で使っているのですが、セーフティーフィルターが邪魔をする時があります。

こちらのようになんでもない文字列なのにフィルターにかかってしまって、テキストの生成がストップすることに悩まされています。早く解決してほしいところです。

まとめ

以上、Gemini-1.5-flashとproを動かすgradioサンプルコードの紹介でした。
このようにAIを使った記事を月1回は書いていますので、フォロー・いいねよろしくお願いいたします。

過去記事