ハッシュタグのカテゴリ分類の比較（BERT vs OpenAI vs ELECTRA）

前回、BERTでテキスト検索を試したときの精度がいまひとつだったので、複数のモデルを使って以下のタスクの結果を比較してみました。

1. ハッシュタグ（短い日本語のテキスト）をベクトル表現に変換する
2. 事前に定義したカテゴリ（日本語）をベクトル表現に変換する
3. ハッシュタグとカテゴリのベクトルのコサイン類似度を計算する
4. コサイン類似度から導出したカテゴリと事前に定めた正解のカテゴリで結果を評価する

上記のタスクを次の3つのモデルでそれぞれ比較しました。

• BERT
• ELECTRA
• OpenAI（Embedding API）

モデル間の分類結果比較

今回は15カテゴリ、695個のハッシュタグを分類しました。
結果は以下のようになりました。

GitHub

Model	Accuracy	Precision	Recall	F1 Score	Elapsed time
BERT	0.12	0.15	0.12	0.12	33.40 seconds
OpenAI Embedding API	0.55	0.58	0.55	0.54	258.63 seconds
ELECTRA Discriminator	0.19	0.19	0.18	0.16	7.14 seconds

予測精度は、OpenAI Embedding APIがすべての評価指標で最も優れた結果になりました。今回はAda v2モデルを使用しました。これはEmbedding APIの中では最も安価なモデルです。他のモデルでは結果が異なるかもしれません。

BERTとその後継のELECTRAを比較すると、やはりELECTRAのほうが性能は高いようです。今回は日本語版のWikipediaで事前学習済みのモデルを使用しました。それでもFine turningなしの状態ではOpenAI Embedding APIのほうが圧倒的に優れている結果となりました。

OpenAI Embedding APIはリクエストのたびに費用がかかりますが、今回695個の日本語のハッシュタグをベクトル変換して消費したトークンは合計3830、計算が正しければ $0.000383（= 0.055 円）なので非常に安価です。

実装詳細

ハッシュタグのベクトル変換

ハッシュタグ文字列をベクトル表現に変換し、DataFrameにその結果を新しい列として追加します。

EmbeddingModelという抽象基底クラス(ABC)に、テキストをベクトルに変換するメソッドを定義しています。このクラスを継承したクラスを、のちほどモデルごとに作成します。

import torch
from abc import ABC, abstractmethod
from sklearn.metrics.pairwise import cosine_similarity
from transformers import BertModel, BertTokenizer

class EmbeddingModel(ABC):
    @abstractmethod
    def convert_to_vector(self, text):
        pass

また、ベクトル変換にかかった時間を計測するための処理を追加します。

import time

def add_vector_column(df: pd.DataFrame, model: EmbeddingModel, column: str) -> pd.DataFrame:
    start_time = time.time()

    # 文字列をベクトルに変換して新しい列に追加
    df.loc[:, column+'_vector'] = df[column].apply(model.convert_to_vector)

    # 処理時間の計測終了と結果の出力
    elapsed_time = time.time() - start_time
    print(f"Elapsed time: {elapsed_time} seconds")

    return df

BERTを使ったEmbeddingModel

BERTModelでは、初期化時にBERTのトークナイザとモデルをロードしています。今回はbert-base-uncasedモデルを使用しています。

class BERTModel(EmbeddingModel):
    def __init__(self, model_name="bert-base-uncased"):
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.model = BertModel.from_pretrained(model_name)
        
    def convert_to_vector(self, text):
        inputs = self.tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)
        vector = outputs.last_hidden_state[0, 0, :].numpy().tolist()
        return vector

OpenAI Embedding APIを使ったEmbeddingModel

OpenAIのEmbedding APIを使用してテキストをベクトルに変換します。

トークナイズが不要で、テキストをAPIに送信するだけなので簡単です。プロンプトを考える手間もありません。

APIはトークン数に応じて課金されるため、消費したトークン数をCSVファイルに保存しています。

class OpenAIModel(EmbeddingModel):
    def __init__(self):
        openai.api_key = os.getenv("OPENAI_API_KEY")
        self.openai = openai
        
    def convert_to_vector(self, text):
        response = self.openai.Embedding.create(
          model="text-embedding-ada-002",
          input=text
        )
        
        # 消費したトークン数をCSVに記録
        filename = 'output/openai_usage_tokens.csv'
        file_exists = os.path.isfile(filename)
        
        with open(filename, mode='a') as file:
            writer = csv.writer(file)
            if not file_exists:
                writer.writerow(['Text', 'Total Tokens'])
            writer.writerow([text, response['usage']['total_tokens']])
        
        # vectorを返却
        return response['data'][0]['embedding']

