Ahogrammer

Deep Dive Into NLP, ML and Cloud

Transformersの事前学習済みモデルをダウンロードして保存する

先日、Transformersの事前学習済みモデルをダウンロードして保存する必要があったのでその方法を紹介。 結論として、モデルのダウンロードはおなじみのfrom_pretrainedで行い、保存はsave_pretrainedで行う。 ちなみに、これをしたかった理由は、モデルをS3にアップロードしてSageMakerのエンドポイントにデプロイしたかったため。

from transformers import AutoTokenizer, AutoModelForQuestionAnswering

# Download a model and a tokenizer.
tokenizer = AutoTokenizer.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
model = AutoModelForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')

# Save the model and the tokenizer.
dir_name = 'model'
tokenizer.save_pretrained(dir_name)
model.save_pretrained(dir_name)

こうしてダウンロードしたモデルとトークナイザを以下のコマンドでtar.gzに圧縮する。こうして圧縮したファイルをS3にアップロードすることで簡単にデプロイすることができる。

cd model/ && tar -zcvf ../model.tar.gz . && cd ..

参考

github.com

stackoverflow.com