はじめに
記事にするようなことでもないかもしれませんが、huggingfaceを使った実装に不慣れな場合、DatasetやDataLoaderの実装の仕方って場合によっては混乱するかなーと思いました（自分がそうだった）。
というのもTokenizerがバッチに関する処理（複数文章を同時にtokenizeしてくれたり、バッチ内のpaddingとかしてくれたり）をしてくれますが、実装の仕方によってはDatasetの段階で文章をTokenizerに通してid列にしておきたいケースもあります。その時のpaddingの処理はどこですべきか、など。
ここで紹介する実装例がベストプラクティスかどうかは全然わかりませんが、私の普段の実装例が誰かの参考になるかも、と思い記事にまとめておこうと思います。
明らかにおかしな実装してるぞこいつ、って感じのつっこみありましたらぜひご指摘ください。
実装例の紹介
準備
とりあえず日本語BERTモデルを使って、livedoorニュースコーパスを使ったカテゴリー分類のタスクを解くことを想定します。
以下のようにlivedoorニュースコーパスを手元に用意します。
Tokenizerも宣言しておきます。
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の処理をしておきたい場合
- Datasetの段階で入力文章1件1件をTokenizerに通してid化しておく
この段階ではpaddingはしない。メモリが非効率だし、huggingfaceが提供するDataCollatorを使えばバッチ毎にきれいにpaddingできるので。ただしモデルの最大長でtruncationはしておく - huggingfaceの
DataCollatorを使って、バッチ毎のpadding処理をDataLoaderのcollate_fnでする
まずはDatasetの実装
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の中身を見るとこんな感じ
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を活用しましょう。
上で実装したDatasetに対して、DataLoaderでバッチを取り出すとき、BERTのインプットに必要なinput_ids、attention_mask、token_type_idsをpaddingしてバッチ内で長さを揃えてからtensorに変換する必要がありますが、huggingfaceのDataCollatorを使えば一発で解決できます。
今回はただpaddingしたいだけなので、DataCollatorWithPaddingを使えばOKです。DataCollatorWithPaddingの中でtokenizer.padが呼ばれバッチ毎のpadding処理をしています。
MLMで事前学習したい、とかそんなときはDataCollatorForLanguageModelingを使えばいいですね。
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がされます。
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のソースコードを確認しましょう。
@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で保持したい場合
- Datasetの段階では、入力文章1件1件をそのままなにもせず格納しておく
- バッチ毎の
padding処理はTokenizerに任せる。そのため、DataLoaderのcollate_fnに渡す関数内でTokenizerを呼び出すようにする。
まずはDatasetの実装。こちらはcollate_fnの関数が実装しやすくなるように辞書型ではなく、特徴量と正解ラベルをタプル形式で返すようにしてます。（そんなことしなくても上と同じように辞書形式ですっきり実装する方法があるとしたら自分が知らないだけです...）
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の中身を確認します。
print(train_dataset[0])
# ('韓国料理といえば、キムチやサムゲタン、チゲ、キムパッなど、リーズナブルな値段で手軽に・・・',
# 3)
DataLoaderのcollate_fnに渡す関数を以下のように実装します。__call__でバッチ単位にまとめられた文章を受け取れるので、それを全部TokenizerでtokenizeすればOK。input_idsもattention_maskも当然バッチ毎にpaddingしてくれます。
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のアウトプットの形状と同じになります。
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を使い倒したい。
おわり