Zenn
📂

日本語要約に特化したLLMをQLoRAを適用したSFTで作ってみる

2024/06/08に公開
lora
LLM
Fine-tuning
sft
tech

インターネットをご覧の皆さん、こんばんは。
皆さんは文章の要約、やっていますか？

やっていますね？

先日 Stability AI よりリリースされた japanese-stablelm-2-base-1_6b をベースモデルとして、 SFT (Supervised Fine-Tuning) を用いて日本語要約に特化した言語モデルを作成してみましたので、本記事ではその学習の流れについて紹介します。

学習に使用したスクリプトは以下のリポジトリに置いています。

https://github.com/ogatayu/summarize-llm-sft

モデルの学習

モデルの学習は以下の流れで行いました。

  1. 学習用データセットの収集
    まず、要約モデルを学習させるために必要なデータセットを収集します。
    要約元となる文章と要約された文章のペアを用意できるようなデータセットを取得します。

  2. 学習用データセットの整形
    収集したデータの前処理を行い、本文と要約のペアを整理します。

  3. モデル学習の実行
    整形したデータを使用してモデルのファインチューニングを行います。
    今回はローカルの環境で学習するために、元モデルを4bitで量子化したのちにQLoRAを適用して学習を行いました。

  4. 学習したモデルで要約を試す
    学習したモデルを用いて、実際に要約を行ってみます。

それでは、各手順の詳細について説明していきます。

学習用のデータセットの収集

なんだかんだで一番頭を悩ませるのが、学習データセットをどのように用意するのかではないでしょうか。
今回は学習に使用するデータセットを KodairaTomonori/ThreeLineSummaryDataset を元に用意しました。

https://github.com/KodairaTomonori/ThreeLineSummaryDataset

こちらのデータセットでは、livedoor ニュースの本文が載っている記事と、その記事の3行要約がペアとして取得できます。

まずは、 livedoorデータセットの使い方 | Hakky Handbook の記述を参考にしながら、記事本文と要約された文章のデータを取得します。

https://book.st-hakky.com/data-science/text-summary-dataset-livedoor/

取得したデータは、例えば以下のような内容になっています。

