BERTを用いたベクトル表現の比較（bank : 銀行と土手）

BERTの「文脈に応じた単語表現を獲得する特徴」を確認するために、bank（銀行）とbank（土手）の単語ベクトル表現を比較する。

BERTを最初から学習させるのではなく、学習済みモデルをロードする。

https://github.com/huggingface/transformers

学習済みモデルのロード

# 学習済みモデルの重みをロード
weights_path = "./weights/pytorch_model.bin"
loaded_state_dict = torch.load(weights_path)

# モデルの用意
net = BertModel(config) # インスタンス作成
net.eval() # netを評価モードに設定

param_names = [] 
# ネットワークのすべてのパラメータとその名前を取得
for name, param in net.named_parameters():
    param_names.append(name)

BERT用のTokenizerの実装

BERTのボキャブラリーファイルを読み込んで、その内容を辞書に格納するための関数`load_vocab`を定義

2つの辞書型変数にボキャブラリーファイルの内容を格納

vocab: 単語をキーとし、対応するIDを値とする
ids_to_tokens: IDをキーとし、対応する単語を値とする

def load_vocab(vocab_file):
    """text形式のvocabファイルの内容を辞書に格納します"""
    vocab = collections.OrderedDict()  # (単語, id)の順番の辞書変数
    ids_to_tokens = collections.OrderedDict()  # (id, 単語)の順番の辞書変数
    index = 0

    with open(vocab_file, "r", encoding="utf-8") as reader:
        while True:
            token = reader.readline()
            if not token:
                break
            token = token.strip()

            # 格納
            vocab[token] = index
            ids_to_tokens[index] = token
            index += 1

    return vocab, ids_to_tokens

単語分割を実行するクラス`BertTokenizer`を定義

tokenize : 文章をサブワードに分割する関数
input : テキスト (text)
output : 分割された単語が入ったリスト (split_tokens)
convert_tokens_to_ids : 分割された単語リストをIDに変換する関数
input : 分割された単語リスト (tokens)
output : 対応するIDのリスト (ids)
convert_ids_to_tokens : IDを単語に変換する関数
input : IDのリスト (ids)
output : 対応する単語のリスト (tokens)

from utils.tokenizer import BasicTokenizer, WordpieceTokenizer

# BasicTokenizer, WordpieceTokenizerは↓
# https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/pytorch_pretrained_bert/tokenization.py
# これらはsub-wordで単語分割を行うクラス

class BertTokenizer(object):
    '''BERT用の文章の単語分割クラスを実装'''

    def __init__(self, vocab_file, do_lower_case=True):
        '''
        vocab_file：ボキャブラリーへのパス
        do_lower_case：前処理で単語を小文字化するかどうか
        '''

        # ボキャブラリーのロード
        self.vocab, self.ids_to_tokens = load_vocab(vocab_file)

        # 分割処理の関数をフォルダ「utils」からimoprt、sub-wordで単語分割を行う
        never_split = ("[UNK]", "[SEP]", "[PAD]", "[CLS]", "[MASK]")
        # (注釈)上記の単語は途中で分割させない。これで一つの単語とみなす

        self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case,
                                              never_split=never_split)
        self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab)

    def tokenize(self, text):
        '''文章を単語に分割する関数'''
        split_tokens = []  # 分割後の単語たち
        for token in self.basic_tokenizer.tokenize(text):
            for sub_token in self.wordpiece_tokenizer.tokenize(token):
                split_tokens.append(sub_token)
        return split_tokens

    def convert_tokens_to_ids(self, tokens):
        """分割された単語リストをIDに変換する関数"""
        ids = []
				# 単語ごとにself.vocabを参照し、単語に対応するIDを取得し、idsリストに追加
        for token in tokens:
            ids.append(self.vocab[token])

        return ids

    def convert_ids_to_tokens(self, ids):
        """IDを単語に変換する関数"""
        tokens = []
        for i in ids:
            tokens.append(self.ids_to_tokens[i])
        return tokens

サブワードによるトークン化

