🐈

大規模小説データをOpenAI Embedding APIで特徴抽出してUMAPで可視化してみた

テラーノベルで機械学習を中心に担当している川尻です。みなさんも大量のデータを特徴抽出してみたけど、どう使ったらいいかわからないなぁということありますよね。とりあえずどういうデータなのか雰囲気だけでもみるために、簡単な可視化から入るかと思います。以前にもグラフデータの可視化について記事を書きました。

https://zenn.dev/tellernovel_inc/articles/53507126aa77bf

テラーノベルでは小説投稿アプリを作っており、OpenAI Embedding APIを使って全作品の特徴量を出しているデータベースがあります。今回はこのデータを例にして、UMAPという可視化に便利な次元削減手法の使い方を紹介します。

UMAPとは

UMAPというのは非線形の次元削減手法で、特徴は次元数やデータ量に対してほとんど一定の計算時間で済むということです。LLMなどの事前学習モデルでよく扱うような1,000次元を超えるような特徴量でも全く問題なく動きます。さらに、Pythonでとても使いやすいパッケージが公開されています。
https://umap-learn.readthedocs.io/en/latest/index.html

UMAPの仕組みなどの解説はこちらの記事が分かりやすくまとまっています。
https://qiita.com/odanny/items/06ab88353bcee7bf6aa7

実行環境とデータの内容

ここからは、google colabを使う前提で進めます。事前に以下のようにumapをインストールしておきます。

!pip install umap-learn

以下のように必要なパッケージをインポートしておきます。

import matplotlib.pyplot as plt
import umap

今回は、BigQueryからpandasのDataFrameの形式で取得して以下のように取得できているとします。

bigqueryを使うためのgoogle認証
from google.colab import auth
auth.authenticate_user()
bigquery magicを使ってデータを取ってくる例
%%bigquery features --use_bqstorage_api
select
... 省略 ...

特徴抽出には、Open AI Embedding APIを使っていて次元数は1536です。UMAPの計算時間は、体感的にはデータサイズに比例します(理論的には計算量は線形よりも若干大きいオーダーです)。今回は試行錯誤もするために1回の次元削減の計算が30秒位で終わるようにレコード数を間引いてあります。もし計算時間を待つことができればメモリーサイズの限界まで増やすことができて、今回使った標準のgoogle colab(RAM=12.7GB)では、10万件では3,4分で計算が完了しましたが、20万件ではRAM不足でできませんでした。また、可視化で一緒に表示するためにジャンル情報(Genre)と公式/インディーズ情報(IsOfficial)のカラムを一緒に取得しています。ジャンル情報について、公式作品については人手である程度チェックして精度が高いですが、インディーズ作品はほとんどが自動判定であるため精度が低いです。それぞれの内訳は以下のようになっています。

公式作品の内訳
features["IsOfficial"].value_counts()
出力
False    20366
True      7113
Name: IsOfficial, dtype: Int64
ジャンルの内訳
features["Genre"].value_counts()
出力
horror_thriller    9227
romance            5845
bl                 5684
fantasy            2513
comedy             1523
drama              1366
unspecified         882
romance_fantasy     439
Name: Genre, dtype: int64

とりあえずUMAPで可視化

次元削減するコードは以下のようになります。

シンプルな次元削減
%%time
u = umap.UMAP(metric="cosine").fit_transform(features["Feature"].to_list())
処理にかかった時間
CPU times: user 51.8 s, sys: 1 s, total: 52.8 s
Wall time: 38.5 s

u は二次元に変化されたデータが入っています。OpenAIのドキュメント [1] によるとコサイン距離が推奨されているので、メトリックを指定しています。

以下のようにジャンルごとに色を分けて可視化してみます。

ジャンルごとに色分け
for i, genre in enumerate(genres):
    idx = features["Genre"] == genre
    plt.scatter(u[idx, 0], u[idx, 1], label=genre, s=0.5, alpha=0.5)
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')

教師なしジャンル

ジャンルごとに大まかには固まっていそうですが、ジャンルが入り乱れているところや、外れている値も結構ありそうですね。

続いて、公式/インディーズで色を付けてみます。

