📕
MLMによるLUKEのドメイン適応

2024/10/27に公開
 はじめに機械学習エンジニアをしています、ishi2kiです。

今回は、日本語の事前学習済みモデルLUKEに対してドメイン適応を行う方法について説明します。
必要知識
Pythonの基礎知識
PyTorchの基礎知識
自然言語処理の基礎知識

 使用技術・用語
 LUKE日本で開発されたTransformerベースのエンコーダモデルです。

Entity (固有表現) の埋め込みを適切に行えるように学習したものですが、Entityに関するタスク以外でも様々なタスクにおいて、他の事前学習済みモデルよりも高い性能を発揮しています。

日本語の事前学習済みモデルも公開されており、今回はluke-japanese-large-liteを使用します。

 MLM (Masked Language Modeling)MLMはエンコーダの学習方法の一種です。

文中のトークンの一部をマスキングして (隠して) 、マスキングされたトークンが何かを予測することでモデルを学習します。

ファインチューニング時には、文とラベルのペア、文と文のペアなど、ラベリングされたデータが必要となりますが、MLMは文さえあれば学習可能である点が特徴です。


 ドメイン適応LUKEの日本語事前学習済みモデルは、日本語Wikipediaで学習されています。

この事前学習で日本語の普遍的な知識 (文法や頻出単語) を獲得することはできますが、専門的な領域の知識を獲得するためには、その領域のデータセットを使って学習する必要があります。

ファインチューニングを行う際にラベル付きのデータを十分に用意できるのであれば問題ないのですが、データが少量しかない場合、ラベル付きデータだけでは十分に知識を学習することができないことがあります。
そこで、事前学習済みモデルを、MLMで追加学習を行うという手法があります。

前述した通り、MLMは文さえ集められればラベリング不要で学習することが可能であり、少ない労力で専門知識を獲得することが期待できます。

もちろん、クラスタリングなどのタスクを行うにはラベリングされたデータも必要となりますが、MLMでの追加学習を行っておくことで、このデータ量も少なくて済むようになります。



今回は、「小説」ドメインにモデルを適応させます。

 実装
 データの用意「小説」にモデルを適応させるためには、小説の文を集めることが必要です。

今回は、青空文庫から小説本文を収集し、文単位に分割してテキストファイル (novel_data.txt) に保存しました。
吾輩は猫である。
名前はまだ無い。
どこで生れたかとんと見当がつかぬ。
何でも薄暗いじめじめした所でニャーニャー泣いていた事だけは記憶している。
...

 データの加工まず、データを加工してDatasetに変換します。

ここでは、以下の処理を行っています。
文をトークンに分割する
モデルの最大入力長を超えないように、ブロック分割する
tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-japanese-large-lite")
block_size = 128


def tokenize(text):
    return tokenizer(text["text"])


def group_texts(examples):
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    if total_length >= block_size:
        total_length = (total_length // block_size) * block_size
    result = {
        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
        for k, t in concatenated_examples.items()
    }
    return result


def main():
    sentences = datasets.load_dataset(
        "text", data_files="novel_data.txt", split="train"
    )
    sentences = sentences.train_test_split(test_size=0.2).flatten()
    tokenized_sentences = sentences.map(
        tokenize, remove_columns=sentences["train"].column_names
    )
    lm_dataset = tokenized_sentences.map(group_texts, batched=True, num_proc=1)
    lm_dataset = lm_dataset.with_format("torch", device=device)

 モデルの学習モデルの学習には、Trainerを使用しています。

また、DataCollatorを使うことで、マスク処理もPyTorchに任せています。
epoch = 4
batch = 1
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def main():
    model_dir = "out_model"
    os.makedirs(model_dir, exist_ok=True)

    tokenizer.pad_token = tokenizer.sep_token
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer, mlm_probability=0.15, return_tensors="pt"
    )
    data_collator = lambda data: dict(data_collator(data))

    model = LukeForMaskedLM.from_pretrained("studio-ousia/luke-japanese-large-lite")
    model = model.to(device)

    training_args = TrainingArguments(
        output_dir=model_dir,
        evaluation_strategy="epoch",
        logging_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=args.epoch,
        learning_rate=5e-5,
        lr_scheduler_type="constant",
        num_train_epochs=args.epoch,
        weight_decay=0.01,
        per_device_train_batch_size=1,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=lm_dataset["train"],
        eval_dataset=lm_dataset["test"],
        data_collator=data_collator,
    )

    trainer.train()
コード全文import os

import datasets
import torch
from transformers import (AutoTokenizer, DataCollatorForLanguageModeling,
                          LukeForMaskedLM, Trainer, TrainingArguments)

tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-japanese-large-lite")
block_size = 128
epoch = 4
batch = 1
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def tokenize(text):
    return tokenizer(text["text"])


def group_texts(examples):
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    if total_length >= block_size:
        total_length = (total_length // block_size) * block_size
    result = {
        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
        for k, t in concatenated_examples.items()
    }
    return result


def main():
    model_dir = "out_model"
    os.makedirs(model_dir, exist_ok=True)

    sentences = datasets.load_dataset(
        "text", data_files="novel_data.txt", split="train"
    )
    sentences = sentences.train_test_split(test_size=0.2).flatten()
    tokenized_sentences = sentences.map(
        tokenize, remove_columns=sentences["train"].column_names
    )
    lm_dataset = tokenized_sentences.map(group_texts, batched=True, num_proc=1)
    lm_dataset = lm_dataset.with_format("torch", device=device)
    tokenizer.pad_token = tokenizer.sep_token
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer, mlm_probability=0.15, return_tensors="pt"
    )
    # solution to transfomer's issue
    # ref: https://github.com/nlp-with-transformers/notebooks/issues/31
    data_collator = lambda data: dict(data_collator(data))

    model = LukeForMaskedLM.from_pretrained("studio-ousia/luke-japanese-large-lite")
    model = model.to(device)

    training_args = TrainingArguments(
        output_dir=model_dir,
        evaluation_strategy="epoch",
        logging_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=epoch,
        learning_rate=5e-5,
        lr_scheduler_type="constant",
        num_train_epochs=epoch,
        weight_decay=0.01,
        per_device_train_batch_size=1,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=lm_dataset["train"],
        eval_dataset=lm_dataset["test"],
        data_collator=data_collator,
    )

    trainer.train()


if __name__ == "__main__":
    main()

 おわりにMLMを使ってLUKEのドメイン適応を行う方法を紹介しました。

ラベル付きデータが十分に用意できなくても、テキストだけであれば集めやすいと思います。

ぜひ一度お試しください。

 参考LUKE: Deep Contextualized Entity Representations with Entity-aware Self-attention
はじめに

使用技術・用語

LUKE

MLM (Masked Language Modeling)

ドメイン適応

実装

データの用意

データの加工

モデルの学習

おわりに

参考

Discussion