サブワードは単語よりも小さな部分文字列

利点

サブワードを使用することで、限られた数の語彙を効率的に扱うことができる
未知語を既存のサブワードに分解し、対応できる
例）"annoyingly" を "annoying" と "ly" に分解されるとか
語形変化への柔軟な対応ができる
例）"running"や"ran"などの単語を、共通のサブワード（例えば、"run"）

サブワードによるトークン化、WordPieceについてはこちらでわかりやすく説明されています

https://note.com/npaka/n/nb08941a36c8b#I3MMH

Bankの文脈による意味変化を単語ベクトルとして求める

まず、文章を入力し単語を分割する。

# 文章1：銀行口座にアクセスしました。
text_1 = "[CLS] I accessed the bank account. [SEP]"

# 文章2：彼は敷金を銀行口座に振り込みました。
text_2 = "[CLS] He transferred the deposit money into the bank account. [SEP]"

# 文章3：川岸でサッカーをします。
text_3 = "[CLS] We play soccer at the bank of the river. [SEP]"

# 単語分割Tokenizerを用意
tokenizer = BertTokenizer(
    vocab_file="./vocab/bert-base-uncased-vocab.txt", do_lower_case=True)

# 文章を単語分割
tokenized_text_1 = tokenizer.tokenize(text_1)
tokenized_text_2 = tokenizer.tokenize(text_2)
tokenized_text_3 = tokenizer.tokenize(text_3)

print(tokenized_text_1)
# ['[CLS]', 'i', 'accessed', 'the', 'bank', 'account', '.', '[SEP]']

次に単語をIDに変換する。

# 単語をIDに変換する
indexed_tokens_1 = tokenizer.convert_tokens_to_ids(tokenized_text_1)
indexed_tokens_2 = tokenizer.convert_tokens_to_ids(tokenized_text_2)
indexed_tokens_3 = tokenizer.convert_tokens_to_ids(tokenized_text_3)

# 各文章のbankの位置を特定
bank_posi_1 = np.where(np.array(tokenized_text_1) == "bank")[0][0]  # 4
bank_posi_2 = np.where(np.array(tokenized_text_2) == "bank")[0][0]  # 8
bank_posi_3 = np.where(np.array(tokenized_text_3) == "bank")[0][0]  # 6

# リストをPyTorchのテンソルに
tokens_tensor_1 = torch.tensor([indexed_tokens_1])
tokens_tensor_2 = torch.tensor([indexed_tokens_2])
tokens_tensor_3 = torch.tensor([indexed_tokens_3])

# bankの単語id
bank_word_id = tokenizer.convert_tokens_to_ids(["bank"])[0]

print(tokens_tensor_1)
# tensor([[  101,  1045, 11570,  1996,  2924,  4070,  1012,   102]])

学習済みモデルをロードしたBERTに入力し、推論する。

# 文章をBERTに入力し、エンコードされたレイヤーを取得
with torch.no_grad():
    encoded_layers_1, _ = net(tokens_tensor_1, output_all_encoded_layers=True)
    encoded_layers_2, _ = net(tokens_tensor_2, output_all_encoded_layers=True)
    encoded_layers_3, _ = net(tokens_tensor_3, output_all_encoded_layers=True)

"""
torch.no_grad(): 勾配の計算を無効にするためのコンテキストマネージャ
これにより、モデルのパラメータが更新されないようになり、メモリ使用量が削減される

output_all_encoded_layers=True: 
各エンコーダーレイヤーの出力をすべて取得するためのオプション
"""

各文章について、1段目と12段目のBertLayerモジュールから出力される、単語bankの位置に対応する特徴量ベクトルを取り出す。

# bankの初期の単語ベクトル表現
# これはEmbeddingsモジュールから取り出し、単語bankのidに応じた単語ベクトルなので3文で共通
bank_vector_0 = net.embeddings.word_embeddings.weight[bank_word_id]

# 文章1のBertLayerモジュール1段目から出力されるbankの特徴量ベクトル
bank_vector_1_1 = encoded_layers_1[0][0, bank_posi_1]