岡むら屋から、期間限定の新メニュー「じゃが肉めし」が登場する	男爵いもなどは味噌ベースで煮こまれ、しっかり味が染み込んでいるとのこと	「岡むら屋特製肉じゃが」と言うべき一品に、仕上がっているという	新橋と秋葉原に店を構える「岡むら屋」。味噌ベースの独自の味つけで牛バラ肉を煮込んだ具材がたっぷり乗ったオリジナル丼『肉めし』で知られる店です。自分も新橋店で『肉めし』を食べたことがありますが、濃いめの味付けでトロットロに煮こまれた牛肉とホカホカご飯は最高の相性！ 味の染み込んだ豆腐も絶品で、腹ぺこボーイズ&ガールズたちの胃袋を満たし続けている丼なのです。昨年開催された『第2回全国丼グランプリ』では、肉丼部門で金賞受賞を獲得しています。そんな「岡むら屋」に期間限定新メニューが登場。その名も『じゃが肉めし』！肉と芋――このまま名作文学のタイトルにもなりそうな、重厚かつ甘美な響き。これって食いしんぼうたちにとっては定番かつ夢の組み合わせではないでしょうか。この組み合わせではすでに「肉じゃが」という殿堂入りメニューがありますが、今回の『じゃが肉めし』は、“岡むら屋特製肉じゃが”と言うべき一品に仕上がっているそうです。北海道産の男爵いもとしらたきを合わせ、味噌ベースの大鍋で他の具材と一気に煮込むことで、しっかり味が染み込んでいるとのこと。商品写真が公開されていますが、ビジュアルを見ただけで、その染みこみ具合いは一目瞭然！丼というステージで繰り広げられる、肉×芋×白米というスーパースターの競演。肉好き必食の1杯と言えそうです。あ、定食も同時販売されるので、気分で選べるのもうれしいですね。詳しい販売期間などは、店舗にお問い合わせください。肉めし「岡むら屋」
東京駅周辺の安くて美味しい「蕎麦ランチ」の名店を紹介している	「越後そば 東京店」では、ミニかき揚げ丼セットがおすすめと筆者	その他には、「手打ちそば 石月」「酢重正之 楽」「鎌倉 一茶庵 丸山」など	名店がひしめく「丸の内・日本橋」エリアで＜うまい蕎麦ランチ＞が食べられるお店を厳選してご紹介。ぴあMOOK『うまい蕎麦の店 首都圏版』が選んだ、とっておきの7店がこちら！ミニかき揚げ丼セット（温冷） 660円関東風の濃いめのつゆと、ふのりを練り込んだ喉ごしの良い自家製のそばが相性抜群。新鮮な油を使ったかき揚げは、サクサク。千代田区丸の内1-9-1 東京駅一番街B1F蕎麦も天ぷらも正統派の味お蕎麦と天丼（二八） 1600円喉ごしや歯応えを楽しむ二八蕎麦のほか、蕎麦の風味を存分に味わえる十割蕎麦と天丼のセットも。天ぷらは鮮度抜群で美味。千代田区丸の内1-6-4 丸の内オアゾ5F信州の郷土料理をアレンジ“信州フランス鴨”セリかも南蛮（温） 1680円和食とフレンチが融合したスタイルの人気店。このメニューは、本格信州そば、信州のフランス鴨、契約農家が育てた野菜など、素材にこだわっている。千代田区丸の内2-7-2 JPタワーKITTE5F納豆と卵白でふんわり食感なっとうそば 1100円取り寄せた蕎麦の実を、職人が毎日石臼で挽き、丹念に手打ちする。たっぷりの納豆と卵白がのって、なめらかな口当たり。千代田区丸の内1-5-1 新丸の内ビルディング5Fやわらかな厚切り鴨肉を堪能鴨せいろ 1945円、さつま揚げ 650円合鴨・ネギ・しめじなどを合わせた濃厚なつけ汁に、程良いコシの蕎麦がよく合う。自家製さつま揚げもぜひ味わいたい一品。千代田区丸の内2-4-1 丸の内ビルディング6Fコシの強い独特な田舎蕎麦が美味とり辛そば（温） 980円信州軽井沢の味噌・醤油屋「酢重正之商店」が手掛ける蕎麦屋の人気メニュー。辛口のつけ汁と太い田舎蕎麦がよく合うと評判。千代田区丸の内1-5-1 新丸の内ビルディングB1F風味豊かな生わさびが決め手ざるそば（生わさび） 710円信州の民家を思わせる店内で打つ自家製麺は歯応え抜群。生わさびを自分でおろして味わう。甘めのつゆがわさびにぴったり。千代田区丸の内1-6-1 丸の内センタービル B1Fうまい蕎麦の店　首都圏版日本人ならいつでも蕎麦が食べたい！
8日から「サンクトガーレン」は、「チョコビール」4種を販売する	ダークな茶色、ほろ苦いビターチョコのような香りが特長のビール	毎年バレンタインシーズン限定で販売され、2016年で11年目となるそう	地ビールメーカー「サンクトガーレン」は、「チョコビール」4種を2016年1月8日から全国の取扱店舗およびオンラインショップで販売します。「チョコビール」は、ビールに使う麦芽を高温で焙煎することで生まれる、ダークな茶色、ほろ苦いビターチョコのような香りが特長のビールです。毎年バレンタインシーズン限定で販売され、今年で11年目となります。発売当初から人気の「インペリアルチョコレートスタウト」（648円）は、通常のお正月飾りにも欠かせない橙（だいだい）を皮ごと使った「オレンジチョコレートスタウト」（540円）は、凝縮されたオレンジの風味と、マーマレードのような皮の苦みが特徴的なビールです。エスプレッソのような濃厚黒ビールにバニラの香りを溶け込ませた「スイートバニラスタウト」（463円）は、甘くまったりとした飲み口で、後味はバニラチョコのよう。そして2016年の限定フレーバー「ストロベリーチョコレートスタウト」（540円）は、甘酸っぱい「とちおとめ」をたっぷり使用した新感覚のビールです。取扱店舗一覧は公式サイトで。また、オンラインショップでは先行予約も受け付け中です。

