Scikit-LLMでテキスト分類しました

あけましておめでとうございます。データサイエンスグループの木下です。

今回は、Scikit-LLMというライブラリを紹介いたします。

Scikit-LLMとは

skllm.beastbyte.ai Scikit-LLMとは、LLMをScikit-Learnのように扱うことができるライブラリです。 Scikit-LearnユーザーになじみやすいAPI設計になっており、学習コストが低いのが特徴です。

上記のドキュメントによると、主なタスクは

Text-Classification
指定したラベルで文章分類ができます。 Zero-shotとFew-Shot、CoT(Chain of Thought)などに対応しています
Text-to-text modelling
要約タスク、翻訳タスクなどができるそうです。
Text Vectorization
テキストの埋め込み表現を作成できます。

Scikit-LLMの実装

では、実際にScikit-LLMを実装してみます。データセットは、料理名とその説明文がセットになったFood Description dataというオープンデータにしました。

本ブログでは、
①英語の説明文を日本語の説明文に翻訳
②日本語の説明文から指定した味覚ラベルに分類
という、２つのタスクを行います。

データの前処理

Food Description dataは下記のようになっており、全部で2602行ありました。今回は、この中からDescriptionの文字列が100文字以上ある65行のデータに絞りました。

Scikit-LLMのインストールとAzureOpenAIのconfig指定

まずは、Scikit-LLMをインストールします。

pip install scikit-llm

今回は、AzureOpenAIを使用します。その場合、ご自身で使われているAzureOpenAIのOPENAI_API_KEYとOPENAI_API_BASEをScikit-LLMのconfigに設定してあげる必要があります。また、モデルの名前には、プレフィックスとして、"azure::"を付けてあげる必要があります。

from skllm.config import SKLLMConfig

# use azure key instead
SKLLMConfig.set_openai_key(OPENAI_API_KEY)
SKLLMConfig.set_azure_api_base(OPENAI_API_BASE)

# start with "azure::" prefix when setting the model name
model_name = "azure::gpt-4"

翻訳タスクの実行

from skllm.models.gpt.text2text.translation import GPTTranslator

t = GPTTranslator(model = model_name, output_language="Japanese")
translated_text = t.fit_transform(df['Description'].tolist())

df['translated'] = translated_text

翻訳タスクは、GPTTranslatorというクラスで行います。使用するgptのモデル名と、翻訳先の言語を指定してインスタンスを作成し、そのインスタンスにfit_transform関数を適用することで、翻訳が完了します。

実際に出力された内容は下記のとおりです。

分類タスクの実行

from skllm.models.gpt.classification.zero_shot import ZeroShotGPTClassifier

clf = ZeroShotGPTClassifier(model = model_name)
labels = ['甘い','しょっぱい','酸っぱい','辛い']
clf.fit(None, labels)

infered = clf.predict(df['translated'].tolist())

df['label'] = infered

今回は、zero-shotの分類を行ったので、 ZeroShotGPTClassifierというクラスを用いました。こちらは、モデル名のみ指定してインスタンスを作成し、 fit関数でインスタンスに分類したいラベルを指定します。その後、predict関数で分類タスクを実行します。

65件の分類結果は、下記のようになりました。