bigquery の text embedding モデルでカテゴリ分類してみる

2024/04/12engineering bigquery embedding bq-ml

なにやらbigqueryのクエリからVertex AI テキストエンベディングの基盤モデルを利用できるとのことで実験してみた

準備

GCP セットアップ

以下記事を参考に、llm の外部接続、iam の作成

https://zenn.dev/cloud_ace/articles/bigquery-llm-rag#事前準備

の権限設定の項まで行う

データセット

iowa_liquor_sales を用いる

bigquery-public-dataにあるデータで、データセットのdescriptionによると、アイオワ州にて小売店にて個人に販売するために購入されたアルコールの卸売注文履歴データだそう

こちらのデータにカテゴリ名と、アイテムの説明のカラムがいるのでそれらを用いてみた

便宜上、カテゴリ名をマスタデータ、アイテムの説明を入力データとよぶことにする

カテゴリの推測

埋め込みモデルを使って入力データとマスタデータの類似度から、カテゴリを推定する

流れとしては、入力データとマスタデータの埋め込みベクトルを取得。入力データに対して、各マスタデータとの類似度を計算する。類似度の高いものを推測されたカテゴリとする。推測されたカテゴリが元レコードのカテゴリと一致するかどうか比較を行い、評価する

必要なデータ格納先のデータセット名は public_test とする

クエリ

すべてbigquery上で完結する

モデル作成

CREATE OR REPLACE MODEL `public_test.remote_llm_embedding`
REMOTE WITH CONNECTION `us.llm_connection`
   OPTIONS (ENDPOINT = 'textembedding-gecko@003' -- 'textembedding-gecko-multilingual' -- 英語以外の場合 
 );

マスタデータの埋め込みベクトル作成

まず、カテゴリのみ埋め込みベクトルのテーブルを作成しておく

-- カテゴリ一覧テーブル
create or replace table `public_test.lowa_liquor_sales_categories`
as (
  SELECT category_name, count(1) as c FROM `bigquery-public-data.iowa_liquor_sales.sales`
  WHERE category_name IS NOT NULL -- embedding が取得できない null を除外
  GROUP BY ALL
)
-- カテゴリ数は 103 あった

-- カテゴリの埋め込みベクトルテーブル作成 
CREATE OR REPLACE TABLE
  `public_test.embedded_category` AS (
SELECT
  *
FROM
  ML.GENERATE_TEXT_EMBEDDING(
    MODEL `public_test.remote_llm_embedding`,
    (
      SELECT
        category_name as content
      FROM
        `public_test.lowa_liquor_sales_categories`
    ),
    STRUCT(TRUE AS flatten_json_output)
  )
  );

入力データの埋め込みベクトル取得、カテゴリの推定

元のテーブルは、28792378 レコード。三千万件、あったので、distinct を取りつつ入力データのサンプリングをした。特にclusteringなどの設定はなさそうなのでlimit句で適当に取得している

WITH item_description_embeddings AS (
  SELECT *
  FROM
    ML.GENERATE_TEXT_EMBEDDING(
      MODEL `public_test.remote_llm_embedding`,
      (
        SELECT distinct
          item_description as content
          ,category_name -- true label

        from `bigquery-public-data.iowa_liquor_sales.sales`
        limit 100
      ),
      STRUCT(TRUE AS flatten_json_output)
    )
),
category_embeddings AS (
  SELECT text_embedding, content
  FROM
    `public_test.embedded_category`
)
SELECT
  c.content AS inferred_category_name,
  i.content AS item_description,
  i.category_name as true_category_name,
  ML.DISTANCE(
    c.text_embedding,
    i.text_embedding,
    'COSINE'
  ) AS sim
FROM
  item_description_embeddings AS i,
  category_embeddings AS c
-- top 1
QUALIFY ROW_NUMBER() OVER (PARTITION BY i.content ORDER BY sim ASC) = 1
ORDER BY i.content, sim

評価

推定したカテゴリと元データとのカテゴリとの一致をみる

結果は 100 件中、9 件だった。一致率は 9%

一致したデータを中心に眺めてみる。単語で判定してるような気もする、というかタスク自体の難易度が高そうだ

| inferred_category_name | item_description | true_category_name | sim | match |
| --- | --- | --- | --- | --- |
| AMERICAN DRY GINS | AVIATION AMERICAN GIN MINI | AMERICAN DRY GINS | 0.2205363973 | 1 |
| WHISKEY LIQUEUR | DR. APPLE WHISKEY MINI | WHISKEY LIQUEUR | 0.2323508543 | 1 |
| WHISKEY LIQUEUR | FIREBALL CINNAMON WHISKEY PARTY BUCKET | WHISKEY LIQUEUR | 0.2686251444 | 1 |
| CANADIAN WHISKIES | HAWKEYE CANADIAN WHISKY MINI | CANADIAN WHISKIES | 0.1854453222 | 1 |
| FLAVORED GIN | INDOGGO STRAWBERRY GIN MINI | FLAVORED GIN | 0.2135974724 | 1 |
| STRAIGHT RYE WHISKIES | WHISTLEPIG 6YR PIGGYBACK RYE | STRAIGHT RYE WHISKIES | 0.2480885032 | 1 |
| JAMAICA RUM | BOMBAY SAPPHIRE | IMPORTED DRY GINS | 0.3321269285 | 0 |

BOMBAY SAPPHIREとかジンの有名ブランドな気もするが、そういう知識は持ち合わせていないのだろうか

まとめ

ほぼ手順を書いて終わった。bigquery上で完結するので、手軽に使えるbigquery-public-dataから選んでみたが、難易度高そうなタスクとなってしまったよう。精度なども見てみたかったが、ドメイン知識のないデータは難しい。今度はいい感じな日本語データセットを探して試してみたいと思った

ちなみに料金は22円だった。入力文字数あたりで課金される

何回か試行錯誤しているので、定かではないが今回は10000レコードくらい試しただろうか

Embeddings for Text - Predictions	Vertex AI 	5,870,820 Characters 	¥22

参考

https://zenn.dev/cloud_ace/articles/bigquery-llm-rag
https://cloud.google.com/blog/ja/products/ai-machine-learning/how-to-use-grounding-for-your-llms-with-text-embeddings
https://cloud.google.com/vertex-ai/generative-ai/pricing?hl=ja