Livebook で Bumblebee の BERT による文章の穴埋めを実行する

はじめに

Bumblebee の公式サンプルに従ってモデルを実行してみるシリーズです

今回は BERT と言う自然言語処理技術を使って、文章の穴埋めを実行します

このシリーズの記事

• 画像分類: ResNet50
• 画像生成: Stable Diffusion
• 文章の穴埋め: BERT　(ここ)
• 文章の判別: BERTweet
• 文章の生成: GPT2
• 質疑応答: RoBERTa
• 固有名詞の抽出: bert-base-NER

Bumblebee の公式サンプル

実装の全文はこちら

実行環境

• MacBook Pro 13 inchi
  • 2.4 GHz クアッドコアIntel Core i5
  • 16 GB 2133 MHz LPDDR3
• macOS Ventura 13.0.1
• Rancher Desktop 1.6.2
  • メモリ割り当て 12 GB
  • CPU 割り当て 6 コア

Livebook 0.8.0 の Docker　イメージを元にしたコンテナで動かしました

コンテナ定義はこちらを参照

セットアップ

必要なモジュールをインストールして EXLA.Backend で Nx が動くようにします

Mix.install(
  [
    {:bumblebee, "~> 0.1"},
    {:nx, "~> 0.4"},
    {:exla, "~> 0.4"},
    {:kino, "~> 0.8"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

コンテナで動かしている場合、キャッシュディレクトリーを指定した方が都合がいいです

※詳細はこの記事を見てください

cache_dir = "/tmp/bumblebee_cache"

モデルのダウンロード

モデルファイルを Haggin Face からダウンロードしてきて読み込みます

必要な場合は cache_dir を指定します

{:ok, bert} =
  Bumblebee.load_model({
    :hf,
    "bert-base-uncased",
    cache_dir: cache_dir
  })

{:ok, tokenizer} =
  Bumblebee.load_tokenizer({
    :hf,
    "bert-base-uncased",
    cache_dir: cache_dir
  })

tokenizer は文章をテンソルに変換するための前処理定義です

サービスの提供

Bumblebee.Text.fill_mask で文章穴埋めサービスを提供します

serving = Bumblebee.Text.fill_mask(bert, tokenizer)

マスクされた文章の準備

テキスト入力を作り、文章を入力します

AI に穴埋めさせたいところを [MASK] にしてください

text_input = Kino.Input.text("マスクされた文章", default: "The most important thing in life is [MASK].")

入力された文章を取得します

text = Kino.Input.read(text_input)

推論の実行

推論してデータテーブルに結果を表示します

serving
|> Nx.Serving.run(text)
|> then(&Kino.DataTable.new(&1.predictions))

他のモデル

Hugging Face のリポジトリーIDを変えれば他のモデルでも実行できます

serve_model = fn repository_id ->
  {:ok, model} =
    Bumblebee.load_model({
      :hf,
      repository_id,
      cache_dir: cache_dir
    })

  {:ok, tokenizer} =
    Bumblebee.load_tokenizer({
      :hf,
      repository_id,
      cache_dir: cache_dir
    })

  Bumblebee.Text.fill_mask(model, tokenizer)
end

"albert-base-v2"
|> serve_model.()
|> Nx.Serving.run(text)
|> then(&Kino.DataTable.new(&1.predictions))

"roberta-base"
|> serve_model.()
|> Nx.Serving.run(text)
|> then(&Kino.DataTable.new(&1.predictions))

モデルによって意見が違っていますね

まとめ

テキスト系の AI も簡単に動かせました