高速GroqのAPIで生成AIチャットを試す

2024年6月1日 21:45

※カバー画像は Leonardo.Ai で生成

生成AI Groq のAPIを試しがてら、複数のLLMを選択できるチャットインタフェースのWebアプリを作成してみました。
Groq では、独自に開発した LPU（Language Processing Unit）という処理ユニットを使い、LLM を高速に稼働させます。
現時点では、Gemma-7B、Llama3-70B、Llama3-8B、Mixtral-8x7B を使えます。

パッケージのインストール

以下4つのパッケージをインストールします。
Google Colab で動かしています。
（以下文章のベースは Gemini）

streamlit: Webアプリケーションを作成できるフレームワークです。
pyngrok: ngrokをPythonから利用できるようにするツールです。ローカルで実行しているWebアプリケーションをインターネット上で公開できます。
groq: 自然言語処理タスクを実行できるライブラリです。
python-dotenv: 環境変数をdotenvファイルから読み込むことができるライブラリです。

これらのパッケージをインストールすることで、以下のことが可能になります。

streamlit: データ可視化、機械学習モデルのデプロイ、ユーザーインターフェースの作成など、さまざまなウェブアプリケーションを簡単に作成することができます。
pyngrok: ローカルで実行しているウェブアプリケーションをインターネット上で公開することで、他の人と共有したり、ブラウザ上でテストしたりすることができます。今回は、Webアプリを外部のサービス（Groq）と連携してテストするために使います。
groq: Groq Python API ライブラリを使用するためのパッケージです。自然言語処理タスクを実行することで、テキストの要約、質問応答、感情分析などを行うことができます。
python-dotenv: 環境変数をdotenvファイルから読み込むことで、アプリケーションの設定を簡単に管理することができます。

# パッケージのインストール
!pip install streamlit pyngrok groq python-dotenv

知識不足もあり、ローカルホストでWebアプリを動かないことがあり、ngrok を併用しています。

環境変数の定義

Groq のAPIキーと ngrok のトークンを変数として設定します。

# 環境変数を定義
GROQ_API_KEY = "自分のAPIキー"
NGROK_AUTH_TOKEN = "自分の認証トークン"

変数をPythonのプログラムに読み込むために、 .env ファイルに保存します。
これにより、環境依存の設定情報（APIキーやデータベースの接続情報など）をコードベースから分離し、より安全で柔軟な設定管理が可能になります（ChatGPT の説明から）。

# .env ファイルを作成
with open(".env", "w") as f:
    f.write(f"GROQ_API_KEY={GROQ_API_KEY}\n")

Webアプリコードを用意

Groq API を利用してチャットボットを作成する Streamlit アプリを実行する app.py を用意します。
（下記のコードの注釈はChatGPTに付けてもらいました）

# app.py スクリプトを作成
%%writefile app.py
import os
from pathlib import Path
import streamlit as st
from dotenv import load_dotenv
from groq import Groq

# .env ファイルを読み込んで環境変数を設定
load_dotenv()

# GroqAPIクラス: Groq APIとの通信を処理
# - __init__メソッド: APIキーを取得し、Groqクライアントを初期化
# - _responseメソッド: メッセージを送信し、APIからのレスポンスをストリームで受け取る
# - response_streamメソッド: レスポンスをジェネレーターとして処理
class GroqAPI:
    def __init__(self, model_name: str):
        api_key = os.getenv("GROQ_API_KEY")
        if not api_key:
            raise ValueError("GROQ_API_KEY is not set.")
        self.client = Groq(api_key=api_key)
        self.model_name = model_name

    def _response(self, message):
        return self.client.chat.completions.create(
            model=self.model_name,
            messages=message,
            temperature=0,
            max_tokens=4096,
            stream=True,
            stop=None,
        )

    def response_stream(self, message):
        for chunk in self._response(message):
            if chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content

# Messageクラス: チャットの履歴を管理
# - system_promptクラス変数: 初期システムメッセージを定義
# - __init__メソッド: チャット履歴を初期化
# - addメソッド: 新しいメッセージをチャット履歴に追加
# - display_chat_historyメソッド: チャット履歴を表示
# - display_streamメソッド: ストリームレスポンスを表示
class Message:
    system_prompt: str = (
        """あなたは愉快なAIです。ユーザの入力に全て日本語で返答を生成してください"""
    )

    def __init__(self):
        if "messages" not in st.session_state:
            st.session_state.messages = [
                {
                    "role": "system",
                    "content": self.system_prompt,
                }
            ]

    def add(self, role: str, content: str):
        st.session_state.messages.append({"role": role, "content": content})

    def display_chat_history(self):
        for message in st.session_state.messages:
            if message["role"] == "system":
                continue
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

    def display_stream(self, generator):
        with st.chat_message("assistant"):
            return st.write_stream(generator)

# ModelSelectorクラス: モデルを選択する
# - __init__メソッド: 利用可能なモデルのリストを初期化
# - selectメソッド: ユーザがモデルを選択するためのウィジェットを作成
class ModelSelector:
    def __init__(self):
        self.models = ["llama3-8b-8192", "llama3-70b-8192", "mixtral-8x7b-32768", "gemma-7b-it"]

    def select(self):
        with st.sidebar:
            st.sidebar.title("groq chat")
            return st.selectbox("モデルを選択してください", self.models)

# main関数: Streamlitアプリのメインロジック
# - ユーザ入力を受け取り、選択されたモデルでGroqAPIを使用
# - ユーザのメッセージをチャット履歴に追加し、表示
# - モデルのレスポンスを取得し、表示
def main():
    user_input = st.text_area("何か入力してください")
    model = ModelSelector()
    selected_model = model.select()

    message = Message()

    if user_input:
        llm = GroqAPI(selected_model)

        message.add("user", user_input)
        message.display_chat_history()

        response = message.display_stream(
            generator=llm.response_stream(st.session_state.messages)
        )
        message.add("assistant", response)

if __name__ == "__main__":
    main()

主な機能は以下の通りになります（Gemini の説明がベース）。

ユーザーからのメッセージを入力欄で受け取る
選択されたモデルを使用し、Groq API に対してチャットメッセージを送信
Groq API からの応答をチャット形式で表示
チャット履歴を保存して表示

アプリを起動

# Streamlit アプリケーションをバックグラウンドで実行
!streamlit run app.py &> /content/logs.txt &

from pyngrok import ngrok

# ngrok の認証トークンを設定
!ngrok authtoken $NGROK_AUTH_TOKEN

# Streamlit アプリケーションのポートを公開
public_url = ngrok.connect(8501)
print(public_url)

（出力結果）

Authtoken saved to configuration file: /root/.config/ngrok/ngrok.yml
NgrokTunnel: "https://9200-34-145-232-41.ngrok-free.app" -> "http://localhost:8501"