Google ColabでModernBERTを動かしてみた

12月にBERTの発展版であるModernBERTが提案されています．このModernBERTはGPTベースのモデルとは異なり，軽量なパラメータ数であり，RAG(Retrieval Augmented Generation)などに利用されることが期待されます[1]．個人的にはBERTをよく触っていたので，エンコーダ型モデルの復活なるか...?!とワクワクしています．

今回の内容

• Google ColabのT4 GPU(無料枠)で動かす
• MASKが複数でも予測できるようにする

今回は上記の2つを行った記事になります．

コード

早速コードの説明をしていきます．
コード全体は以下のリンクから開くことができます．
https://github.com/kokist/modern-bert-sample

基本的に，実装のベースはこちらのHugging Faceのページを参考にしてます[2]．

環境構築

!pip install git+https://github.com/huggingface/transformers.git
!pip install triton

ここはtritonのインストールも必要でした．

Hugging FaceのページだとGPU環境がサポートされているなら，flash-attnもインストールすることを勧めています．しかし，これを試しましたが，T4 GPUだと以下のエラーが出ました．GPUのアーキテクチャ的に厳しそうなのでインストールを断念してます．

RuntimeError: FlashAttention only supports Ampere GPUs or newer.

MASK推論

ベースのコードだと，毎回モデルをダウンロードすることや，文に複数のMASK箇所がある場合に対応していなかったため書き換えをしています．

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

model_id = "answerdotai/ModernBERT-base"
tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForMaskedLM.from_pretrained(model_id, torch_dtype=torch.float16)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

print(device)
print(torch.cuda.get_device_name(0))  # GPUのモデル名を表示

テキストのみの変更で推論させるため以下を関数化

def predict_masked_tokens(text):
    """
    入力文に含まれるすべての[MASK]トークンの予測を行う。
    :param text: [MASK]トークンを含む文字列
    :return: 予測されたトークンのリスト
    """
    # 入力文のトークナイズ
    inputs = tokenizer(text, return_tensors="pt")
    inputs = {key: value.to(device) for key, value in inputs.items()}

    # 推論実行
    with torch.no_grad():
        with torch.autocast(device_type=device.type, dtype=torch.float16):
            outputs = model(**inputs)

    # 入力文中のすべての[MASK]トークンのインデックスを取得
    input_ids = inputs["input_ids"][0].tolist()
    mask_token_indices = [i for i, token_id in enumerate(input_ids) if token_id == tokenizer.mask_token_id]

    # 各[MASK]トークンに対する予測トークンを取得
    predicted_tokens = []
    for masked_index in mask_token_indices:
        logits = outputs.logits[0, masked_index]
        predicted_token_id = logits.argmax(axis=-1).item()
        predicted_token = tokenizer.decode(predicted_token_id)
        predicted_tokens.append(predicted_token)

    return predicted_tokens

予測

text1 = "The capital of France is [MASK]."
text2 = "The capital of France [MASK] [MASK]."
text3 = "No pain, [MASK] gain."
text4 = "No pain, no [MASK]."

print(f"入力: {text1} -> 予測: {predict_masked_tokens(text1)}")
print(f"入力: {text2} -> 予測: {predict_masked_tokens(text2)}")
print(f"入力: {text3} -> 予測: {predict_masked_tokens(text3)}")
print(f"入力: {text4} -> 予測: {predict_masked_tokens(text4)}")

予測結果

入力: The capital of France is [MASK]. -> Paris
入力: The capital of France [MASK] [MASK]. -> is, Paris
入力: No pain, [MASK] gain. -> no
入力: No pain, no [MASK]. -> gain

この予測結果を見る限りだとかなり良さそうでした．連続する複数MASKにも正解できていました．
なお，日本語でも試してみましたが，出力自体はできますが使える感じではなさそうでした．
入力: 石の上にも[MASK]年. -> 今
「石の上にも今年」 だそうです．

終わりに

ModernBERTの予測をColabで実行し，複数MASKでも試してみました．英語ではかなり精度が良さそうだったので，今後，日本語での学習済みモデルが出てくるのを期待ですね．
ModernBERTが「石の上にも今年」と言っていたので，1年，1年の意識なら辛抱強く何か続けられそうだなと思いました．それではお疲れ様でした～

参考

[1] ModernBERTの解説記事, https://huggingface.co/blog/modernbert
[2] コードの参考にしたHugging Faceのページ, https://huggingface.co/answerdotai/ModernBERT-base