LLMにおいてPEFTを使ってLoRAでテキスト分類用の学習を行う際のメモ

2023年5月30日
機械学習

目次

はじめに

ここでは、PytorchやTransformersを用いた機械学習のやり方の基本は理解しているものとして、LoRAで学習をする際の要点だけ記載する。

トークナイズについて

まずは、どんなトークナイズが必要かを知っておくのが大事。（適当にやっていると沼にハマる）
というわけで、スペシャルトークンを確認する。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-small")
print(tokenizer.special_tokens_map)
print("bos_token :", tokenizer.bos_token, ",", tokenizer.bos_token_id)
print("eos_token :", tokenizer.eos_token, ",", tokenizer.eos_token_id)
print("unk_token :", tokenizer.unk_token, ",", tokenizer.unk_token_id)
print("pad_token :", tokenizer.pad_token, ",", tokenizer.pad_token_id)

semt = tokenizer(line[1]+"<|endoftext|>",return_tensors='pt',truncation=True,max_length=512,padding="max_length",)["input_ids"][0]
{'bos_token': '<|endoftext|>', 'eos_token': '<|endoftext|>', 'unk_token': '<|endoftext|>', 'pad_token': '<|padding|>'}
bos_token : <|endoftext|> , 0
eos_token : <|endoftext|> , 0
unk_token : <|endoftext|> , 0
pad_token : <|padding|> , 1

eos_tokenは0、pad_tokenは1となっているので、最終的に以下のようなトークナイズがされていれば良いということになる。よく確認せずに、BERTとかのコードをモデル名だけ変えて流用すると酷いことになる。

tensor([ 4981,  1317, 10340,  1847,   247,     0,     1,     1,     1,     1,
            ..., 1])

よって、トークナイズを行うコードは以下のようにする。

tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-small", add_eos_token=True)
semt = tokenizer("今日はいい天気ですね。"+"<|endoftext|>",return_tensors='pt',truncation=True,max_length=512,padding='max_length')["input_ids"][0]

LoRAを用いた学習

まず、テキスト分類を行う際にはEmbeddingを線形層に与える必要があるが、自作で行うときはベースモデルの構造とかをよく調べて実装する。

ここでは、Hugging FaceのTransformersを用いて簡単にテキスト分類を行う。
この場合は、AutoModelForSequenceClassificationやLlamaForSequenceClassificationなどのXXXForSequenceClassificationとなっているものを使用する。これは、ベースモデルの後ろにクラス分類のための線形層を追加したものである。

GPT-NeoXやLlamaの場合はCLSではなく文の最後のトークンのEmbeddingをクラス分類に使用しているようなので、eos_tokenが文の最後にあることが重要だと思われる。パディングがきちんとできていれば、トークン列の最後または、pad_tokenの直前のトークンをクラス分類のためのEmbeddingとして使用してくれるみたい…

また、LoRAを適用するのには同じくHugging FaceのPEFTを用いると簡単に実装できる。
これらを用いた、モデルの読み込みとLoRAの設定は以下のようにする。

from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel
from peft import (
    LoraConfig,
    get_peft_model,
    get_peft_model_state_dict,
)

model = AutoModelForSequenceClassification.from_pretrained(
    "cyberagent/open-calm-small",
    num_labels = 3 # 分類するクラスの数
)
model.config.pad_token_id = 1 # 上で調べたpad_tokenを設定する

config = LoraConfig(
    r=8,
    lora_alpha=16,
    # target_modules=["q_proj", "v_proj"], # Llamaを用いる場合
    target_modules=['query_key_value'], # GPT-NeoXを用いる場合
    lora_dropout=0.05,
    bias="none",
    task_type="SEQ_CLS", # テキスト分類の場合は"SEQ_CLS"となる
)
model = get_peft_model(model, config)

model.print_trainable_parameters()

target_modulesはベースとなるモデルによって変わるので注意する。
task_typeはテキスト分類の場合は、”SEQ_CLS”となる。”TOKEN_CLS” や “SEQ_2_SEQ_LM” 、 “Causal_LM” など色々あるので間違わないようにする。

また、model.print_trainable_parameters()でパラメータ全体のうちどれくらいを学習させているかを確認できる。ここで、学習させているパラメータが全てでない（＝LoRAで置き換えた部分以外は学習させていない）ことを確認しておくと良い。

あとは、普通に学習させれば良い。
モデルの保存は以下のようにする。

model.save_pretrained("model_dir_path")

保存したモデルのロード

以下のように、モデルをロードして推論に使用する。

from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel, PeftConfig

# config = PeftConfig.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(
    "cyberagent/open-calm-small",
    num_labels = 3 # 分類するクラスの数
)
model.config.pad_token_id = 1 # 上で調べたpad_tokenを設定する(ここでも必要かは不明)
model = PeftModel.from_pretrained(model, "model_dir_path")