💬

LlamaForSequenceClassificationでテキスト分類

2023/10/17に公開

分類

Hugging Face

生成 AI

llama2

LlamaForSequenceClassification

tech

はじめに

生成AIを使う機会が最近多いのですが、分類予測に関しては、いろいろな記事を見ても強化学習や、文章生成中のプロンプト文の形式を指定して分類予測させているのを用いているのをよく目にします。
そんな中、huggingfaceのドキュメントを見てみると、テキスト分類用のライブラリが存在しました。
こちらのライブラリの記事が国内外問わずほとんどなかったので、動作確認をしてみました。

ライブラリのインストール

!pip install transformers 
!pip install sentencepiece

HuggingFaceのログイン

今回使用するLlama2は利用申請が必要です。
1.こちらのリポジトリ内の使用するモデルへの利用申請をしてください。
2.こちらから、公式のMetaのフォームにも記入してください。
上記の２つの申請が通るとアクセスが許可されます。

# HuggingFaceのログイン(事前登録が必要です。)
!huggingface-cli login
'''
※こんな 面倒臭いことしたくないという方は
"model_name = 'openlm-research/open_llama_3b'"にモデルを修正すると、
上記のコードは不要になります！
'''

ライブラリのインポート

from transformers import LlamaForSequenceClassification, LlamaTokenizer
import sentencepiece
import torch

モデルの読み込み

model_name = 'meta-llama/Llama-2-7b-hf' 
tokenizer = LlamaTokenizer.from_pretrained(model_name)
# デフォルトのモデルは２値分類の感情予測です（０：ネガティブ、１：ポジティブ）
model = LlamaForSequenceClassification.from_pretrained(model_name)

モデルの予測

ポジティブコメントの予測

# テキストをトークン化してエンコードします。
text = "love Llama love you"

encoded_input = tokenizer(text, return_tensors='pt')

# モデルに入力し、予測を取得します。
output = model(**encoded_input)

# 予測を取得します。
prediction = torch.argmax(output.logits)
print(prediction)

# output
# tensor(0)

ネガティブコメントの予測

# テキストをトークン化してエンコードします。
text = "Bad Llama you are bad guy"

encoded_input = tokenizer(text, return_tensors='pt')

# モデルに入力し、予測を取得します。
output = model(**encoded_input)

# 予測を取得します。
prediction = torch.argmax(output.logits)
print(prediction)

# output
# tensor(1)