見出し画像

end-to-endの文書画像認識モデルDonutをファインチューニングする

oshizo

DonutはOCRを使わないend-to-endの文書理解モデルです。
Vision Encoder Decoder Modelになっており、OCRエンジンに依存せずに視覚的な文書分類や情報抽出を高い精度で行うことができます。

Donutは日本語を含む4言語で学習されたモデルnaver-clova-ix/donut-baseが公開されており、日本語で何かしたいときにファインチューニングして使えそうだなと思っていました。
今回、AIキャラクターと一緒にノベルゲームをプレイするために、ノベルゲーム風画面の合成データセットでdonut-baseをファインチューニングしました。

以下を目標として作成しました。

  • <unk>になる漢字をvocabに追加して学習する

  • 選択肢、名前、メッセージを別々に認識し、jsonを出力する

  • SKIP、LOADなどのUIの文字、日付表示などを読み取らない

  • ルビを無視する。ルビに影響されずに文字認識する

漢字の認識精度を高くすることは目標にしておらず、そのため学習量はかなり少なめです（50kサンプル、ColabのT4で18h程度）。
認識結果のjsonをChatGPTに投げて会話することが目的なので、ある程度誤認識しても雰囲気でやってくれることを期待しているためです。

学習済みモデルと認識結果の例

モデルはでHubで公開しています。

動かし方や注意点は上記のHubのreadmeと、Colabで動かせるnotebookサンプルを見てください。
oshizo/donut-base-japanese-visual-novel

認識結果の例1

{'options': '', 'names': '結月', 'messages': 'この神社には古い言い伝えがあるの。神樹の下で誓いを立てると、その願いは必ず叶うという。心を開いて、自分の想いを信じてみて。'}

選択肢なし、ルビあり、UI要素あり

認識結果の例2

{'options': ['行こう!', '今回は見送る', '準備を整えるまで待って(会話から抜けます)', '旅の目的について詳しく教えてください'], 'names': 'リリアン', 'messages': '私たちの使命は、新たな発見と交流を通じて地球と宇宙の未来を築くこと。この壮大な旅に参加する準備はできているかしら?'}

選択肢2x2レイアウト、日付表示

認識結果の例3

{'options': ['全力で攻撃する!勝利をつかめ!', '堅実に守り、敵の隙を待とう。'], 'names': '', 'messages': '敵を誘い込んで、戦術を駆使せよ。'}

選択肢のみ（名前なし、メッセージなし）

この例は3つ目の選択肢「敵を～」が誤ってmessagesキーとして読み取られてしまいました。

認識結果の例4

{'options': 'もちろん、手伝います!', 'names': '下尾崎 菊欠郎', 'messages': 'この書斎は重要な手がかりが隠されているかもしれない。君も協力してくれるか?'}

選択肢1つ、名前がメッセージボックスの下にあるパターン

名前はメッセージの下にあっても読み取れるように訓練しましたが、このレイアウトの認識率はあまり高くないです。

データセット

データセットの作成方法

ノベルゲーム風の画像と教師ラベルになる文字列のペアを集めるために、合成データ生成器を実装しました。データ生成器のリポジトリはoshizo/VisualNovelVDUDataGeneratorで公開しています。

Belval/TextRecognitionDataGeneratorを参考にしており、PIL.ImageFontやPIL.ImageDrawを使ってプログラム的に画像を作成する方式です。
フォントファイル、背景画像、透過pngのキャラクター画像、メッセージテキストを与えて画像を合成します。

今回は画像に含まれるルビを無視するモデルにしたかったので、↓の画像のように、ルビやUIをレンダリングしつつ、それらは教師ラベルには追加しないようにしました。

教師ラベル … {'options': [], 'names': [], 'messages': ['南長柄には美しい自然が広がっています。']}

背景・キャラクター画像

Stable Diffusionなどを使い、適当な背景画像1,000枚とキャラクター画像500枚を生成しskytnt/anime-remove-backgroundを使って透過pngを作りました。
画像のクオリティは重要ではないので、わざわざ生成する必要はなかったと思います（認識精度というより、楽しさの問題）

テキスト

DonutのdecoderはTransformerが使われているので、テキストは支離滅裂な文字の羅列ではなく、日本語として自然な文章になっていることが望ましいです。

今回は、以前rinna/japanese-gpt2-mediumをキャラクターのセリフに特化してファインチューニングしたモデルを使って5万件程度のテキストを生成しました。
ランダムな生成だと文字のカバー率が低くなってしまうのでデータセットに含めたい文字の単語をpromptにして、その後ろに自然なセリフを生成する方法を取りました。

こういうテキストが生成されます（太字がprompt）

驚愕に値しない......というか、これはただの嘘ではないのですか?
疥癬虫を駆除して、治療薬を塗ったところ......治りました!
謄写版だ。インクを乾かすと、文字がくっきり出るんだよ。
蔗糖の製錬が上手くいっているのは、お主らが居るからだ。

（テキストの先頭に難しい漢字が偏るのが気になりますが…）

日本語トークンの追加

donut-baseのtokenizerはXLMRobertaTokenizerで、これは57k程度の語彙を持つSentencePieceが使われています。
多言語を57kに詰め込んでいるため日本語の漢字カバー率は微妙で、「暮」「択」など、普通に目にする漢字も<unk>になってしまいます。

