❄️

Streamlit in Snowflake (SiS) でトークン数チェックアプリを作ってみた

2024/07/09に公開

はじめに

はじめまして。Snowflake でセールスエンジニアをしています。日頃試してみたことを皆様にお裾分けできればと思い色々と投稿していければと思います。今回は Streamlit in Snowflake で Cortex LLM のトークン数をチェックしてコスト試算ができるアプリを作ってみたいと思います。

※ 本記事は私個人の発信となります。Snowflake を代表しての発言では無いためご理解いただければ幸いです。

Streamlit in Snowflake (SiS) とは?

まず Stremlit は Python のライブラリです。Streamlit を使うとシンプルな Python の記述だけで Web UI を作成することができます。本来 Web UI は HTML / CSS / JavaScript などを用いて記述する必要があるのですが、Streamlit がそれらの機能を代替してくれるイメージとなります。

こちらの App Gallery にユーザーが投稿したサンプルがあるためイメージしやすいです。

Streamlit in Snowflake は Snowflake 上で Streamlit の Web アプリを開発したり実行することが可能です。Snowflake のアカウントさえあれば触れる手軽さも魅力ですが、何より Snowflake 上のテーブルデータなどを簡単に Web アプリに組み込める点が良いところだと思います。

Streamlit in Snowflake について (Snowflake 公式ドキュメント)

Snowflake Cortex とは?

Snowflake Cortex とは Snowflake の生成 AI の機能群です。その中の Cortex LLM は Snowflake 上で稼働する大規模言語モデルを呼び出せる機能であり、SQL や Python などからシンプルな関数で色んな種類の LLM を利用することが可能となります。

Large Language Model (LLM) Functions (Snowflake Cortex) (Snowflake 公式ドキュメント)

機能概要

イメージ図

芥川龍之介先生の『蜘蛛の糸』から文章を一部お借りしました。

機能一覧

  • ユーザーが Cortex LLM のモデルを選ぶ
  • ユーザーが入力した文章の文字数やトークン数を表示
    • (LLM はモデル毎に搭載するトークナイザーが異なるため、同じ文字数でもトークン数が異なります)
  • トークン数と文字数の比率を表示
  • Snowflake のクレジット単価を掛けて概算費用を算出
    • (Cortex LLM は Snowflake のクレジット数ベースでの単価設定となっているため)

Cortex LLM の価格表 (PDF)

前提条件

  • Cortex LLM が利用できる Snowflake アカウント
  • snowflake-ml-python 1.1.2 以降

Cortex LLM のリージョン対応表 (Snowflake 公式ドキュメント)

ソースコード

import streamlit as st
from snowflake.snowpark.context import get_active_session
import snowflake.snowpark.functions as F

# 現在のセッションの取得
session = get_active_session()

# アプリケーションのタイトル
st.title("Cortex AI トークン数チェックアプリ")

# AI設定
st.sidebar.title("AI の設定")
lang_model = st.sidebar.radio("使用したい言語モデルを選んでください",
                              ("snowflake-arctic", "reka-flash", "reka-core", "mistral-large", "mixtral-8x7b", 
                               "llama3-70b", "llama2-70b-chat", "llama3-8b", "mistral-7b", "gemma-7b"))

# トークン数を計算する関数 (Cortex用のトークンカウント関数が用意されている)
def count_tokens(model, text):
    result = session.sql(f"SELECT SNOWFLAKE.CORTEX.COUNT_TOKENS('{model}', '{text}') as token_count").collect()
    return result[0]['TOKEN_COUNT']

# トークン数チェックと料金計算
st.header("トークン数チェックと料金計算")

input_text = st.text_area("左ペインから言語モデルを選択して、トークン数をチェックしたいテキストを入力してください:", height=200)

# ユーザーに1クレジットあたりの単価を入力させる
credit_price = st.number_input("Snowflakeの1クレジットあたりの単価 (ドル) を入力してください:", min_value=0.0, value=2.0, step=0.01)

# モデルごとの1Mトークンあたりのクレジット数 (2024/7/9時点では arctic は無料)
model_credits = {
    "snowflake-arctic": 0,
    "reka-flash": 0.45,
    "reka-core": 5.5,
    "mistral-large": 5.1,
    "mixtral-8x7b": 0.22,
    "llama3-70b": 1.21,
    "llama2-70b-chat": 0.45,
    "llama3-8b": 0.19,
    "mistral-7b": 0.12,
    "gemma-7b": 0.12
}

if st.button("トークン数を計算"):
    if input_text:
        # 文字数の計算
        char_count = len(input_text)
        st.write(f"入力テキストの文字数: {char_count}")

        if lang_model in model_credits:
            # トークン数の計算
            token_count = count_tokens(lang_model, input_text)
            st.write(f"入力テキストのトークン数: {token_count}")
            
            # 文字数とトークン数の比率
            ratio = token_count / char_count if char_count > 0 else 0
            st.write(f"トークン数 / 文字数の比率: {ratio:.2f}")
            
            # 料金計算
            credits_used = (token_count / 1000000) * model_credits[lang_model]
            cost = credits_used * credit_price
            
            st.write(f"使用クレジット数: {credits_used:.6f}")
            st.write(f"概算料金: ${cost:.6f}")
        else:
            st.warning("選択されたモデルはSnowflakeのトークン数計算機能に対応していません。")
    else:
        st.warning("テキストを入力してください。")

最後に

特に日本語を利用していると文字数とトークン数にギャップがあり見積もりが難しくなりがちです。本アプリをいつでも使えるようにすることで、サッと LLM ワークロードのコスト試算がしやすくなるのではないでしょうか?皆様のお役に立てば幸いです。

宣伝

X で Snowflake の What's New Bot 始めてみました。

Discussion