RustでBERT！？rust-bertで遊んでみた (前半)

最近、勉強を始めたRust初心者の「よこ」です。
Rustには rust-bert なるクレートがあるんですよ！！！期待しか無いですよね。
実際に、 rust-bert でHuggingface上のモデルを使ってみたいなと思い、
いろいろ触ってみました。

流れ

• この記事→　環境構築＆準備(モデルの変換)
• 次の記事→　GPT-2でTextGenerationタスク

今回、PyTorch用→Rust用への変換から解説しています。
モデルの変換をせず、推論だけ試してみたい方は、

• libtorchのインストール
• Rust用に変換
  を飛ばしてください！

使用環境

OS: Pop!_OS 21.10（Ubuntu）
CUDA: 11.3
Python: 3.8.9
torch: 1.10.2+cu113
(CPU用で書いてます！)

Rust環境の準備

Rustの日本語ドキュメント通りに進めて、cargoが使えるようにします。
（割愛します）
https://doc.rust-jp.rs/book-ja/ch01-01-installation.html

libtorchのインストール

PyTorchのダウンロードページから、LibTorchをダウンロードしてきます。

https://pytorch.org/

Download here (cxx11 ABI):のURLをコピーして、wgetで取ってきました。

bash

wget https://download.pytorch.org/libtorch/cu113/libtorch-xxx.zip
unzip libtorch-xxx.zip
sudo cp -r libtorch/include/* /usr/local/include/
sudo cp -r libtorch/lib/* /usr/local/lib/
sudo cp -r libtorch/share/* /usr/local/share/
sudo ldconfig
sudo reboot now

※こちらの記事を参考にさせていただきました。
pythonで学習したDNNモデルをC++から利用する（PyTorch & libtorch版）

変換用モデルの準備

(自前のモデル + vocabファイルなどが準備出来ている方はスキップしてください。)
今回はHugging Faceに上がっている、gpt2をRust用に変換してみました。
※既にrust_model.otが準備されていますが、練習用ということで...

bash

// git-lfsが入っていない人用
sudo apt install git-lfs

// モデルのダウンロード
git lfs install
git clone https://huggingface.co/gpt2

Rust用に変換

PyTorch用のpytorch_model.binはそのままでは使えないので、Rust用のrust_model.otに変換する必要があります。
参照
Loading pretrained and custom model weights

まずはrust-bertのリポジトリを取ってきます。

bash

git clone https://github.com/guillaume-be/rust-bert.git

このようなディレクトリ構造を想定しています。

bash

workspace
    ├── gpt2/(もしくは自前のmodel)
    │    └──...
    └── rust-bert/
            ├──...
            ├──utils/
            ├──...

さっそく変換していきます。
(torchなどのモジュールが足りないなどは各自インストールしてください。)

bash

python ./rust-bert/utils/convert_model.py ./gpt2/pytorch_model.bin

convert_model.pyの中を見たらわかりますが、cargo runが実行されており、必要なクレートのダウンロード + ビルドが自動で行われてます。

bash

ls gpt2/
... ... ... rust_model.ot ... ... ...

やったー！！！(gpt2には既に用意されてたので、上書きしただけですが。)

まとめ

今回はHugginfaceのgpt2を変換してみましたが、rinna社のjapanese-gpt2-smallも同様に変換できました。

次の記事では実際にTextGenerationタスクをしてみます。

後半はこちらから