# 文章1のBertLayerモジュール最終12段目から出力されるのbankの特徴量ベクトル
bank_vector_1_12 = encoded_layers_1[11][0, bank_posi_1]

# 文章2、3も同様に
bank_vector_2_1 = encoded_layers_2[0][0, bank_posi_2]
bank_vector_2_12 = encoded_layers_2[11][0, bank_posi_2]
bank_vector_3_1 = encoded_layers_3[0][0, bank_posi_3]
bank_vector_3_12 = encoded_layers_3[11][0, bank_posi_3]

取り出した単語ベクトル表現について、コサイン類似度を計算する。

# コサイン類似度を計算
import torch.nn.functional as F

print("bankの初期ベクトル と 文章1の1段目のbankの類似度：",
      F.cosine_similarity(bank_vector_0, bank_vector_1_1, dim=0))
print("bankの初期ベクトル と 文章1の12段目のbankの類似度：",
      F.cosine_similarity(bank_vector_0, bank_vector_1_12, dim=0))

print("文章1の1層目のbank と 文章2の1段目のbankの類似度：",
      F.cosine_similarity(bank_vector_1_1, bank_vector_2_1, dim=0))
print("文章1の1層目のbank と 文章3の1段目のbankの類似度：",
      F.cosine_similarity(bank_vector_1_1, bank_vector_3_1, dim=0))

print("文章1の12層目のbank と 文章2の12段目のbankの類似度：",
      F.cosine_similarity(bank_vector_1_12, bank_vector_2_12, dim=0))
print("文章1の12層目のbank と 文章3の12段目のbankの類似度：",
      F.cosine_similarity(bank_vector_1_12, bank_vector_3_12, dim=0))

# 出力 (値が1に近いほど類似している)

bankの初期ベクトル と 文章1の1段目のbankの類似度： tensor(0.6814, grad_fn=<DivBackward0>)
bankの初期ベクトル と 文章1の12段目のbankの類似度： tensor(0.2276, grad_fn=<DivBackward0>)
# BertLayerの繰り返しで単語ベクトルの表現が変化していることがわかる。

文章1の1段目のbank と 文章2の1段目のbankの類似度： tensor(0.8968)
文章1の1段目のbank と 文章3の1段目のbankの類似度： tensor(0.7584)
# 1段目では差はそれほどない

文章1の12段目のbank と 文章2の12段目のbankの類似度： tensor(0.8796)
文章1の12段目のbank と 文章3の12段目のbankの類似度： tensor(0.4814)
# 12段目から出力される単語ベクトルだと、
# 1と2は類似しているが、1と3はあまり類似していないことがわかる。

BERTではBertLayerを12段通過する間に、同じ単語でも文章内の周辺の単語の情報をSelf-Attentionで拾ってきて演算処理し、文脈に応じたベクトル表現を獲得して出力していることが確認できる。

BERTの学習・推論、判定根拠の可視化を実装

映画レビューデータのポジ・ネガを判定する感情分析モデルをBERTを用いて構築して、学習・推論を行う。推論時にSelf-Attentionの重みを可視化し、どのような単語に着目して推論しているか確認する。

IMDbデータを読み込み、DataLoaderを作成（BERTのTokenizerを使用）

まず、文章の前処理と単語分割を行う関数を実装

# 前処理と単語分割をまとめた関数を作成
import re
import string
from utils.bert import BertTokenizer

def preprocessing_text(text):
    '''IMDbの前処理'''
    # 改行コードを消去
    text = re.sub('<br />', '', text)

    # カンマ、ピリオド以外の記号をスペースに置換
    for p in string.punctuation:
        if (p == ".") or (p == ","):
            continue
        else:
            text = text.replace(p, " ")

    # ピリオドなどの前後にはスペースを入れておく
    text = text.replace(".", " . ")
    text = text.replace(",", " , ")
    return text

