fairseq WMT19翻訳システムを🤗 Transformersに移植

TL;DR

Hugging Faceは、fairseq WMT19翻訳システムをTransformersライブラリに移植し、英語‑ロシア語および英語‑ドイツ語翻訳のためのすぐに使えるモデルを提供しました。これらのモデルはAutoTokenizerとAutoModelForSeq2SeqLMでロードできます。

準備とファイルレイアウト

ポート作業は、作業ディレクトリをセットアップし、必要なリポジトリ（fairseq, mosesdecoder, fastBPE、およびdev extras付きのtransformersライブラリ）をインストールすることから始まりました。著者は~/portingフォルダーを作成し、各リポジトリをクローンして編集可能モードでインストールしました。fairseq WMT19モデルは、4つのチェックポイント（model1.pt–model4.pt）、ソースおよびターゲット辞書（dict.en.txt, dict.ru.txt）、およびBPEコードファイル（bpecodes）を提供します。これらのファイルを調べて、モデルのチェックポイント、語彙、トークン化アーティファクトを理解しました。

トークナイザーの移植

トークナイザーのエンコーダーは、既存のtokenization_xlm.pyファイルを適応させて移植されました。著者はtokenization_xlm.pyをコピーしてtokenization_fsmt.pyとし、クラス名をXLMからFSMTに変更し、未使用のコードを削除しました。WMT19モデルは別々のソースおよびターゲット語彙を使用しているため、トークナイザーのget_vocabおよびvocab_sizeプロパティはソース語彙を返すようにオーバーライドされました。BPEの処理は、fastBPEスタイル（@@ が非最終サブワード）からTransformersスタイル（ が最終サブワード）に変更され、語彙のリマッピングはfairseq.data.dictionary.Dictionary.loadを使用して正しいIDマッピングを取得することで行われました。デコーダーは後に、出力IDを文字列に変換し、BPEマーカーを削除し、Mosesデトークナイザーを適用することで完成させました。

モデル変換とアーキテクチャ

変換スクリプトconvert_fsmt_original_pytorch_checkpoint_to_pytorch.pyは、BART変換スクリプトから始めて必要な部分を徐々に追加して作成されました。モデルの重みは、fairseqハブAPIを使用してfairseqチェックポイントから抽出され、これにより古い結合されたin_proj重みを別々のk/q/v投影に変換する処理も行われます。設定引数は、fairseqの引数からTransformers FSMTConfigにマッピングされ、activation_dropout、attention_dropout、d_model、dropout、max_position_embeddings、num_hidden_layers、src_vocab_size、tgt_vocab_size、およびbos、pad、eosのトークンIDが含まれます。モデルアーキテクチャはmodeling_bart.pyから導出され、fairseqのTransformerEncoderとTransformerDecoderに合わせて層が調整されました（例：未使用の層を削除し、欠けている層を追加し、ソースおよびターゲット語彙サイズの正しい使用を確保）。サインusoidal位置埋め込みは、TorchScriptの要件を満たすために通常のnn.Embeddingサブクラスとして再実装され、決定論的な重みが保存されないようにしました。

テストと評価

トークナイザーとモデリングコンポーネントのユニットテストが追加され、既存のBARTテストスイートをベースにし、デュアル語彙セットアップに適応させました。高速CIテストのためにランダムな重みを持つ小さなモデルが生成されました。手動検証スクリプトは、デバッガーを使用して中間結果を合わせながら、fairseqとTransformersの実装の出力をトークン単位および文単位で比較しました。ビームサーチの動作は調整され、ポートされたモデルではearly_stopping=Falseを使用しており、ビームサイズ5でのfairseqのデフォルトであるearly_stopping=Trueよりも高いBLEUスコアを得られることがわかりました。WMT19テストセットをsacrebleuで評価したところ、ビームサイズ5および長さペナルティ1.1でのru‑en方向のBLEUスコアは39.0498となりました。著者は、元のfairseq論文がより高いスコアを報告しているのは、4つのチェックポイントのアンサンブルと再ランク付けステップを使用しているためであり、これらは移植版では再現されていないと指摘しました。

アップロード、統合、および自動化

変換後、モデルファイルは著者のアカウントのHugging Face S3にアップロードされ、その後facebookおよびallenai組織に移動されました。モデルは標準のAPIを使ってロードできます。たとえば、FSMTTokenizer.from_pretrained("facebook/wmt19-en-ru))のようにです。AutoConfig、AutoTokenizer、およびAutoModelWithLMHeadはfsmtモデルタイプを認識するように更新され、パイプラインスタイルの使用が可能になりました。各バリアントについて、言語ペア、ライセンス、データセット、評価指標を詳述したモデルカードが作成されました。ドキュメントは既存のBARTドキュメントをFSMTに適応させて追加され、make docsでビルドプロセスが検証されました。

含意と締めくくりの考え

WMT19システムをTransformersに移植することで、ダウンロードサイズは概ね13 GB（オプティマイザー状態を含む）から1モデルあたり約1.1 GBに削減され、ダウンストリームでの利用がより容易になりました。移植版は元の4つのチェックポイントのアンサンブルをサポートしていませんが、シングルチェックポイントモデルは依然として強い翻訳品質を達成しています。この取り組みは、既存のTransformersコンポーネント（BARTベースのモデリング、XLMベースのトークナイゼーション、変換ユーティリティ）を再利用および適応させて高品質なfairseqモデルをライブラリに組み込む方法を示しており、著者はPRレビュー過程でのSam ShleiferからのメンターシップおよびLysandre DebutとSylvain Guggerからの貢献に感謝しています。