概要

前回の続き

Github に pretrain のソースコードを上げた。
https://github.com/elm200/bert_pretrainer/blob/master/pretrain.py

蛇足かもしれないが、このソースコードに解説を加えていく。
HuggingFace Transformers のおかげで、比較的簡単に実装できる。
Tokenizer としては、東北大日本語BERTモデルのものを使用させていただいている。

メイン部分(抄)

Next Sentence Prediction (NSP) をするかどうかの選択

BertForPreTrainingWithoutNSP の実体は、transformers からコピーした modeling_bert.py の BertForPreTraining クラスであり、この中の forward メソッドの一部を書き換えている。具体的には、NSPの部分を loss から外している。
実験した結果 NSP が悪影響を及ぼしていることがあったので、こういうオプションを追加した。

メインループ

create_sent_pairs() で BATCH_SIZE 文の文ペアのリストを作り、encode_sent_pairs() で token id にエンコード。あとは、普通に loss を逆伝播しているだけ。

    if USE_NSP:
        model = BertForPreTraining(config)
    else:
        model = BertForPreTrainingWithoutNSP(config)
    model.to(device)

    optimizer = AdamW(model.parameters(), lr=2e-5)
    model.train()

    for i in range(1, MAX_STEPS + 1):
        optimizer.zero_grad()
        sent_pairs = create_sent_pairs(sents_list, batch_size=BATCH_SIZE)
        encoded = encode_sent_pairs(sent_pairs)
        res = model(
            encoded["input_ids"].to(device),
            token_type_ids=None,
            attention_mask=encoded["attention_mask"].to(device),
            labels=encoded["labels"].to(device),
            next_sentence_label=encoded["next_sentence_label"].to(device),
        )
        loss = res.loss       
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

create_sent_pairs()

sents_listの各要素には日本語 Wikipedia の1つの項目の各文がJSON配列の形で入っている（このデータの作り方は後日説明予定）。seq_pair_ratioの確率で、出力される2文は連続したものとなり、そうでないときは、全く無関係の2文が選ばれる。これは Next Sentence Prediction タスクのための工夫である。

def create_sent_pairs(
    sents_list: list[str], batch_size: int, seq_pair_ratio: float = 0.5
) -> list[tuple[str, str, int]]:
    n_seq_pair = int(batch_size * seq_pair_ratio)
    n_random_pair = batch_size - n_seq_pair
    pairs = []

    for i in range(n_seq_pair):
        while True:
            sents = json.loads(random.choice(sents_list))
            if len(sents) >= 2:
                break
        st = random.randint(0, len(sents) - 2) if len(sents) > 2 else 0
        pair = tuple(sents[st : st + 2]) + (0,)
        pairs.append(pair)

    for i in range(n_random_pair):
        ss1, ss2 = random.sample(sents_list, 2)
        sents1 = json.loads(ss1)
        sents2 = json.loads(ss2)
        pair = (random.choice(sents1), random.choice(sents2), 1)
        pairs.append(pair)

    pairs = random.sample(pairs, len(pairs))
    return pairs

mask_token_ids_and_labels()

Masked Language Model(MLM) タスクのための教師データを自動的に作るメソッド。
トークン1つごとに確率 mask_rate で入力がマスク化され、さらにそのマスク化されたトークンが

80%の確率で[MASK]に置換
10%の確率でランダムに他のトークンに置換
10%の確率でそのまま

となる。 MLM では、マスク化されたトークンに関してのみ loss を集計するようにする。

def mask_token_ids_and_labels(
    token_ids: list[int], mask_rate: float = 0.15
) -> tuple[list[int], list[int]]:
    indexes = []
    for i in range(len(token_ids)):
        r = random.random()
        if r < mask_rate:
            indexes.append(i)
    labels = [IGNORE_ID] * len(token_ids)
    if len(indexes) == 0:
        return token_ids, labels
    labels = np.array(labels)
    indexes = np.array(indexes)
    token_ids = np.array(token_ids)
    masked_token_ids = np.array(token_ids)
    labels[indexes] = token_ids[indexes]
    masked_token_ids[indexes] = MASK_ID
    for i in indexes:
        r = random.random()
        if r < 0.1:
            masked_token_ids[i] = token_ids[i]
        elif r < 0.2:
            masked_token_ids[i] = random.choice(token_ids)
    return masked_token_ids.tolist(), labels.tolist()

次回予告

Wikipedia のダンプからどうやって pretrain 用のデータを作るかということを解説したい。