# 単語分割用のTokenizerを用意
tokenizer_bert = BertTokenizer(
    vocab_file="./vocab/bert-base-uncased-vocab.txt", do_lower_case=True)

# 前処理と単語分割をまとめた関数を定義
# 単語分割の関数を渡すので、tokenizer_bertではなく、tokenizer_bert.tokenizeを渡す
def tokenizer_with_preprocessing(text, tokenizer=tokenizer_bert.tokenize):
    text = preprocessing_text(text)
    ret = tokenizer(text)  # tokenizer_bert
    return ret
# 分割された単語のリストを返す

データを読み込んだときに、読み込んだ内容に対して行う処理を定義

# 256単語に設定
max_length = 256

TEXT = torchtext.data.Field(sequential=True, tokenize=tokenizer_with_preprocessing, use_vocab=True,
                            lower=True, include_lengths=True, batch_first=True, fix_length=max_length, init_token="[CLS]", eos_token="[SEP]", pad_token='[PAD]', unk_token='[UNK]')
LABEL = torchtext.data.Field(sequential=False, use_vocab=False)

# sequential: データの長さが可変か？文章は長さがいろいろなのでTrue.ラベルはFalse
# tokenize: 文章を読み込んだときに、前処理や単語分割をするための関数を定義
# use_vocab：単語をボキャブラリーに追加するかどうか
# lower：アルファベットがあったときに小文字に変換するかどうか
# include_length: 文章の単語数のデータを保持するか
# batch_first：ミニバッチの次元を先頭に用意するかどうか
# fix_length：全部の文章を指定した長さと同じになるように、paddingします
# init_token, eos_token, pad_token, unk_token：文頭、文末、padding、未知語に対して、
# どんな単語を与えるかを指定

Datasetを作成

# IMDbを読み込む
train_val_ds, test_ds = torchtext.data.TabularDataset.splits(
    path='./data/', train='IMDb_train.tsv',
    test='IMDb_test.tsv', format='tsv',
    fields=[('Text', TEXT), ('Label', LABEL)])

# torchtext.data.Datasetのsplit関数で訓練データとvalidationデータを分ける
# 80%のデータをトレーニングデータとして使用し、残りの20%をバリデーションデータとして使用
train_ds, val_ds = train_val_ds.split(
    split_ratio=0.8, random_state=random.seed(1234))

学習時に使用するボキャブラリーを準備

# BERTが持つ全単語でBertEmbeddingモジュールを作成しているので、
# ボキャブラリーとしては全単語を使用
# そのため訓練データからボキャブラリーは作成しない

from utils.bert import BertTokenizer, load_vocab

# BERT用の単語辞書を辞書型変数に格納
vocab_bert, ids_to_tokens_bert = load_vocab(
    vocab_file="./vocab/bert-base-uncased-vocab.txt")

# TEXTフィールドのボキャブラリーを構築
# min_freq=1: 最低頻度の制限を設けずにすべての単語をボキャブラリーに含める
TEXT.build_vocab(train_ds, min_freq=1)
TEXT.vocab.stoi = vocab_bert
# IDと単語を紐づける辞書が得られた

DataLoaderを作成

# DataLoaderを作成
batch_size = 32  # BERTでは16、32あたりを使用する

# 訓練データのDataLoader
# torchtext.data.Iteratorを使用してデータをバッチ化
train_dl = torchtext.data.Iterator(
    train_ds, batch_size=batch_size, train=True)

# バリデーションデータのDataLoader
val_dl = torchtext.data.Iterator(
    val_ds, batch_size=batch_size, train=False, sort=False)

# テストデータのDataLoader
test_dl = torchtext.data.Iterator(
    test_ds, batch_size=batch_size, train=False, sort=False)

# 辞書オブジェクトにまとめる
dataloaders_dict = {"train": train_dl, "val": val_dl}

"""
train=True
訓練データを対象としてデータローダーを作成
この場合、データはシャッフルされ、バッチ化され、必要な前処理（例：パディング）が適用される

train=False
検証データやテストデータなどの訓練以外のデータを対象としてデータローダーを作成
データはシャッフルされず、バッチ化されるだけで、前処理は適用されない
"""

