huggingface / transformersにおいてDatasetの段階ではstringのままにしておきたい場合とTokenizerで処理した状態にしておきたい場合のDatasetとDataLoaderの実装例

# はじめに

記事にするようなことでもないかもしれませんが、huggingfaceを使った実装に不慣れな場合、DatasetやDataLoaderの実装の仕方って場合によっては混乱するかなーと思いました（自分がそうだった）。
というのもTokenizerがバッチに関する処理（複数文章を同時にtokenizeしてくれたり、バッチ内のpaddingとかしてくれたり）をしてくれますが、実装の仕方によってはDatasetの段階で文章をTokenizerに通してid列にしておきたいケースもあります。その時のpaddingの処理はどこですべきか、など。

ここで紹介する実装例がベストプラクティスかどうかは全然わかりませんが、私の普段の実装例が誰かの参考になるかも、と思い記事にまとめておこうと思います。

明らかにおかしな実装してるぞこいつ、って感じのつっこみありましたらぜひご指摘ください。

# 実装例の紹介

## 準備

とりあえず日本語BERTモデルを使って、livedoorニュースコーパスを使ったカテゴリー分類のタスクを解くことを想定します。
以下のようにlivedoorニュースコーパスを手元に用意します。
Tokenizerも宣言しておきます。

```python
import pickle
import pandas as pd
from sklearn.model_selection import train_test_split
from transformers import AutoTokenizer

# 東北大BERTのTokenizerをお借りします。
tokenizer = AutoTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')

# 事前に手元にカテゴリー名、タイトル、本文でまとめたlivedoorニュースコーパスのDataFrameを用意しておく。
with open('./livedoor_data.pickle', 'rb') as r:
    df = pickle.load(r)

# カテゴリー名をIDに変換した列を追加
categories = df['category'].unique().tolist()
category2id = {k:categories.index(k) for k in categories}
df['category_id'] = df['category'].map(lambda x: category2id[x])

# 学習データとテストデータに分ける
train_df, test_df = train_test_split(df, train_size=0.8)
print('train size', train_df.shape)
print('test size', test_df.shape)
train_df.sample(3)

# train size (5900, 4)
# test size (1476, 4)
# 	category	title	body	category_id
# 2707	topic-news	大学ブランド偏差値、首都圏トップ30を発表	9日、日経BPコンサルティングは、今年8月に実施した「大学ブランド・イメージ調査 2011-...	3
# 523	movie-enter	ルーニー・マーラ、可憐さでセクシーなドレスで『ドラゴン・タトゥーの女』プレミアに登場	『セブン』『ソーシャル・ネットワーク』の鬼才デヴィド・フィンチャーの最新作『ドラゴン・タトゥ...	0
# 3960	peachy	LINEにも登場！資生堂「ワタシプラス」でスキンケアサンプルをプレゼント	「私の肌には、どんな化粧品を使ったらいいの？」「私の顔には、どんなメイクが似合うんだろう…」...	5
```

## Case1. Datasetの段階でTokenizerの処理をしておきたい場合

1. Datasetの段階で入力文章1件1件をTokenizerに通してid化しておく  
この段階では`padding`はしない。メモリが非効率だし、huggingfaceが提供する`DataCollator`を使えばバッチ毎にきれいに`padding`できるので。ただしモデルの最大長で`truncation`はしておく
1. huggingfaceの`DataCollator`を使って、バッチ毎の`padding`処理をDataLoaderの`collate_fn`でする


まずはDatasetの実装

```python
class SampleDataset(Dataset):
    def __init__(self, df, tokenizer, max_length=512):
        self.features = []
        for row in tqdm(df.itertuples(), total=df.shape[0]):
            text = row.body
            category_id = row.category_id
            # Tokenizerの引数のreturn_tensors='pt'を指定しちゃうと
            # 戻り値のサイズが(1, 文章の長さ)と先頭にバッチサイズ1が入っちゃうので、
            # return_tensorsは指定せず配列で保持するようにしておく。
            encoding = tokenizer(text, truncation=True, max_length=max_length)
            # 上記はattention_maskやtoken_type_idsも含まれていますが、
            # input_idsだけにしたい場合はこんな感じにしておけばOK
            # encoding = {'input_ids': tokenizer(text, truncation=True, max_length=max_length)['input_ids']}
            
            encoding['category_ids'] = category_id
            self.features.append(encoding)
    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        return self.features[idx]

train_dataset = SampleDataset(train_df, tokenizer)
test_dataset = SampleDataset(test_df, tokenizer)
```

Datasetの中身を見るとこんな感じ

```python
print(train_dataset[0].keys())
print(train_dataset[0])
# dict_keys(['input_ids', 'token_type_ids', 'attention_mask', 'category_ids'])
# {'input_ids': [2, 775, 32, 6, 57, 10374, 51, 12, 265, 9, 76, 247, 13, 7736, 15, 16, 33, 15235,...
# 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,...
# 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,...
# 'category_ids': 6}
```


DataLoaderの実装はhuggingfaceの`DataCollator`を活用しましょう。

- https://huggingface.co/docs/transformers/main_classes/data_collator

上で実装したDatasetに対して、DataLoaderでバッチを取り出すとき、BERTのインプットに必要な`input_ids`、`attention_mask`、`token_type_ids`をpaddingしてバッチ内で長さを揃えてからtensorに変換する必要がありますが、huggingfaceの`DataCollator`を使えば一発で解決できます。
今回はただ`padding`したいだけなので、`DataCollatorWithPadding`を使えばOKです。`DataCollatorWithPadding`の中で`tokenizer.pad`が呼ばれバッチ毎の`padding`処理をしています。
MLMで事前学習したい、とかそんなときは`DataCollatorForLanguageModeling`を使えばいいですね。