ELECTRA Discriminatorを使ったEmbeddingModel

ELECTRA（Efficiently Learning an Encoder that Classifies Token Replacements Accurately）を使用してテキストをベクトルに変換します。実装方法はBERTとほぼ一緒でした。

from transformers import ElectraModel, ElectraTokenizer

class ELECTRAModel(EmbeddingModel):
    def __init__(self, model_name="Cinnamon/electra-small-japanese-discriminator"):
        self.tokenizer = ElectraTokenizer.from_pretrained(model_name)
        self.model = ElectraModel.from_pretrained(model_name)
        
    def convert_to_vector(self, text):
        inputs = self.tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)
        vector = outputs.last_hidden_state[0, 0, :].numpy().tolist()
        return vector

コサイン類似度が最大のカテゴリを取得

ハッシュタグとカテゴリそれぞれDataFrameを引数として渡し、各ハッシュタグがどのカテゴリにもっとも類似しているかを予測します。

類似度はベクトル間のコサイン類似度で計算しています。計算結果は-1~1の値を取り、1に近いほど2つのベクトルが類似していることを意味します。

最後にハッシュタグのDataFrameに予測されたカテゴリとそのコサイン類似度の列を追加しています。

import numpy as np

def predict_category(df_hashtag: pd.DataFrame, df_category: pd.DataFrame) -> pd.DataFrame:
    # ベクトルのリストを取得し、NumPy配列に変換
    hashtag_vectors = np.array(df_hashtag['hashtag_vector'].tolist())
    category_vectors = np.array(df_category['category_vector'].tolist())

    # 各ハッシュタグベクトルと全カテゴリベクトルのコサイン類似度を計算
    cosine_similarities = cosine_similarity(hashtag_vectors, category_vectors)

    # 最も類似度が高いカテゴリのインデックスを取得
    predicted_category_indices = np.argmax(cosine_similarities, axis=1)

    # 最大類似度を取得
    max_similarities = np.max(cosine_similarities, axis=1)

    # 予測カテゴリを取得
    predicted_categories = df_category['category'].iloc[predicted_category_indices].values

    # df_hashtagに新しい列を追加
    df_hashtag['predict_category'] = predicted_categories
    df_hashtag['cosine_similarity'] = max_similarities

    return df_hashtag

カテゴリ分類結果の評価

予測結果を、以下4つの評価指標を使って評価しています。

• Accuracy（精度）: モデルがクラスを正しく予測した割合。各クラスのサンプル数がバランスしていない場合（不均衡データ）には、この指標はあてにならない。
• Precision（適合率）: モデルが正と予測したサンプルのうち、実際に正であったサンプルの割合。偽陽性を避ける指標。
• Recall（再現率）: 実際の正のサンプルのうち、モデルが正と予測したサンプルの割合。偽陰性（実際は正だが、予測は負）を避ける指標。
• F1 Score（F1スコア）: F1 PrecisionとRecallの調和平均をとったもので、これら2つの評価指標のバランス。F1スコアは、PrecisionとRecallの両方が重要な場合に特に有用。

さらに、混同行列を作成して、実際のカテゴリと予測カテゴリの組み合わせごとの頻度をヒートマップで可視化しています。

from sklearn.metrics import confusion_matrix, accuracy_score, precision_score, recall_score, f1_score
import seaborn as sns
import matplotlib.pyplot as plt

def evaluation(df_predicted: pd.DataFrame):
    # 予測と実際のカテゴリを取得
    y_true = df_predicted['category']
    y_pred = df_predicted['predict_category']

    # 評価指標の計算
    accuracy = accuracy_score(y_true, y_pred)
    precision = precision_score(y_true, y_pred, average='macro')
    recall = recall_score(y_true, y_pred, average='macro')
    f1 = f1_score(y_true, y_pred, average='macro')

    print(f'Accuracy: {accuracy:.2f}\nPrecision: {precision:.2f}\nRecall: {recall:.2f}\nF1 Score: {f1:.2f}')

    # 混同行列の作成
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(10, 7))
    sns.heatmap(cm, annot=True, fmt="d", cmap='Blues')
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.show()

t-SNEを使ったハッシュタグクラスターの可視化

t-SNE（t-Distributed Stochastic Neighbor Embedding）を用いて高次元のハッシュタグベクトルを2次元に可視化し、予測されたカテゴリごとのクラスタをプロットしています。

予測カテゴリごとに一意な色を割り当てているので、モデルがハッシュタグをうまくカテゴリ分類できている場合は、同じ色が2次元空間上で近くに配置されます。