作成したデータセットには文字が5200種類ほど含まれているのですが、そのうち1500種類が<unk>になることが分かったのでこれらのtokenをそれぞれ1文字1tokenで追加しました。

語彙の追加方法

具体的にはtokenizer.jsonの"model"の下の"vocab"に、SentencePieceの情報（token文字列とスコア値）がリストとして埋め込まれているので、ここに直接追加しjsonファイルを上書きしました。
下の画像のリスト末尾に1,500文字分の行を追加しました。

tokenizer.json

↓以下余談

Transformersモデルに語彙を追加するときは、tokenizer.add_tokenメソッドが使えるのですが「add_tokenに気を付けろ！」という記事があったり、add_tokenで追加した語彙はtokenizer.encodeしたときに追加したtokenの後ろに半角スペースtokenが付く仕様が気になり、文字の追加ではadd_tokenは使いませんでした。
この仕様は、追加したtokenを含む文字列がそのtokenに分割されることを保証するための処理なのかと想像しましたが…区切り文字がスペースの言語なら問題なくても日本語だといちいちスペースが挿入され、対処法が分かりませんでした。
donutのくずし字モデルはadded_tokens.json（naver-clova-ix/donut-base-finetuned-kuzushiji）に語彙が追加されているようなので、私のやり方は適切ではなさそうな気がします。

Token Embeddingの設定

語彙追加をした後、DecoderのEmbedding層を追加した語彙のサイズに合わせて拡張しておく必要があります。
これはmodel.decoder._resize_token_embeddingsメソッドで行うことができます。

このメソッドではパラメタがランダムに初期化されますが、ファインチューニングはあまり長期間行うつもりがなかったのでランダムだと学習しきれなそうな気がし、ランダムではなく近いtokenの埋め込みをコピーして初期化してみることにしました。

今回の場合の近いtokenは（tokenが1文字の前提で）見た目の形が似ている文字と考えることにました。
GPT-4に聞いてみてcv2.findContoursで輪郭を抽出してcv2.matchShapescontourするとか、HOG特徴量を比べるなど教えてもらってやってみたのですがいまいちうまくいかず、最終的には未知の文字の画像をdonut-baseに読ませ、間違えて出力するtokenを近いtokenとみなすことにしました。

以下のような画像をdonut-baseに読ませて、「こんにちは」の後ろにくるtokenを使いました。「こんにちは」を付けているのはこれが日本語と認識させるためで、中国語tokenを予測しづらくなります。

donut-baseに読ませて「暮」に近い文字を探す

結果の抜粋です（左が追加した文字→右が初期埋め込みとして使った近いtokenの文字）。結構納得感あると思います。

  • 「溝」 → 「薄」を使用

  • 「洟」 → 「湊」を使用

  • 「努」 → 「祭」を使用

  • 「芟」 → 「支」を使用

  • 「腺」 → 「服」を使用

こうして作った追加するtokenと既存の似たtokenのペアを使い、無理やりinput_embeddingを上書きしました。

input_embeddings = model.get_input_embeddings()
for new_token_id, similar_existing_token_id in zip(new_token_ids, similar_existing_token_ids):
    input_embeddings.weight.data[new_token_id] = input_embeddings.weight.data[similar_existing_token_id].clone()
model.set_input_embeddings(input_embeddings)

学習

Donutを学習させるときの注意事項ですが…公式のリポジトリ直下にあるtrain.pyではなくNielsRogge/Transformers-Tutorialsにあるサンプルnotebookを使うことをおすすめします。

Donutは公開後しばらくしてHuggingfaceに取り込まれたのですが、公式リポジトリのtrain.pyはHuggingface版に対応していません。

NielsRogge/Transformers-TutorialsのサンプルnotebookはHuggingface版を使っており、説明もわかりやすいです。

今回の場合、FullHDの画像でデータセットを作ると1枚あたり2MB前後、50k件100GBになりちょっとでかすぎるということで、あらかじめ生成せずにDatasetの__getitem__の中で動的に画像生成する改変を行っています。

学習はColabのT4を使いました。
入力は1920x1280の解像度で、fp16混合精度、batch sizeは1でVRAM13.1GB
1.36sec/it程度で、50kサンプルを学習するのに18時間程度でした。

課題

  • 選択肢の数が多い場合に、選択肢の一部が認識されないことが多い

  • 文字数の少ない選択肢がかなり認識されづらい

  • 複数行のメッセージの場合、1行目が無視されることがある

  • アスペクト比が1920:1080でない画像を含めて学習する

  • 名前やメッセージが複数ある画面レイアウト（チャット形式など）や、メッセージが画面の上半分にあるレイアウトなど、学習していないレイアウトを学習データに含めて学習

認識率はまだ十分ではないので、使う方の実装（AIキャラクターと楽しくノベルゲームを遊ぶ）を進めつつ、その中で問題になった項目に対処してv2にしようと考えています。
ただ、ChatGPTとプレイヤー（私）が会話しながら遊ぶ場合、認識されなかった文字列はチャットで教えてあげれば何とかなるので、現状でもそこそこ使えそうな気はしています。

この記事が気に入ったらサポートをしてみませんか？