```python
from transformers import DataCollatorWithPadding
padding_collator = DataCollatorWithPadding(tokenizer)
train_loader = DataLoader(train_dataset, collate_fn=padding_collator, batch_size=8, shuffle=True)
test_loader = DataLoader(test_dataset, collate_fn=padding_collator, batch_size=8, shuffle=False)
```

DataLoaderから取り出されるバッチを確認しましょう。tensorの形状がきれいに揃っていることが確認できます。
下は`max_length`の512になってますが、バッチ内のデータの最大長が512未満であればそのサイズまでで`padding`がされます。

```python
batch = next(iter(train_loader))
print(batch.keys())
print(batch['input_ids'].shape)
print(batch['attention_mask'].shape)
print(batch['token_type_ids'].shape)
print(batch['category_ids'].shape)

# dict_keys(['input_ids', 'token_type_ids', 'attention_mask', 'category_ids'])
# torch.Size([8, 512])
# torch.Size([8, 512])
# torch.Size([8, 512])
# torch.Size([8])
```

### (補足)huggingfaceのcollatorは`label`もしくは`label_ids`という名前のキーを`labels`に変換します

上で使った`DataCollatorWithPadding`のソースコードを確認しましょう。

```python
@dataclass
class DataCollatorWithPadding:
# コメントは省略しました。
    tokenizer: PreTrainedTokenizerBase
    padding: Union[bool, str, PaddingStrategy] = True
    max_length: Optional[int] = None
    pad_to_multiple_of: Optional[int] = None
    return_tensors: str = "pt"

    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
        batch = self.tokenizer.pad(
            features,
            padding=self.padding,
            max_length=self.max_length,
            pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors=self.return_tensors,
        )
        if "label" in batch:
            batch["labels"] = batch["label"]
            del batch["label"]
        if "label_ids" in batch:
            batch["labels"] = batch["label_ids"]
            del batch["label_ids"]
        return batch
```

なんでこんなことするんだ、って話ですが、huggingfaceのTrainerクラスとかが`labels`という名前のキーを参照したりするし、ライブラリー全体の整合性を取るためかなーと想像してます。


## Case2. Datasetの段階ではstringで保持したい場合

1. Datasetの段階では、入力文章1件1件をそのままなにもせず格納しておく
1. バッチ毎の`padding`処理はTokenizerに任せる。そのため、DataLoaderの`collate_fn`に渡す関数内でTokenizerを呼び出すようにする。

まずはDatasetの実装。こちらは`collate_fn`の関数が実装しやすくなるように辞書型ではなく、特徴量と正解ラベルをタプル形式で返すようにしてます。（そんなことしなくても上と同じように辞書形式ですっきり実装する方法があるとしたら自分が知らないだけです...）

```python
class SampleDataset(Dataset):
    def __init__(self, df, tokenizer, max_length=512):
        self.texts = []
        self.category_ids = []
        for row in tqdm(df.itertuples(), total=df.shape[0]):
            text = row.body
            label = row.category_id
            self.texts.append(text)
            self.category_ids.append(label)
            
    def __len__(self):
        return len(self.category_ids)

    def __getitem__(self, idx):
        return self.texts[idx], self.category_ids[idx]

train_dataset = SampleDataset(train_df, tokenizer)
test_dataset = SampleDataset(test_df, tokenizer)
```

Datasetの中身を確認します。

```python
print(train_dataset[0])
# ('韓国料理といえば、キムチやサムゲタン、チゲ、キムパッなど、リーズナブルな値段で手軽に・・・',
#  3)
```

DataLoaderの`collate_fn`に渡す関数を以下のように実装します。`__call__`でバッチ単位にまとめられた文章を受け取れるので、それを全部TokenizerでtokenizeすればOK。`input_ids`も`attention_mask`も当然バッチ毎に`padding`してくれます。

```python
class SampleCollator():
    def __init__(self, tokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __call__(self, examples):
        texts, category_ids = list(zip(*examples))
        encoding = self.tokenizer(texts, padding=True, truncation=True, max_length=self.max_length, return_tensors='pt')
        encoding['category_ids'] = torch.tensor(category_ids)
        return encoding

sample_collator = SampleCollator(tokenizer)
train_loader = DataLoader(train_dataset, collate_fn=sample_collator, batch_size=8, shuffle=True)
test_loader = DataLoader(test_dataset, collate_fn=sample_collator, batch_size=8, shuffle=False)
```

もちろんCase1.で見たDataLoaderのアウトプットの形状と同じになります。

```python
batch = next(iter(train_loader))
print(batch.keys())
print(batch['input_ids'].shape)
print(batch['attention_mask'].shape)
print(batch['token_type_ids'].shape)
print(batch['category_ids'].shape)

# dict_keys(['input_ids', 'token_type_ids', 'attention_mask', 'category_ids'])
# torch.Size([8, 512])
# torch.Size([8, 512])
# torch.Size([8, 512])
# torch.Size([8])
```


# おわりに

huggingfaceを使った実装ってhuggingfaceならではの実装の仕方になりがちな気がするけど、使いこなすと色々と凄い便利なんだろうなーと思っている次第です。
huggingfaceを使い倒したい。

おわり