感情分析用のBERTモデルを構築

BERTモデルに対して、学習済みパラメータをロードし、ポジ・ネガ分類用のアダプターモジュールを取り付けることで、感情分析用BERTモデルを構築する。

BERTモデルを構築し、学習済みのパラメータを読み込む

from utils.bert import get_config, BertModel, set_learned_params

# モデル設定のJOSNファイルをオブジェクト変数として読み込む
config = get_config(file_path="./weights/bert_config.json")

# BERTモデルを作成
net_bert = BertModel(config)

# BERTモデルに学習済みパラメータをセット
net_bert = set_learned_params(
    net_bert, weights_path="./weights/pytorch_model.bin")
# net_bertには学習済みのBERTモデルが格納される

IMDbのポジティブ・ネガティブを判定するためのBertForIMDbモデルを定義

class BertForIMDb(nn.Module):
    '''BERTモデルにIMDbのポジ・ネガを判定する部分をつなげたモデル'''

    def __init__(self, net_bert):
        super(BertForIMDb, self).__init__()

        # BERTモデル
				# BERTの基本モデル部分を取得し、入力データを処理
        self.bert = net_bert

        # ポジネガ予測（全結合層）
        # 入力はBERTの出力特徴量の次元、出力はポジ・ネガの2つ
        self.cls = nn.Linear(in_features=768, out_features=2)

        # 重み初期化
        nn.init.normal_(self.cls.weight, std=0.02)
        nn.init.normal_(self.cls.bias, 0)

    def forward(self, input_ids, token_type_ids=None, attention_mask=None, output_all_encoded_layers=False, attention_show_flg=False):
        '''
        input_ids： [batch_size, sequence_length]の文章の単語IDの羅列
        token_type_ids： [batch_size, sequence_length]の、各単語が1文目なのか、2文目なのかを示すid
        attention_mask：Transformerのマスクと同じ働きのマスキングです
        output_all_encoded_layers：最終出力に12段のTransformerの全部をリストで返すか、最後だけかを指定
        attention_show_flg：Self-Attentionの重みを返すかのフラグ
        '''

				"""
	        BERTの基本モデル部分の順伝搬

	        BERTモデルに入力データを渡し、
					エンコードされた各レイヤーの特徴量と最終的なプールされた出力を取得
        """
				if attention_show_flg == True:
            '''attention_showのときは、attention_probsもリターンする'''
            encoded_layers, pooled_output, attention_probs = self.bert(
                input_ids, token_type_ids, attention_mask, output_all_encoded_layers, attention_show_flg)
        elif attention_show_flg == False:
            encoded_layers, pooled_output = self.bert(
                input_ids, token_type_ids, attention_mask, output_all_encoded_layers, attention_show_flg)

        # 入力文章の1単語目[CLS]の特徴量を使用して、ポジ・ネガを分類します
        vec_0 = encoded_layers[:, 0, :]
        vec_0 = vec_0.view(-1, 768)  # sizeを[batch_size, hidden_sizeに変換
        out = self.cls(vec_0)

        # attention_showのときは、attention_probs（1番最後の）もリターンする
        if attention_show_flg == True:
            return out, attention_probs
        elif attention_show_flg == False:
            return out

モデル構築

net = BertForIMDb(net_bert)

# 訓練モードに設定
net.train()

ファインチューニング用設定

# 勾配計算を最後のBertLayerモジュールと追加した分類アダプターのみ実行

# 1. まず全部を、勾配計算Falseにしてしまう
for param in net.parameters():
    param.requires_grad = False

# 2. 最後のBertLayerモジュールを勾配計算ありに変更
for name, param in net.bert.encoder.layer[-1].named_parameters():
    param.requires_grad = True

# 3. 識別器を勾配計算ありに変更
for name, param in net.cls.named_parameters():
    param.requires_grad = True