（livedoor ニュースの記事と要約文章から引用しています）

各行のフォーマットとしては、以下の形式になっています。

要約1\t要約2\t要約3\t本文

次のステップでは、収集できたこちらのデータセットを学習に使いやすいように整形していきます。

学習用データセットの整形

要約元の文章と要約された文章がペアになったデータが取得できたので、次は学習用にデータの整形をしていきます。

収集されたデータには一部不完全な文章が含まれていたりしますので、前処理として以下のような方針でフィルタリングを行います。

  • 要約と本文の類似度を計算し、類似度が高いペアのみを学習データとして使用
    • 要約と本文のペアが適切であるかを確認するために要約と本文の類似度を計算し、類似度が高いペアのみを学習データとして使用します
  • 要約や本文が短すぎるペアは除外
  • 要約と本文の長さの比率が一定の割合を満たさないペアは除外
  • 要約と本文のの長さが一定の範囲内に収まるペアを選定

上記を踏まえて、以下のスクリプトを作成しました。

こちらのスクリプトでは、収集したデータセットから要約と本文のペアを作成し、適宜フィルタリングを行ってから、最終的にJSON形式で保存しています。

モデル学習の実行

学習用のデータセットが準備できたので、これからは実際にモデルのファインチューニングを行っていきます。

今回はローカル環境で学習を行いたかったため、学習に必要なリソースを抑えられるように元モデルを4bit量子化してから、QLoRAを適用して学習を行います。

学習に使用したスクリプトは以下のファイルです。

ここからは上記スクリプトを一部抜粋しながら、学習の流れについて説明していきます。

モデルとトークナイザーの準備

まずは、モデルとトークナイザーを読み込みます。

以下のコードでは、stabilityai/japanese-stablelm-2-base-1_6b のモデルをBnB 4bitで量子化してロードしています。
（モデルを使用するには事前にHugging Face上で連絡先の登録と各条項への同意が必要になりますのでご注意ください）

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

MODEL_NAME = "stabilityai/japanese-stablelm-2-base-1_6b"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    useFast=False,
    trust_remote_code=True,
)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map={"":0},
    trust_remote_code=True,
    torch_dtype="auto",
)

print(model)

モデルの準備

ロード & 量子化したモデルに対して、 PEFT を使用して QLoRA の適用を行います。

from peft import prepare_model_for_kbit_training

model.enable_input_require_grads()
model.gradient_checkpointing_enable()

model = prepare_model_for_kbit_training(model)

from peft import LoraConfig, get_peft_model

