BERTにラベルの埋め込みを追加する方法

# はじめに
BERTを使う時にテキストデータと一緒にラベルデータを使いたいときはありませんか？
例えば以下のような`テキストデータから書いた人の年齢を当てる`タスクを解くとします．
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/842766/b887573a-1063-be95-aa07-33b8eac463c0.png)
この時にモデルに対してそれぞれの文節がどのような特徴を持つかを明示する，つまり以下のような文節の種類を表すラベルデータを一緒にモデルに入力出来れば精度が上がりそうな気がしますね．（実際に上がるかはかなりタスク依存）
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/842766/814b1a13-c729-05c3-c393-85ae536d9540.png)
huggingfaceのBERTではこれが簡単に実装できるようになっているため，今回の記事ではその方法を紹介していきたいと思います．

例以外にも要約タスクの共参照関係や係り受けの関係性を明示するなどいろいろなタスクに用いられます．

# 手法
## インポート

```python
import torch
import torch.nn as nn
from transformers import BertModel
from transformers import BertJapaneseTokenizer
```

## データの用意
簡単のためにバッチなどは一旦置いておいて，テキストデータとラベルデータを簡単に用意します．
ここでラベルは以下のように定義します．

| ラベルID | ラベル | 説明 |
|:-:|:-:|:-:|
| 0 | 挨拶 | こんにちは，おはようなどの挨拶 |
| 1 | 事実説明 | 現在位置を言うなど事実をただ述べているだけ |
| 2 | 自己開示 | 自身の趣味や職業などの開示 |

```python
sentences = [
    ("こんにちは[SEP]私は小学生の頃にマリカDSをプレイしていました。[SEP]今は新卒2年目の会社員です。"),
    ("初めまして！[SEP]今日はいい天気ですね。[SEP]天気が良い日は散歩がしたくなります。"),
]

# 各文節に対応するラベルデータ
labels =[
    [0, 2, 2],
    [0, 1, 2],
]

# 教師データ（年齢）
ages = [25, 43]
```

## 埋め込みの獲得
### テキストデータ
huggingface transformersでBERTを使っている方にはおなじみの`BertJapaneseTokenizer`を使ってトークナイズします．

```python
tokenizer = BertJapaneseTokenizer.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking")
tokenized_sentences = tokenizer.encode_plus(sentences, return_tensors="pt")
```

通常はここでBERTに入力してしまうのですが，今回はラベル埋め込みを加算する関係上，自前で埋め込みを作ります．
ここではわかりやすさのために事前学習済みモデル`cl-tohoku/bert-base-japanese-whole-word-masking`の語彙サイズ(32000)と埋め込みサイズ(768)を設定していますが，BertConfig等を使って`config.vocab_size`のように指定してあげるとよりスマートです．
bert_embedingにトークナイズしたテキストデータを与えてテキストデータの埋め込みを作ります．

```python
bert_embeding = nn.Embeding(32000, 768)
text_embeds = bert_embeding(tokenized_sentences["input_ids"])
```

pytorchのEmbedingsが分からない方はこちらの分かりやすい記事がおすすめです．

https://gotutiyan.hatenablog.com/entry/2020/09/02/200144

### ラベル
続いてラベルの埋め込みです．こちらは元からラベルIDの列として定義しているので先程と同じ手順で埋め込みを計算します．
新しく`nn.Embedings`のインスタンスを作成してあげないとテキストデータの埋め込みと競合してしまうので，新しく定義し直します．
ただしこちらの語彙サイズは **3**とします．これは用いるラベルの種類の数です．埋め込みサイズは先程と同様に768とします．

```python
label_embeding = nn.Embeding(3, 768)
label_embeds = label_embeding(tokenized_sentences)
```

### 埋め込みの加算
単純ですがここが重要な部分です．
今回やりたい「`BERTに対してテキストデータと対応するラベルデータを一緒に入力する`」というのはテキストデータの埋め込みとラベルの埋め込みを単純に加算することで実現できます．

```python
inputs_embeds = text_embeds + label_embeds
```

## BERTへの入力
ここまでで入力データを埋め込みに変換できたので，最後にモデルに入力して予測を行いたいと思います．
`BertForSequenceClassification`は`num_labels=1`とすると回帰予測が出来るので今回はこれを使います．
そして`inputs_embeds`引数に先程作った埋め込みのinputs_embedsを渡します．
これでテキストデータとラベルデータを考慮した予測を行うことが出来ます．
ただし，今回使っている東北大BERTはそんな事前学習は一切されていないので，性能が上がるどころか下がってしまうと思います．（未検証）
そのため，テキストデータとラベルデータを考慮したモデルを作りたい場合はファインチューニングをすることを推奨します．

```python
model = BertForSequenceClassification.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking', num_labels=1)
outputs = model(inputs_embeds=inputs_embeds, labels=ages)
```

BertForSequenceClassificationの使い方については，タスクは違いますが以下の記事がおすすめです．

https://qiita.com/nekoumei/items/7b911c61324f16c43e7e#bertforsequenceclassification%E3%81%AB%E3%81%A4%E3%81%84%E3%81%A6

# まとめ
ここまでありがとうございました．
テキストデータとラベルデータの両方を考慮させたい！となった時に「BERTのアーキテクチャから自作しないと．．．」と思ったら便利な機能があったので記事にしてみました．
いいねを貰えると励みになるのでこの記事がいいなと思ったらぜひいいねをお願いします．
ミスなどがありましたらコメントください．