Snowflake Arcticに入力したプロンプトのトークン数を確認する

Snowflake Arcticとは

Snowflake社が開発したオープンなLLMです。
トレーニングや推論のコストを低減する様々な技術が導入されていてコスト効率がいいことが特長です。
また、SQL文の生成やプログラムのコーディングが得意で、プロンプトによる指示を遵守する能力が高いとされています。

Snowflake Cortexから最小限の引数で呼び出してみる

Snowflake Cortexという機能群を用いることでLLMを簡単に呼び出すことができます。
例えば以下のようなプログラムからSnowflake Arkticを呼び出すことができるのですが、最小限の引数で呼び出すと単純に文字列のみが返ってきます。

arctic_simple.py

response = snowflake.cortex.Complete("snowflake-arctic", user_prompt)
print(type(response))
print('---')
print(response)
print('---')

以下は実際の出力です。

<class 'str'>
---
このプロンプトは、"このプロンプトは何トークンですか？" という文章で構成されています。この文章は、日本語で書かれており、トークンという単語が含まれて います。トークンは、自然言語処理や暗号技術などの分野で、文章やデータを解析する際に用いられる単位です。

この文章に含まれるトークンを分析すると、以下のようになります。

1. この: これは、日本語の指示詞で、「これ」という意味があります。
2. プロンプト: これは、英語の単語で、「提案、案、提示」という意味があります。日本語でも、プロンプトという言葉が用いられることがあります。
3. は: これは、日本語の助詞で、「は」という意味があります。主に、話題の説明や対比をする際に用いられます。
4. 何: これは、日本語の単語で、「何」という意味があります。「何」は、不定冠詞として、名詞の前に置かれることがあります。
5. トークン: これは、英語の単語で、「記号、単位」という意味があります。日本語でも、トークンという言葉が用いられることがあります。
6. です: これは、日本語の助動詞で、「です」という意味があります。主に、文の終わりに用いられ、文の整合性を保ちます。
7. か: これは、日本語の助詞で、「か」という意味があります。主に、疑問文や条件文に用いられます。

以上のように、このプロンプトは7つのトークンで構成されています。
---

あるいはSnowflakeのWeb UIであるSnowsightからSQLでSnowflake Arcticを呼び出してみても同様に文字列だけが返ってきます。

トークン数を確認する

options引数を指定することでトークン数を含む詳細な返り値を取得することができます。

When the options argument is given, a string representation of a JSON object containing the following keys.

Pythonであれば、例えば以下のようなプログラムで確認することができます。

arctic_metadata.py

# 環境構築コマンド
# conda create -n arctic python=3.8
# conda activate arctic
# pip install snowflake-connector-python snowflake-snowpark-python snowflake-ml-python[all] streamlit tiktoken

# 実行コマンド
# streamlit run arctic_st.py

import streamlit as st
import snowflake.connector
import snowflake.snowpark as snowpark
import snowflake.cortex as cortex
import json
import tiktoken

# ローカルPython環境からSnowflakeに接続するための関数
@st.cache_resource(ttl=7200)
def connect_snowflake():
    # Snowflakeに接続する
    # 注意！本来はここに指定するパラメータをハードコーディングしてはいけない
    connection = snowflake.connector.connect(
        user='sakatoku',
        password='ここにパスワードを記載',
        account='XXXXXXX-TEST_OREGON (アカウント識別子のドット.をハイフン-に変えたものを記載)',
        role='cortex_user_role',
        warehouse='compute_wh')

    # Snowparkセッションを作成する
    session = snowpark.Session.builder.configs({"connection": connection}).create()
    return session 

# 画面レイアウトを設定
st.set_page_config(layout="wide")

# デフォルトのプロンプト
prompt = '''You are an expert planner.
Your task is to consider the combination of snacks.
The combination shall be output includes an array of [ID, quantity] and total price.
The total price must exceed 400 yen and must not exceed 500 yen.
You should output the combination by JSON.
You should be concerned with variety of snacks in combination.
Information of snacks is below, that format is (ID), (name), (price):
- ID1, Umai-bo, 10 yen
- ID2, Coron, 20 yen
- ID3, Dorayaki chocolate, 15 yen
- ID4, Pockey, 100 yen
- ID5, Sio senbei, 5 yen
- ID6, Potato chips, 50 yen
- ID7, Banana chocolate, 30 yen
- ID8, Cider tablet, 10 yen
- ID9, Fruit gumi, 80 yen
- ID10, Salami, 50 yen

So, let's start planning!'''

# TikTokenでトークン数を見積もる。Snowflake Arcticのトークナイザではないので参考値
def estimate_token_count(prompt):
    encoding1 = tiktoken.encoding_for_model("gpt-4")
    encoding2 = tiktoken.encoding_for_model("gpt-3.5-turbo")
    tokens1 = encoding1.encode(prompt)
    tokens2 = encoding2.encode(prompt)
    return len(tokens1), len(tokens2)

# ユーザプロンプトを取得する。ユーザプロンプトが空の場合はここで終了
user_prompt = st.text_area("プロンプト", value=prompt)
if not user_prompt:
    st.stop()

# 推定トークン数を表示する
token_count1, token_count2 = estimate_token_count(user_prompt)
st.write(f'推定トークン数(GPT-4): {token_count1}')
st.write(f'推定トークン数(GPT-3.5): {token_count2}')

# ユーザプロンプトを表示
with st.chat_message("user"):
    st.write(user_prompt)

# Snowflakeに接続する
session = connect_snowflake()

# SQLでSnowflake Cortex LLM Functionsを呼び出す
def cortex_complete(session, model_name, prompt):
    # 注意！本来はSQLインジェクション攻撃ができないように適切に文字列を処理することが望ましい
    # プロンプトをエスケープする
    prompt = prompt.replace('\'', '\\\'')
    # JSONで入力すると出力のときにメタデータが得られる
    # cmd = f"select snowflake.cortex.complete('{model_name}', [{{'role': 'user', 'content': '{prompt}'}}], {{'temperature': 0.1}}) as RESPONSE"
    cmd = f"select snowflake.cortex.complete('{model_name}', '{prompt}') as RESPONSE"
    # session.sqlの出力はROWのリストとして得られる
    # ROWの"RESPONSE"キーにJSON文字列が格納されている。キーは必ずアッパーケースになることに注意
    df_response = session.sql(cmd).collect()
    print(df_response)
    json_response = json.loads(df_response[0]["RESPONSE"])
    return json_response

# Snowflake Cortex LLM FunctionsのJSONからレスポンスだけを抽出する
def parse_llm_response(j):
    return j["choices"][0]["messages"]

# Snowflake Cortex LLM Functionsを呼び出してレスポンスを取得する
response = cortex_complete(session, "snowflake-arctic", user_prompt)

# レスポンスを表示
with st.chat_message("arctic"):
    st.write(parse_llm_response(response))
    with st.expander("詳細", expanded=False):
        st.write(response)

実行結果は以下のようになりました。

今回入力したプロンプトのトークン数は272トークンのようです。
トークナイザが違うから当たり前なのですが、TikTokenで推定したトークン数とはかなり差がありました。
ちょっとした用途であればとりあえず実行してトークン数を計測する、でもあまり困らないとは思うのですが…トークン数を推定するいい方法をご存じの方がいたら教えてください。