『Transformerによる自然言語処理』のRoBERTa事前訓練のコードを、データをhuggingface/datasetsで読み込むように書き直す

はじめに

今日も素振りにとりくーみこ！¹ nikkieです！

先日、『Transformerによる自然言語処理』の中のRoBERTaの事前訓練を写経したという記事を書きました：

"考えながら写経"していて、いくつか掘り下げたい事項が出てきています。
今回はデータの読み込みにフォーカスします。

今回解消する積み残し

datasetは、🤗的にはdatasetsを使ってロードする方法に置き換えたいようです
LineByLineTextDatasetはdeprecatedっぽい雰囲気なので、datasetsでの読み込みの仕方を調べたい（以上、写経記事より）

『Transformerによる自然言語処理』3章では、カントの著作のテキストデータを使いました。
テキストファイルとして用意したデータをLineByLineTextDatasetクラスのインスタンスとします。

dataset = LineByLineTextDataset(
    tokenizer=tokenizer, file_path="kant.txt", block_size=128
)

これをdatasetsというライブラリを使って書き換えてみます。

なお、この書き換えはLineByLineTextDatasetをインスタンス化したときにWarningで案内されています²。

This dataset will be removed from the library soon, preprocessing should be handled with the 🤗 Datasets library.
You can have a look at this example script for pointers: https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_mlm.py

参考例：examplesのlanguage-modeling/run_mlm.py

Warningで案内されたスクリプトを参考にします。

データの読み込みをたどると、大雑把に以下の流れと分かりました。

datasets.load_datasetで読み込み、raw_datasetsとする
raw_datasetsをtokenized_datasetsに変換（トークナイズ）
tokenized_datasetsは辞書のように扱える（DatasetDict）
- tokenized_datasets["train"]がtrain_dataset
- tokenized_datasets["validation"]がeval_dataset

動作環境

Colab (Python 3.7系)
transformers 4.18.0
tokenizers 0.12.1
datasets 2.1.0
torch 1.11.0+cu113

datasetsライブラリで書き換え

ドキュメントを引きつつ、書き換えたのがこちら！

token_dir = Path("KantaiBERT")
tokenizer = RobertaTokenizer.from_pretrained(str(token_dir), max_length=512)
text_column_name = "text"


def tokenize_function(examples):
    examples[text_column_name] = [
        line
        for line in examples[text_column_name]
        # 空行や空白文字だけからなる行を除くことでline by lineにしている
        if len(line) > 0 and not line.isspace()
    ]
    return tokenizer(
        examples[text_column_name],
        padding=False,
        truncation=True,
        max_length=512,  # from_pretrainedのmax_lengthと揃えた
        return_special_tokens_mask=True,
    )


raw_datasets = load_dataset("text", data_files="kant.txt")
tokenized_datasets = raw_datasets.map(
    tokenize_function,
    batched=True,
    num_proc=None,
    remove_columns=[text_column_name],
    load_from_cache_file=True,
    desc="Running tokenizer on dataset line_by_line",
)
dataset = tokenized_datasets["train"]

書き換え解説

そもそもdatasetsライブラリにおけるdatasetは、以下を含むディレクトリとのことです³。

some data files in generic formats (JSON, CSV, Parquet, text, etc.)

and optionally a dataset script, if it requires some code to read the data files. This is used to load any kind of formats or structures.

今回はテキストファイルの読み込みなので、load_datasetの第1引数に"text"を指定して読み込みます。
ref: https://huggingface.co/docs/datasets/loading#text-files

3章の範囲では、trainとvalidationのsplitは考えていないので、load_datasetのsplit引数は指定しません。
すると、load_datasetはDatasetDictを返します⁴。

返されたraw_datasets（DatasetDict）のキーを確認すると、trainだけを持ちます（繰り返しますが、trainとvalidationのsplitは考えていません）。
raw_datasets["train"]の各要素は辞書で、{"text": "kant.txtの1行"}という形式です（キーは自動でtextとなります）。
ここには空行も1要素として含んでいます。

DatasetDictはmapメソッドで各データセットを変換できます。
tokenizerを使ってトークナイズする関数を定義し、mapメソッドに渡します。
この関数でLineByLineTextDatasetの場合と同じ要素数に揃います。

