MLXでT5モデルを利用する

株式会社United Code

T5とは？

T5（Text-To-Text Transfer Transformer）は、Google Researchが開発した自然言語処理（NLP）モデルです。T5は、「Text-to-Text」の名前の通り、あらゆるNLPタスクをテキスト入力からテキスト出力への形式に統一することを目指しています。これにより、翻訳、要約、質問応答、感情分析など、様々なタスクに同じアーキテクチャを利用できるようになります。

この記事ではMLXで、T5を利用してみたいと思います。
環境はApple M1のMacBookProです。

インストール

cd t5 
pip install -r requirements.txt

モデルのダウンロード

convert.pyスクリプトを使用して、指定されたモデルを変換し、MLXが読み取ることができるようにします。

モデルは下記から指定できます。

  • t5-small

  • t5-base

  • t5-large

  • t5-3b

  • t5-11b

今回はt5-smallにしました。

python convert.py --model t5-small

テキスト生成

モデルのダウンロードが完了したら、T5モデルを使ってテキストを生成できるようになります。

サンプルプロンプトを実施してみます。

python t5.py --model t5-small --prompt "translate English to German: A tasty apple"

下記のような結果が出力されました。

Input: translate English to German: A tasty apple
Ein leckerer Apfel
Time: 0.78 seconds, tokens/s: 8.92