config = LoraConfig(
    target_modules=[
        #"embed_tokens",
        #"lm_head",
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
    r=16, 
    lora_alpha=32,
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
    inference_mode=False,
)

model = get_peft_model(model, config)

データセットのロードとトークン化

学習データセットをロードし、トークナイザーを用いてトークン化します。
japanese-stablelm-2-base-1_6b は指示応答学習をされていないベースモデルであるため、プロンプトのフォーマットは適当に決めています。

import datasets

prompt_template = """以下の文章を要約してください。

### 文章:
{input}

# 要約結果:
"""

def encode(sample):
    prompt = prompt_template.format(input=sample["input"])
    target = sample["summary"] + tokenizer.eos_token
    
    input_ids_prompt, input_ids_target = tokenizer([prompt, target]).input_ids
    input_ids = input_ids_prompt + input_ids_target
    
    labels = input_ids.copy()
    labels[:len(input_ids_prompt)] = [-100] * len(input_ids_prompt)
    
    return {"input_ids": input_ids, "labels": labels}

def get_collator(tokenizer, max_length):
    def collator(batch):
        batch = [{ key: value[:max_length] for key, value in sample.items() } for sample in batch ]
        batch = tokenizer.pad(batch, padding=True)
        batch["labels"] = [ e + [-100] * (len(batch["input_ids"][0]) - len(e)) for e in batch["labels"] ]
        batch = { key: torch.tensor(value) for key, value in batch.items() }
        return batch
    return collator

dataset = datasets.load_dataset('json', data_files='./dataset/three_line_summary.json')
dataset = dataset.map(encode)
dataset = dataset["train"].train_test_split(0.2)
train_dataset = dataset["train"]
eval_dataset = dataset["test"]

#  評価用のデータセットを保存しておく
eval_dataset.remove_columns(['input_ids', 'labels']).save_to_disk(f"./{session_path}/eval_dataset")

データセットは、学習用と評価用に8:2の割合で分割しています。
本記事では触れていませんが後ほどモデルを評価するときに使用できるように、評価用に分割したデータセットを別途保存しています。

学習の実行

transformers ライブラリの Trainer クラスを設定し、学習を実行します。

学習性能の向上を期待して NEFTune を適用していますが、適用すると出力される文章が長くなる傾向があるようなので、要約タスクの学習に適用するのは微妙かもしれないです。

import transformers
training_arguments = transformers.TrainingArguments(
    output_dir=f"./{session_path}/checkpoints",
    learning_rate=2e-5,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    per_device_eval_batch_size=1,
    num_train_epochs=1,
    logging_strategy='steps',
    logging_steps=10,
    save_strategy='epoch',
    evaluation_strategy='epoch',
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    save_total_limit=2,
    optim="paged_adamw_8bit",
    neftune_noise_alpha=5,
)

trainer = transformers.Trainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    args=training_arguments,
    data_collator=get_collator(tokenizer, 2048)
)

model.config.use_cache = False

trainer.train()

model = trainer.model
peft_model_name = f"./{session_path}/model"
model.save_pretrained(peft_model_name)

私の環境(RTX4070 12GB)では、1時間程度で学習が完了しました。