トークナイズする関数のシグネチャは、function(batch: Dict[List]) -> Union[Dict, Any]となります。⁵
batched引数にTrueを指定していて、with_indices引数はデフォルト値のFalseとなるためです。

そして、mapメソッドが返したtokenized_datasetsからtrainのデータセットを取り出しました（tokenized_datasets["train"]の要素の形式については後述します）。

今回validationは用意しませんでしたが、上記のコードに少し手を入れるだけでvalidationのトークナイズもできるでしょう！（early stoppingを試したいと思っています）

書き換えて3章

上記の書き換えを使い、『Transformerによる自然言語処理』3章の内容を実施したnotebookはこちらです。

書き換えたことの検証

LineByLineTextDatasetからdatasetsライブラリを使うように書き換え、デグレていないかを確認しました。

書き換え前後でdatasetの長さが同じか
datasetの各要素について、トークナイズの結果（input_ids）が等しいか

from datasets import load_dataset
from transformers import LineByLineTextDataset, RobertaTokenizer

model_dir = "KantaiBERT"
text_column_name = "text"

tokenizer = RobertaTokenizer.from_pretrained(model_dir, max_length=512)


def tokenize_function(examples):
    examples[text_column_name] = [
        line
        for line in examples[text_column_name]
        if len(line) > 0 and not line.isspace()
    ]
    return tokenizer(
        examples[text_column_name],
        padding=False,
        truncation=True,
        max_length=512,
        return_special_tokens_mask=True,
    )


deprecated_dataset = LineByLineTextDataset(
    tokenizer=tokenizer, file_path="kant.txt", block_size=128
)

raw_datasets = load_dataset("text", data_files="kant.txt")
tokenized_datasets = raw_datasets.map(
    tokenize_function,
    batched=True,
    num_proc=None,
    remove_columns=[text_column_name],
    load_from_cache_file=True,
    desc="Running tokenizer on dataset line_by_line",
)
train_dataset = tokenized_datasets["train"]

assert len(deprecated_dataset) == len(train_dataset)

for old, new in zip(deprecated_dataset, train_dataset):
    assert old["input_ids"].tolist() == new["input_ids"]

LineByLineTextDatasetの要素は以下の辞書
- キー: input_ids
- input_idsの値はtorch.tensor
datasetsライブラリを使った実装では要素は以下の辞書
- キー: input_ids, special_tokens_mask, attention_mask
- input_idsの値はlist

スクリプトを実行したところ、AssertionErrorは送出されず、検証はパスしました。
デグレなしと言えると思います（検証項目の見落としに気づいた方はお知らせください）。
なお、書き換えた3章のコードを実行したところ、RoBERTaも事前訓練できていそうです✌️

終わりに

deprecatedと思われるLineByLineTextDatasetから、datasetsライブラリを使ったデータの読み込みに書き換えました！
主観でしかないですが、ちょっとだけモダンになった感じがします。

『Transformerによる自然言語処理』の訳者あとがきには、以下のようにあります。

（略）今後も技術発展とともに変更(※)が生じることが予想される。そのつもりで読んで対処していただきたい。(p.283)

(※)補足すると、コードやツールの変更

transformersを始めとするライブラリはたしかに変化が早いですね。
他の箇所でも新しい書き方ができることに気づいたら、またアウトプットしたいと思います。

Transformerによる自然言語処理

作者:Denis Rothman
朝倉書店

Amazon

川崎定演をご一緒した方には通じるかと思います↩
https://github.com/huggingface/transformers/blob/v4.18.0/src/transformers/data/datasets/language_modeling.py#L121-L126 ↩
https://huggingface.co/docs/datasets/v2.1.0/en/package_reference/loading_methods#datasets.load_dataset ↩
「if split is None, a datasets.DatasetDict with each split.」 ref: https://huggingface.co/docs/datasets/v2.1.0/en/package_reference/loading_methods#datasets.load_dataset そして、split引数のデフォルト値はNoneです↩
https://huggingface.co/docs/datasets/v2.1.0/en/package_reference/main_classes#datasets.DatasetDict.map ↩

nikkie-ftnextの日記

イベントレポートや読書メモを発信