from sklearn.manifold import TSNE
import matplotlib.cm as cm
import ast

def visualized_in_2d(df_predicted: pd.DataFrame):
    df_predicted['hashtag_vector'] = df_predicted['hashtag_vector'].apply(ast.literal_eval)
    df_predicted['hashtag_vector'] = df_predicted['hashtag_vector'].apply(lambda x: [float(i) for i in x])

    # df_predictedの'hashtag_vector'列から配列を作成
    matrix = np.array(df_predicted['hashtag_vector'].tolist())

    tsne = TSNE(n_components=2, perplexity=15, random_state=42, init="random", learning_rate=200)
    vis_dims2 = tsne.fit_transform(matrix)

    x = [x for x, y in vis_dims2]
    y = [y for x, y in vis_dims2]

    # predict_categoryの一意な値とその色のマッピング
    unique_categories = df_predicted['predict_category'].unique()
    colors = cm.rainbow(np.linspace(0, 1, len(unique_categories)))
    category_color_mapping = {category: color for category, color in zip(unique_categories, colors)}

    for category, color in category_color_mapping.items():
        xs = np.array(x)[df_predicted['predict_category'] == category]
        ys = np.array(y)[df_predicted['predict_category'] == category]
        plt.scatter(xs, ys, color=color, alpha=0.3)

        avg_x = xs.mean()
        avg_y = ys.mean()

        plt.scatter(avg_x, avg_y, marker="x", color=color, s=100)

    plt.title("Hashtag Clusters visualized in 2D using t-SNE")
    plt.show()

OpenAIを使ったカテゴリ分類結果の詳細

ベクトル変換にかかる時間

openai_model = OpenAIModel()

# ハッシュタグをベクトル変換してcsv出力
df_openai = add_vector_column(df, openai_model, 'hashtag')
df_openai.to_csv('output/hashtags_openai.csv', index=False)

# カテゴリをベクトル変換してcsv出力
df_category_openai = add_vector_column(df_category, openai_model, 'category')
df_category_openai.to_csv('output/categories_openai.csv', index=False)

df_openai_usage_tokens = pd.read_csv('output/openai_usage_tokens.csv')
total_tokens = df_openai_usage_tokens["Total Tokens"].sum()
print(f"Total tokens: {total_tokens}, ${(total_tokens/1000)*0.0001}")

Elapsed time: 253.96260690689087 seconds
Elapsed time: 4.664952278137207 seconds
Total tokens: 3830, $0.00038300000000000004

カテゴリ予測

# df_openaiの`hashtag_vector`列をリストに変換し、内部の文字列を浮動小数に変換
df_openai['hashtag_vector'] = df_openai['hashtag_vector'].apply(ast.literal_eval)
df_openai['hashtag_vector'] = df_openai['hashtag_vector'].apply(lambda x: [float(i) for i in x])

# df_category_openaiの`category_vector`列も同様に変換
df_category_openai['category_vector'] = df_category_openai['category_vector'].apply(ast.literal_eval)
df_category_openai['category_vector'] = df_category_openai['category_vector'].apply(lambda x: [float(i) for i in x])

# 予測関数を呼び出す
df_openai_predicted = predict_category(df_openai, df_category_openai)
df_openai_predicted.to_csv('output/hashtags_openai_predicted.csv', index=False)

df_openai_predicted.head()

hashtag	category	hashtag_vector	predict_category	cosine_similarity
サッカー	スポーツ	[-0.010748780332505703, -0.0013573989272117615...	スポーツ	0.904832
野球	スポーツ	[-0.013168673031032085, -0.0203936118632555, 0...	スポーツ	0.881114
ラグビー	スポーツ	[-0.01908210478723049, 0.00013070361455902457,...	スポーツ	0.859531
マラソン	スポーツ	[-0.018582643941044807, -0.010240363888442516,...	スポーツ	0.838464
バレーボール	スポーツ	[-0.021505270153284073, -0.01171257346868515, ...	スポーツ	0.854930

予測結果の評価

evaluation(df_openai_predicted)

Accuracy: 0.55
Precision: 0.58
Recall: 0.55
F1 Score: 0.54

分類に使ったハッシュタグ

予測精度が高かったカテゴリ。

• 13（サイエンス・科学）
• 11（教育・学習）
• 14（ブック・文学）

予測精度が低かったカテゴリ。

• 8（アート・デザイン）
• 9（ファッション・ビューティ）
• 6（クッキング・レシピ）

visualized_in_2d(df_openai_predicted)