# 最適化手法の設定
# 設定はBERTの論文で推奨されているものを使用
optimizer = optim.Adam([
    {'params': net.bert.encoder.layer[-1].parameters(), 'lr': 5e-5},
    {'params': net.cls.parameters(), 'lr': 5e-5}
], betas=(0.9, 0.999))

# 損失関数の設定
criterion = nn.CrossEntropyLoss()

学習・検証

# モデルを学習させる関数を作成
def train_model(net, dataloaders_dict, criterion, optimizer, num_epochs):

    # ミニバッチのサイズ
    batch_size = dataloaders_dict["train"].batch_size

    # epochのループ
    for epoch in range(num_epochs):
        # epochごとの訓練と検証のループ
        for phase in ['train', 'val']:
            if phase == 'train':
                net.train()  # モデルを訓練モードに
            else:
                net.eval()   # モデルを検証モードに

            epoch_loss = 0.0  # epochの損失和
            epoch_corrects = 0  # epochの正解数
            iteration = 1

            # 開始時刻を保存
            t_epoch_start = time.time()
            t_iter_start = time.time()

            # データローダーからミニバッチを取り出すループ
            for batch in (dataloaders_dict[phase]):
                # batchはTextとLableの辞書型変数

                # GPUが使えるならGPUにデータを送る
                inputs = batch.Text[0].to(device)  # 文章
                labels = batch.Label.to(device)  # ラベル

                # optimizerを初期化
                optimizer.zero_grad()

                # 順伝搬（forward）計算
                with torch.set_grad_enabled(phase == 'train'):

                    # BertForIMDbに入力
                    outputs = net(inputs, token_type_ids=None, attention_mask=None,
                                  output_all_encoded_layers=False, attention_show_flg=False)

                    loss = criterion(outputs, labels)  # 損失を計算

                    _, preds = torch.max(outputs, 1)  # ラベルを予測

                    # 訓練時はバックプロパゲーション
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                        if (iteration % 10 == 0):  # 10iterに1度、lossを表示
                            t_iter_finish = time.time()
                            duration = t_iter_finish - t_iter_start
                            acc = (torch.sum(preds == labels.data)
                                   ).double()/batch_size
                            print('イテレーション {} || Loss: {:.4f} || 10iter: {:.4f} sec. || 本イテレーションの正解率：{}'.format(
                                iteration, loss.item(), duration, acc))
                            t_iter_start = time.time()

                    iteration += 1

                    # 損失と正解数の合計を更新
                    epoch_loss += loss.item() * batch_size
                    epoch_corrects += torch.sum(preds == labels.data)

            # epochごとのlossと正解率
            t_epoch_finish = time.time()
            epoch_loss = epoch_loss / len(dataloaders_dict[phase].dataset)
            epoch_acc = epoch_corrects.double(
            ) / len(dataloaders_dict[phase].dataset)

            print('Epoch {}/{} | {:^5} |  Loss: {:.4f} Acc: {:.4f}'.format(epoch+1, num_epochs,
                                                                           phase, epoch_loss, epoch_acc))
            t_epoch_start = time.time()

    return net

学習したネットワークパラメータを保存と、テストデータでの正解率を確認する

# 学習したネットワークパラメータを保存
save_path = './weights/bert_fine_tuning_IMDb.pth'
torch.save(net_trained.state_dict(), save_path)

# テストデータでの正解率を求める
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

net_trained.eval()   # モデルを検証モードに
net_trained.to(device)  # GPUが使えるならGPUへ送る

# epochの正解数を記録する変数
epoch_corrects = 0

for batch in tqdm(test_dl):  # testデータのDataLoader
    # batchはTextとLableの辞書オブジェクト
    # GPUが使えるならGPUにデータを送る
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    inputs = batch.Text[0].to(device)  # 文章
    labels = batch.Label.to(device)  # ラベル

    # 順伝搬（forward）計算
    with torch.set_grad_enabled(False):

        # BertForIMDbに入力
        outputs = net_trained(inputs, token_type_ids=None, attention_mask=None,
                              output_all_encoded_layers=False, attention_show_flg=False)

        loss = criterion(outputs, labels)  # 損失を計算
        _, preds = torch.max(outputs, 1)  # ラベルを予測
        epoch_corrects += torch.sum(preds == labels.data)  # 正解数の合計を更新