学習後のLoRAモデルは ./session/***_japanese-stablelm-2-base-1_6b/model 以下に保存しています。

学習したモデルで要約を試す

学習したモデルを用いて実際に推論を実行し、日本語記事の要約を行ってみます。

推論は下記のようなコードで実行できます。

import torch
import peft
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import time

# ベースモデルとLoRAモデルを選択
MODEL_NAME = "stabilityai/japanese-stablelm-2-base-1_6b"
peft_model_path = "./session/***_japanese-stablelm-2-base-1_6b"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    padding_side="left",
    trust_remote_code=True,
)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map={"":0},
    trust_remote_code=True,
)
model = peft.PeftModel.from_pretrained(model, f"{peft_model_path}/model")

prompt_template = """以下の文章を要約してください。

### 文章:
{input}

# 要約結果:
"""

input_texts = [
    # https://ja.stability.ai/blog/japanese-stable-lm-2-16b
    """# 日本語大規模言語モデル「Japanese Stable LM 2 1.6B」をリリースしました
製品 9 May

## ポイント

Japanese Stable LM 2 1.6B（JSLM2 1.6B）は16億パラメータで学習した日本語の小型言語モデルです。

JSLM2 1.6Bのモデルサイズを16億パラメータという少量にすることによって、利用するために必要なハードウェアを小規模に抑えることが可能であり、より多くの開発者が生成AIのエコシステムに参加できるようにします。

ベースモデルとしてJapanese Stable LM 2 Base 1.6Bと、指示応答学習（Instruction tuning）済みのJapanese Stable LM 2 Instruct 1.6Bを提供します。両モデルともStability AI メンバーシップで商用利用が可能です。また、どちらのモデルもHugging Faceからダウンロードすることができます。

    - Japanese Stable LM 2 Base 1.6B
    - Japanese Stable LM 2 Instruct 1.6B

‘A beautiful anime-like hummingbird flying with the text "Japanese Stable LM 2" below it, with a lofi anime landscape of Mount Fuji forming the outline of the text "Japanese Stable LM 2"’ — Stable Diffusion 3で生成

Stability AI Japanは16億パラメータで学習した日本語の言語モデルJapanese Stable LM 2 1.6B（JSLM2 1.6B）のベースモデルと指示応答学習済みモデルをリリースしました。ベースモデルの学習ではWikipediaやCulturaX等の言語データを利用、指示応答学習ではjaster、Ichikara-Instruction、Ultra Orca Boros v1の日本語訳等、商用データおよび公開データを利用しました。今回のJSLM2 1.6B では、言語モデリングにおける最新のアルゴリズムを活用し、適度なハードウェアリソースで迅速な実験を繰り返すことを可能にし、スピードと性能を両立しました。

## 性能評価

Nejumiリーダーボードを用いて、他の小規模パラメータのモデルと比較したJSLM2 1.6Bの性能は以下のとおりです。今回はllm-leaderboard（の社内Fork）のcommit c46e165を用いています。


16億パラメータという小型モデルでありながら、40億パラメータ以下のモデルのスコアよりも高いスコアを達成し、70億パラメータのモデルに近いスコアを獲得しています。

高性能な小型言語モデルをリリースすることで、言語モデル開発の敷居を下げ、より高速に実験を反復することを可能にします。なお、少ないパラメータ数の小型モデルであるため、より規模の大きいモデルで発生しうるハルシネーションや間違いをおかす可能性があります。アプリケーションでのご利用の際には適切な対策を取るようご注意ください。JSLM2 1.6Bのリリースを通じて、日本語LLMのさらなる開発と発展に貢献できると幸いです。

## 商用利用について

JSLM2 1.6BはStability AI メンバーシップで提供するモデルのひとつです。商用でご利用したい場合は、Stability AIメンバーシップページから登録し、セルフホストしてください。



Stability AI の最新情報は公式X、Instagram をチェックしてください。
"""
]

print(f"### {peft_model_path} ###")

with open(peft_model_path + "/test.log", "w", encoding="utf-8") as f:
    f.write(f"### {peft_model_path} ###\n")
    
for input_text in input_texts:
    print("===")
    
    start_time = time.time()

    prompt = prompt_template.format(num="100", input=input_text)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        tokens = model.generate(
            **inputs,
            max_new_tokens=1024,     
            repetition_penalty=1.1
        )

    output = tokenizer.decode(tokens[0], skip_special_tokens=False)
    print(output)
        
    output = output.split(f"# 要約結果:\n")[1]
    
    end_time = time.time()
    elapsed_time = end_time - start_time

    print("Elapsed time:", elapsed_time, "seconds")
    
    with open(peft_model_path + "/test.log", "a", encoding="utf-8") as f:
        f.write("===\n")
        f.write(f"{output}\n")
        f.write(f"Elapsed time: {elapsed_time} seconds\n")

要約元の文章は、Japanese Stable LM 2 1.6Bがリリースされた際の記事から引用しています。

https://ja.stability.ai/blog/japanese-stable-lm-2-16b

生成された要約の文章は以下のようになりました。

日本語の小型言語モデル「Japanese Stable LM 2 1.6B」をリリースした。16億パラメータで学習した日本語の言語モデル。ベースモデルとしてJapanese Stable LM 2 Base 1.6Bと、指示応答学習済みのモデルを提供する。

いい感じに要約できているように見えます。やったー！

おわりに

今回は、Stability AIよりリリースされた日本語特化の言語モデル「Japanese Stable LM 2 1.6B」をベースに、SFTを用いて日本語の記事を要約するモデルを作成しました。
最終的には、それなりに違和感なく要約を行えるモデルにできたかと思います。

全体的に省リソースの環境でも実行できるような構成で学習を行ったため、RTX4070 12GBのようなミドルクラスのGPUでも差し支えなくモデルの学習を実行できました。
最近は小規模で高性能なモデルが各所からリリースされ始めてきており、LLMのファインチューニングを手軽に試せる環境になってきていますね。

本記事では触れませんでしたが、作成したモデルの要約性能の定量評価も行ってみましたので、そちらの結果についてはまた別の記事で紹介できればと思います。

それではさようなら。また会いましょう。

参考にした記事

ここまでの作業を行うにあたり、以下の記事を参考にさせていただきました。

Discussion