公式/インディーズで色分け
idx = (features["IsOfficial"] == False).to_list()
plt.scatter(u[idx, 0], u[idx, 1], label="Indies", s=0.5, alpha=0.5)
idx = (features["IsOfficial"] == True).to_list()
plt.scatter(u[idx, 0], u[idx, 1], label="Official", s=0.5, alpha=0.5)
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')

教師なし公式/インディーズ

こちらのほうがよりはっきりと分かれそうで、公式作品のようなインディーズ作品を判定する事もできそうです。

教師ありUMAP

UMAPは教師ありUMAPも対応していて、同じラベルが付いたものは近くになるように次元削減します。以下のようにラベルを整数に変換して渡します。

教師ありUMAP
%%time
s = umap.UMAP(metric="cosine").fit_transform(features["Feature"].to_list(), features["Genre"].map(lambda x: genres.index(x)).to_list()) 
処理時間
CPU times: user 49.8 s, sys: 2.43 s, total: 52.3 s
Wall time: 34.7 s
ジャンルごとに色分け
for i, genre in enumerate(genres):
    idx = features["Genre"] == genre
    plt.scatter(s[idx, 0], s[idx, 1], label=genre, s=0.5, alpha=0.5)
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')

教師ありUMAP

先程と比べると、はっきりとジャンルごとに分かれるようになりました。特にデータ量が多いジャンルではしっかり集まっていて判定しやすそうですが、データが少ないジャンルでは他のジャンルと混ざっていて推定するのが難しそうです。

準教師ありUMAP

実際のデータでは全部のデータについて欲しい情報がついていることはあまりなく、しっかりと判定モデルを作成する前にどれくらい判定できそうかのあたりを付けたいことが多いかと思います。つづいては、準教師ありUMAPでこれをやってみます。

すでに説明したように、インディーズ作品のジャンルは作成した判定モデルで自動分類したものです。判定モデルを作成する際には、公式作品についてだけ人手でつけたジャンルを教師データとして使いました。そのため以下のように、公式作品のデータだけラベル情報として与えて、インディーズ作品はラベルなしとして処理される、-1をつけます。

準教師ありUMAP
%%time
features["Label"] = -1  # ラベルなしは -1
official_idx = (features["IsOfficial"] == True).to_list()
features["Label"][official_idx] = features["Genre"][official_idx].map(lambda x: -1 if x == "unspecified" else genres.index(x))  # unspecifiedも判定できなかったものなのでラベルなしに
ss = umap.UMAP(metric="cosine").fit_transform(features["Feature"].to_list(), features["Label"].to_list())
処理時間
CPU times: user 1min, sys: 1.58 s, total: 1min 2s
Wall time: 39 s

ラベルなしデータはグレーにして、可視化してみます。

unlabeled_idx = (features["Label"] == -1).to_list()
plt.scatter(ss[unlabeled_idx, 0], ss[unlabeled_idx, 1], label="unknown", s=0.5, c="gray", alpha=0.1)

labeled_idx = (features["Label"] != -1).to_list()
for genre in genres:
    idx = features[labeled_idx]["Genre"] == genre
    plt.scatter(ss[labeled_idx][idx, 0], ss[labeled_idx][idx, 1], label=genre, s=0.5, alpha=0.5)

plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')

準教師ありUMAP

こちらをみてもやはり、データが多いジャンルではうまく分離できそうなことが見て取れますね。一方で、ラベルなしのインディーズ作品だけがまとまっているところもあるので、インディーズ特有のジャンルがあったりしそうです。必要に応じて、クラスタリングなどして作品を抜き出して、定性的に確認していくといいかもしれないです。

まとめ

可視化のための次元削減の手法としてUMAPを紹介して、実際にざっくりした分析をしてみました。ここまでの説明でお気づきの方もいるかと思いますが、UMAPの次元削減は可視化だけではなくて、メトリック学習やクラスタリング、異常値検知などにも応用できます。UMAPの公式ドキュメントにも分かりやすく説明が書いてあるのでぜひ見てみてください。

脚注
  1. embedding APIの説明: https://platform.openai.com/docs/guides/embeddings/limitations-risks ↩︎

テラーノベル テックブログ

Discussion