NERモデルをMMOCRで学習させてみる

# NERとは
Named Entity Recognitionの頭文字を取ったもので、文中から固有表現 (Named Entity) を抽出し、それを固有名詞（人名、組織名、地名など）や日付、時間表現、数量、金額、パーセンテージなどのあらかじめ定義された固有表現分類へと分類するもの[[wiki参照](https://ja.wikipedia.org/wiki/%E5%9B%BA%E6%9C%89%E8%A1%A8%E7%8F%BE%E6%8A%BD%E5%87%BA)]を指します

## NERを試してみる
MMOCRで配布されているNERのdemoを動かしてみます。
そのために少し準備が必要です。
MMOCRのサイトから[NERの学習済みcheckpoint](https://mmocr.readthedocs.io/en/latest/ner_models.html)、[Dataset](https://mmocr.readthedocs.io/en/latest/datasets/ner.html)をダウンロードし、これを以下のように配置します。
```
mmocr
    ├-- mmocr
    ├-- demo
        ├-- ner_demo.py
    ├-- configs
        ├-- ner
            ├-- bert_softmax
                ├-- bert_softmax_cluener_18e.py
    ├-- checkpoints(なければmkdir)
        ├-- bert_softmax_cluener-eea70ea2.pth(これを配置)
    ├-- data
        ├-- cluener2020(これを配置)
            ├-- cluener_predict.json
            ├-- dev.json
            ├-- README.md
            ├-- test.json
            ├-- train.json
            ├-- vocab.txt
```
次に以下のコマンドを実行します。
```
python demo/ner_demo.py \
    configs/ner/bert_softmax/bert_softmax_cluener_18e.py \
    checkpoints/bert_softmax_cluener-eea70ea2.pth
```
実行すると入力を求められるのでNERを試したい文章を入力します。
ただしこのモデルは[CLUENER2020](https://github.com/CLUEbenchmark/CLUENER2020)という中国語のデータセットで学習させたものですので中国語（CLUENER2020のサイトに載っていた文）を入力してみます。
```
Please enter a sentence you want to test: 北京勘察设计协会副会长兼秘书长周荫如
organization: 北京勘察设计协会
position: 副会长
position: 秘书长
name: 周荫如
```
組織名や役職、人物名に分類されています。
ここで少し述べておくと、CLUENER2020は以下のラベルに分類するデータセットになっています。
```
address, book, company, game, government, movie, name, organization, position, scene
```

# Fine-Tuning
MMOCRはNERのためのモデルとしてBERTのみを提供しています。
以下では付属のDataset、任意のDatasetでこの既存モデルをFine-Tuningする方法について記します。
## MMOCR付属のDatasetを用いたFine-Tuning
MMOCRが用意しているDatasetはCLUENER2020のみです。このDatasetを用いたFine-Tuningについては既にconfig fileが用意されていますので、それを使ってFine-Tuningしてみます。
```
python tools/train.py \
    configs/ner/bert_softmax/bert_softmax_cluener_18e.py \
    --cfg-options work_dir="work_dirs/cluener"
```
上記のコマンドを実行することで学習を開始することが出来ます。`work_dir`に設定したファイルに学習結果の重み、logなど学習結果が格納されます。
```
python demo/ner_demo.py \
    work_dirs/cluener/bert_softmax_cluener_18e.py \
    work_dirs/cluener/latest.pth
```
で学習結果を確認できます。
（tools/test.pyの動作は確認できていないです）

## 任意のDatasetを用いたFine-Tuning
### Datasetの準備
任意のDatasetを用いる場合は以下のようなファイルが必要になります。
```
train.json
dev.json
test.json(使われていない？)
vocab.txt
```
CLUENER2020のDatasetを元にそれぞれのファイルの構造を見てみます。
```:train.json
...
{"text": "生生不息CSOL生化狂潮让你填弹狂扫", "label": {"game": {"CSOL": [[4, 7]]}}}
{"text": "那不勒斯vs锡耶纳以及桑普vs热那亚之上呢？", "label": {"organization": {"那不勒斯": [[0, 3]], "锡耶纳": [[6, 8]], "桑普": [[11, 12]], "热那亚": [[15, 17]]}}}
...
```
```:dev.json
...
{"text": "光大三家银行信用卡专业人士分别就境外购买1000欧元、1000英镑和1万港币的商品，", "label": {"company": {"光大": [[0, 1]]}}}
{"text": "4、雷吉纳vsac米兰推荐：0", "label": {"organization": {"雷吉纳": [[2, 4]], "ac米兰": [[7, 10]]}}}
...
```
```:test.json
...
{"id": 6, "text": "央视新址文化中心外立面受损严重"}
{"id": 7, "text": "单看这张彩票，税前总奖金为5063992元。本张票面缩水后阿森纳的结果全部为0，斯图加特全部为1，"}
...
```
`train.json`、`dev.json`には文とそれに対するラベルが対応するアノテーションファイルとなっています。構造としては辞書の入れ子となっていることがわかります。一方`test.json`には文が羅列されています。
（ぱっと見ですがtestとdevの命名が逆な気がします）
```:vocab.txt
...
request
##gence
qt
##っ
1886
347
363
q7
##zzi
...
```
また、vocab.txtは語彙の羅列になっています。
#### wiki Dataset
これらの構造を考えながら日本語のデータセットである[ner-wiki-dataset](https://github.com/stockmarkteam/ner-wikipedia-dataset)で学習を行ってみます。
[ner-wiki-dataset](https://github.com/stockmarkteam/ner-wikipedia-dataset)から`ner.json`をダウンロードしてきて、その構造を確認します。
```:ner.json
[
    {
        "curid": "3572156",
        "text": "SPRiNGSと最も仲の良いライバルグループ。",
        "entities": [
            {
                "name": "SPRiNGS",
                "span": [
                    0,
                    7
                ],
                "type": "その他の組織名"
            }
        ]
    },
    {
        "curid": "2415078",
        "text": "レッドフォックス株式会社は、東京都千代田区に本社を置くITサービス企業である。",
        "entities": [
            {
                "name": "レッドフォックス株式会社",
                "span": [
                    0,
                    12
                ],
                "type": "法人名"
            },
            {
                "name": "東京都千代田区",
                "span": [
                    14,
                    21
                ],
                "type": "地名"
            }
        ]
    },
    ...
]
```
GithubのREADMEによると以下のような構造をしているらしいです。
>curidはデータ元のWikipediaのページID
textはタグ付を行う対象のテキスト
entitiesは固有表現のリスト
nameは固有表現名
spanはtextでの位置
typeは固有表現のタイプ

`train.json`と比較してみると
<table>
  <thead>
    <tr>
      <th>Wiki Dataset</th> <th>Cluener2020(MMOCR)</th>
    </tr>
  </thead>
  <tr>
    <td> text </td> <td>text</td>
  </tr>
  <tr>
    <td> enitities </td> <td> label </td>
  </tr>
  <tr>
    <td> {"name": , "span": , "type", } </td> <td> {"type": {"name": ["span"]}} </td>
  </tr>
</table>

となっています。また`train.json`はJSON形式ですが`ner.json`はNDJSON形式です。
これらの違いを踏まえて **Wiki Dataset**をMMOCRで利用できるように変換します。
```python:converter.py
import random
import json
import ndjson
import sys

def convert_from_wiki_dataset(wiki_dataset_file, train_ratio=0.9):
    '''
    construct dataset from https://github.com/stockmarkteam/ner-wikipedia-dataset
    Parameters:
    -----------
    wiki_dataset_file: str
        path to the json file
    train_ratio: float
        separate dataset into train and test with this ratio
    -----------
    '''
    with open(wiki_dataset_file, 'r', encoding='UTF-8') as f:
        annotated_texts = json.load(f)
    results = []
    for annotated_text in annotated_texts:
        result = {'text': None, 'label': None}
        text = annotated_text['text']
        result['text'] = text

        entities = annotated_text['entities']
        label = {}
        for entity in entities:
            entity_type = entity['type'] # str ex. 施設名, イベント名
            name = entity['name'] # str ex. 日本橋支店, ウッドストック・フェスティバル
            span = entity['span'] # list ex. [49, 54], [24, 39]
            span[1] -= 1
            tmp_dict = {name: [span]}
            if entity_type in label:
                label[entity_type].update(tmp_dict)
            else:
                label[entity_type] = tmp_dict
        result['label'] = label
        results.append(result)
    
    random.shuffle(results)
    train_data = results[0:int(len(annotated_texts)*train_ratio)]
    test_data = results[int(len(annotated_texts)*train_ratio):]

    with open('train.json', 'w', encoding='UTF-8') as f:
        ndjson.dump(train_data, f, ensure_ascii=False)
    with open('test.json', 'w', encoding='UTF-8') as f:
        ndjson.dump(test_data, f, ensure_ascii=False)

def main():
    convert_from_wiki_dataset(sys.argv[1])

if __name__ == '__main__':
    main()
```
上記のコードを実行すると
```
python converter.py <ner.jsonのパス>
```
9:1に分割された`train.json`、`test.json`が出力されます。
出力された`train.json`を見てみるとMMOCRで利用できる形に変換されていることがわかります。
```:train.json(wiki dataset)
{"text": "標高1,600-2,200メートルの山地にある森林に生息する。", "label": {}}
{"text": "2017年から第5代メタルワン代表取締役社長を務め、2018年には住友商事との間で国内鋼管事業統合により新会社住商メタルワン鋼管の設立を行うことなどで合意した。", "label": {"法人名": {"メタルワン": [[10, 14]], "住友商事": [[33, 36]], "住商メタルワン鋼管": [[55, 63]]}}}
...
```
これらのファイルを以下のように配置します。
```
mmocr
    ├-- mmocr
    ├-- demo
    ├-- configs
        ├-- ner
            ├-- bert_softmax
                ├-- bert_softmax_wiki_18e.py(後ほど配置)
    ├-- checkpoints
    ├-- data
        ├-- wiki(これを配置)
            ├-- test.json
            ├-- train.json
            ├-- vocab.txt
```
`Cluener2020`の`vocab.txt`には日本語の語彙も含まれているため、同じものを利用します。

### config fileの記述
次にconfig fileを記述していきます。
使うモデルがCLUENER2020のときと同じことから`configs/ner/bert_softmax/bert_softmax_cluener_18e.py`を継承して`configs/ner/bert_softmax/bert_softmax_wiki_18e.py`を作成します。
```python:bert_softmax_wiki_18e.py
_base_ = [
    './bert_softmax_cluener_18e.py'
]

categories = [
    '人名', '法人名', '政治的組織名', 'その他の組織名', '地名', '施設名', '製品名', 'イベント名'
]

test_ann_file = 'data/hogehoge/test.json'
train_ann_file = 'data/hogehoge/train.json'
vocab_file = 'data/hogehoge/vocab.txt'

max_len = 256
```
上記の`categoriesには**Wiki Dataset**のラベルをすべて列挙します。
<table>
  <thead>
    <tr>
      <th>タイプ</th> <th>固有表現数</th> <th>備考</th>
    </tr>
  </thead>
  <tr>
    <td> 人名 </td> <td> 2980 </td> <td></td>
  </tr>
  <tr>
    <td> 法人名 </td> <td> 2485 </td> <td> 法人又は法人に類する組織 </td>
  </tr>
  <tr>
    <td> 政治的組織名 </td> <td> 1180 </td> <td> 政治的組織名、政党名、政府組織名、行政組織名、軍隊名、国際組織名 </td>
  </tr>
  <tr>
    <td> その他の組織名 </td> <td> 1051 </td> <td> 競技組織名、公演組織名、その他 </td>
  </tr>
  <tr>
    <td> 地名 </td> <td> 2157 </td> <td></td>
  </tr>
  <tr>
    <td> 施設名 </td> <td> 1108 </td> <td></td>
  </tr>
  <tr>
    <td> 製品名 </td> <td> 1215 </td> <td> 商品名、番組名、映画名、書籍名、歌名、ブランド名等 </td>
  </tr>
  <tr>
    <td> イベント名 </td> <td> 1009 </td> <td></td>
  </tr>
</table>

`test_ann_file`、`train_ann_file`には変換によって得られた`test.json`、`train.json`を設定します。
`max_len`には`text`の入力として何文字までを許容するかを記します。今回であれば大凡400文字が`train.json`内の最大文字数となります。ただし、ほとんどの`text`がそれほど大きくないため今回は`max_len`を256に設定してあります。

:::note warn
実行時に256以上のtextが入力されるとWarningを吐きますが、この入力はスキップされ処理は問題なく継続されます。
:::


### train.pyの実行
```
python tools/train.py \
    configs/ner/bert_softmax/bert_wiki_18e.py \
    --cfg-options work_dir="work_dirs/wiki"
```
上記のコマンドを実行することで学習を開始することが出来ます。
### 学習結果の確認
```
python demo/ner_demo.py \
    work_dirs/hogehoge/bert_softmax_wiki_18e.py \
    work_dirs/wiki/latest.pth
```
で学習結果を確認できます。
試しに適当な分を入力してみるとしっかりNERされています。
```
Please enter a sentence you want to test: イギリスはリーマンショック直後の2008年10月にイングランド銀行のバランスシートを一気に3倍近く増やした後、2008年11月から2009年3月にかけて段階的に縮小させていった。
地名: イギリス
イベント名: リーマンショック
政治的組織名: イングランド銀行
```

# 番外編
## NERとOCRの接続
MMOCRの`NER`はMMOCRの`DET`、`RECOG`とは接続されていません。（NERはOCRに直接関係ないからでしょうか）
これをうまく接続することで画像の入力に対し固有表現抽出できるようになります。
`NER`、`DET`、`RECOG`を実行するためのコードは以下のファイルに記述されています。
・`NER` : `demo/ner_demo.py`
・`DET`、`RECOG` : `mmocr/utils/ocr.py`

### コードの解析
`demo/ner_demo.py`はconfigとcheckpointをコマンドライン引数として取り、その後入力された文に対して固有表現抽出を行うものとなっており、簡単なコードです。
一方、`mmocr/utils/ocr.py`は複雑な構造をしているため動作を簡単に解説します。
`mmocr/utils/ocr.py`の主要なコマンドライン引数は以下の通りになっており、この引数の有無で動作が変化します。
<table>
  <thead>
    <tr>
      <th>引数</th> <th>説明</th>
    </tr>
  </thead>
  <tr>
    <td> --det </td> <td> MMOCRが用意しているdetectionモデル名 </td>
  </tr>
  <tr>
    <td> --det-config </td> <td> MMOCRが用意しているdetectionモデルについて独自で学習させたconfigファイルのパス </td>
  </tr>
  <tr>
    <td> --det-ckpt </td> <td> MMOCRが用意しているdetectionモデルについて独自で学習させたckptファイルのパス </td>
  </tr>
  <tr>
    <td> --recog </td> <td> MMOCRが用意しているrecogモデル名 </td>
  </tr>
  <tr>
    <td> --recog-config </td> <td> MMOCRが用意しているrecogモデルについて独自で学習させたconfigファイルのパス </td>
  </tr>
  <tr>
    <td> --recog-ckpt </td> <td> MMOCRが用意しているrecogモデルについて独自で学習させたckptファイルのパス </td>
  </tr>
  <tr>
    <td> --kie </td> <td> MMOCRが用意しているkieモデル名 </td>
  </tr>
  <tr>
    <td> --kie-config </td> <td> MMOCRが用意しているkieモデルについて独自で学習させたconfigファイルのパス </td>
  </tr>
  <tr>
    <td> --kie-ckpt </td> <td> MMOCRが用意しているkieモデルについて独自で学習させたckptファイルのパス </td>
  </tr>
</table>

`--**-config`、`--**-ckpt`は独自で学習させた結果を用いる際のオプションとなります。
また、MMOCRが用意していないモデルの使用は想定されていません。
これらの引数の有無とその動作、呼び出される関数を以下に示します。
<table>
  <thead>
    <tr>
      <th>引数</th> <th>動作</th>
    </tr>
  </thead>
  <tr>
    <td> det、recogどちらかのみ指定 </td> <td> 指定されたどちらかの操作のみを実行、single_inference()</td>
  </tr>
  <tr>
    <td> det、recogを指定 </td> <td> det後にrecogを実行、det_recog_kie_inference() </td>
  </tr>
  <tr>
    <td> det、recog、kieを指定 </td> <td> det、recog、kieを順に実行、det_recog_kie_inference() </td>
  </tr>
</table>

これらのことから以下を実装すれば`NER`と`OCR`を接続できます。
1. コマンドライン引数に`--ner`、`--ner-config`、`--ner-ckpt`を追加（独自モデルの利用を想定（`--ner`が`None`でなければOCR結果を`NER`する））
2. コマンドライン引数の処理を`--ner`、`--ner-config`、`--ner-ckpt`について追加
3. `demo/ner_demo.py`と`det_recog_kie_inference()`をもとに`det_recog_ner_inference()`を実装
4. コマンドライン引数の有無による処理分岐を実装

### 実装
上記の4項目を実装します。
まずは、1.を実装します。引数を指定している`parse_args()`に以下のように追記します
``` python:ocr.py
def parse_args():
    ...
    parser.add_argument(
        '--ner',
        type=str,
        default='',
        help='Pretrained named entity recognition algorithm')
    parser.add_argument(
        '--ner-config',
        type=str,
        default='',
        help='Path to the custom config file of the selected ner model. It'
        'overrides the settings in ner')
    parser.add_argument(
        '--ner-ckpt',
        type=str,
        default='',
        help='Path to the custom checkpoint file of the selected ner model. '
        'It overrides the settings in ner')
    ...
```
次に2.を実装します。コマンドラインの処理は`__init__()`に追記します。
``` python:ocr.py
def __init__(self,
             ...
             kie_config='',
             kie_ckpt='',
             ner='', # 追記
             ner_config='', # 追記
             ner_ckpt='', # 追記
             config_dir=os.path.join(str(Path.cwd()), 'configs/'),
             device=None,
             **kwargs):
    ...
    self.td = det
    self.tr = recog
    self.kie = kie
    self.ner = ner # 追記
    self.device = device
    ...
    self.kie_model = None
    if self.kie:
        ...
    
    self.ner_model = None # 追記
    if self.ner: # 追記
        self.ner_cfg = Config.fromfile(ner_config) # 追記
        # build the model from a config file and a checkpoint file
        self.ner_model = init_detector(ner_config, ner_ckpt, device=self.device) # 追記
```
次に3.を実装します。
`det_recog_kie_inference()`をコピーし、以下のように変更を加えます。
``` python:ocr.py
def det_recog_ner_inference(self, det_model, recog_model, ner_model): # 変更
    ...
    for filename, arr, bboxes, out_file in zip(self.args.filenames,
                                               self.args.arrays,
                                               bboxes_list,
                                               self.args.output):
        ...
        if self.args.merge:
            img_e2e_res['result'] = stitch_boxes_into_lines(
                img_e2e_res['result'], self.args.merge_xdist, 0.5)
        
        # 以下を追記
        for i in range(len(img_e2e_res['result'])):
            text = img_e2e_res['result'][i]['text']
            ner_result = text_model_inference(ner_model, text)
            print('')
            print('recognized text in box:', text)
            # show the results
            for pred_entities in ner_result:
                for entity in pred_entities:
                    print(f'{entity[0]}: {text[entity[1]:entity[2] + 1]}')
                    if entity[0] in img_e2e_res['result'][i]:
                        img_e2e_res['result'][i][entity[0]].append(text[entity[1]:entity[2]])
                    else:
                        img_e2e_res['result'][i][entity[0]] = []
                        img_e2e_res['result'][i][entity[0]].append(text[entity[1]:entity[2]])
    ...
```
最後に4.を実装します。条件分岐は`readtext()`に記載されており、これを変更します。
``` python:ocr.py
def readtext(...):
    ...
    pp_result = None

    # Send args and models to the MMOCR model inference API
    # and call post-processing functions for the output
    if self.detect_model and self.recog_model and self.ner_model: # 追記
        det_recog_result = self.det_recog_ner_inference( # 追記
            self.detect_model, self.recog_model, ner_model=self.ner_model) # 追記
    elif self.detect_model and self.recog_model: # 変更
        det_recog_result = self.det_recog_kie_inference(
            self.detect_model, self.recog_model, kie_model=self.kie_model)
        pp_result = self.det_recog_pp(det_recog_result)
    else:
        for model in list(
                filter(None, [self.recog_model, self.detect_model])):
            result = self.single_inference(model, args.arrays,
                                           args.batch_mode,
                                           args.single_batch_size)
            pp_result = self.single_pp(result, model)

    return pp_result
```
### 実行
これらの変更を加えて、以下のコマンドにより`OCR`+`NER`を実行してみます。
また、本記事で学習させたwikiモデルを使用します。
```
python mmocr/utils/ocr.py demo/任意の画像 \
    --det PS_CTW \
    --recog SAR \
    --ner 適当な文字列 \
    --ner-config configs/ner/bert_softmax/bert_softmax_wiki_18e.py \
    --ner-ckpt work_dirs/wiki/epoch_18.pth
```

実行結果として入力画像に対するBBOX内のテキストとそれに対するNER結果が表示されます。

# 本記事について
本記事はTRUST SMITH株式会社でのインターン中に書きました。
TRUST SMITH株式会社は東大発のAIベンチャーです。
インターンやポスト、提供サービスに興味がおありでしたら以下から会社HPにアクセスしていただけます。
[TRUST SMITH株式会社](https://www.trustsmith.net/)
[SMITH&VISION](https://www.smith-vision.net/)