# 正解率
epoch_acc = epoch_corrects.double() / len(test_dl.dataset)

print('テストデータ{}個での正解率：{:.4f}'.format(len(test_dl.dataset), epoch_acc))

テストデータ25000個での正解率：0.9038

Attentionの可視化

# batch_sizeを64にしたテストデータでDataLoaderを作成
batch_size = 64
test_dl = torchtext.data.Iterator(
    test_ds, batch_size=batch_size, train=False, sort=False)

# BertForIMDbで処理

# ミニバッチの用意
batch = next(iter(test_dl))

_, preds = torch.max(outputs, 1)  # ラベルを予測

可視化用の関数を用意

def highlight(word, attn):
    "Attentionの値が大きいと文字の背景が濃い赤になるhtmlを出力させる関数"

    html_color = '#%02X%02X%02X' % (
        255, int(255*(1 - attn)), int(255*(1 - attn)))
    return '<span style="background-color: {}"> {}</span>'.format(html_color, word)

def mk_html(index, batch, preds, normlized_weights, TEXT):
    "HTMLデータを作成する"

    # 指定されたインデックスのデータ（文章、ラベル、予測結果）を抽出
    sentence = batch.Text[0][index]  # 文章
    label = batch.Label[index]  # ラベル
    pred = preds[index]  # 予測

    # ラベルと予測結果を文字に置き換え
    if label == 0:
        label_str = "Negative"
    else:
        label_str = "Positive"

    if pred == 0:
        pred_str = "Negative"
    else:
        pred_str = "Positive"

    # 表示用のHTMLを作成する
    html = '正解ラベル：{}<br>推論ラベル：{}<br><br>'.format(label_str, pred_str)

    # Self-Attentionの重みを可視化。Multi-Headが12個なので、12種類のアテンションが存在
    for i in range(12):

        # indexのAttentionを抽出と規格化
        # 0単語目[CLS]の、i番目のMulti-Head Attentionを取り出す
        # indexはミニバッチの何個目のデータかをしめす
        attens = normlized_weights[index, i, 0, :]
        attens /= attens.max()

        html += '[BERTのAttentionを可視化_' + str(i+1) + ']<br>'
        for word, attn in zip(sentence, attens):

            # 単語が[SEP]の場合は文章が終わりなのでbreak
            if tokenizer_bert.convert_ids_to_tokens([word.numpy().tolist()])[0] == "[SEP]":
                break

            # 関数highlightで色をつける、関数tokenizer_bert.convert_ids_to_tokensでIDを単語に戻す
            html += highlight(tokenizer_bert.convert_ids_to_tokens(
                [word.numpy().tolist()])[0], attn)
        html += "<br><br>"

    # 12種類のAttentionの平均を求める。最大値で規格化
    all_attens = attens*0 # all_attensという変数を作成する
		for i in range(12):
				all_attens += normlized_weights[index, i, 0, :]
		all_attens /= all_attens.max()
		
		html += '[BERTのAttentionを可視化_ALL]<br>'
		for word, attn in zip(sentence, all_attens):
		
		    # 単語が[SEP]の場合は文章が終わりなのでbreak
		    if tokenizer_bert.convert_ids_to_tokens([word.numpy().tolist()])[0] == "[SEP]":
		        break
		
		    # 関数highlightで色をつける、関数tokenizer_bert.convert_ids_to_tokensでIDを単語に戻す
		    html += highlight(tokenizer_bert.convert_ids_to_tokens(
		        [word.numpy().tolist()])[0], attn)
		html += "<br><br>"
		
		return html

from IPython.display import HTML

index = 61  # 出力させたいデータ
html_output = mk_html(index, batch, preds, attention_probs, TEXT)  # HTML作成
HTML(html_output)  # HTML形式で出力

参考