Cloudflare Constellationで日本語言語モデルを動かしたかった

2023年5月にCloudflare Constellationが発表されました。これを使うと、エッジアプリケーションでAIモデルを動かすことができるようになります。デモとして、画像分類アプリケーションが動かせるようになっていますが、これを使って日本語の言語モデルを動かそうとしてみました。しかし、ちょっとうまくいきませんでした。開発の途中でも気をつけるべきところがいくつかあったので、本稿に書き残します。

なおCloudflare Constellationは2023年6月現在プライベートベータ版であり、動かすにはwaitlistに登録が必要です。また以下に書いてあることは正式リリース時には変更されている可能性があります。

環境

• cloudflare/constellation: 0.0.12
• xenova/transformers: 2.3.0

実装

2023/05/15の公式日本語ブログ記事に記載の構築手順は一部足りない情報やtomlのパースエラーが含まれていたりするので参考にせず、公式ドキュメントの方を参考にして準備していきます。

Cloudflare Constellationで使えるモデルの大きさには現状厳しい制限がありますが、6月のアップデートで50MBまでのモデルがアップロードできるようになりました。京都大学の公開しているBERTモデルであるku-nlp/deberta-v2-tiny-japanese-char-wwmであればサイズが小さいので、載せてみました。

実装したコードの全体像は以下のようになりました。

src/index.ts

import { BertTokenizer } from '@xenova/transformers';

import tokenizer_json from './my_fast_onnx_bert/tokenizer.json';
import tokenizer_config from './my_fast_onnx_bert/tokenizer_config.json';

// 中略

export default {
  async fetch(request: Request, env: Env): Promise<Response> {
    const input = request.headers.get("text");
    const tokenizer = await new BertTokenizer(tokenizer_json, tokenizer_config);

    const { input_ids, attention_mask, token_type_ids } = await tokenizer(input, {return_tensor: false});

    // convert ? (344) to tokenizer.mask_token_id (4)
    input_ids[input_ids.indexOf(344)] = tokenizer.mask_token_id;

    const getTensor = ids => new Tensor("int64", [1, ids.length], ids);
    const output = await run(
      env.BIND,
      "id",
      [getTensor(input_ids), getTensor(attention_mask), getTensor(token_type_ids)]
    );

    const logits = output.logits;
    const mask_token_index = input_ids.indexOf(tokenizer.mask_token_id);

    const predicteds = logits.getItem(0).getItem(mask_token_index).value;
    const predicted_token_id = predicteds.indexOf(Math.max(...predicteds));
    const answer = tokenizer.decode([predicted_token_id]);

    return new Response(JSON.stringify({
      input,
      answer,
    }));
  },
};

これで$ npx wrangler dev --remoteで起動して、適当な入力をしてみたのですが

あまり正しい出力のようには見えません。

実際には以下のように、人をあらわすような名詞っぽいものが並んでくれるはずです。

動作するまでにつまづいた点

モデルのアップロード

Cloudflare Constellationは現状ONNXモデルをアップロードする必要があります。
Hugging Face Optimumでrinna GPTモデルをONNXに変換するなどを参考に変換してからアップロードします。$ wrangler constellation model uploadでアップロードできます。

トークナイザーの初期化

Transfoermers.jsはFastTokenizer形式のみを受け入れるようです。convert_slow_tokenizerで変換して、tokenizer.jsonを入手します。

適当なPython環境

from transformers import convert_slow_tokenizer
MODEL_NAME = 'ku-nlp/deberta-v2-tiny-japanese-char-wwm'
tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)
tokenizer = convert_slow_tokenizer.convert_slow_tokenizer(tokenizer)

tokenizer.save("tokenizer.json")

tokenizer.jsonとtokenizer_config.jsonを配置して、読み込みます。

src/index.ts

import { BertTokenizer } from '@xenova/transformers';

import tokenizer_json from './my_fast_onnx_bert/tokenizer.json';
import tokenizer_config from './my_fast_onnx_bert/tokenizer_config.json';

const tokenizer = await new BertTokenizer(tokenizer_json, tokenizer_config);

mask tokenのトークナイズ

Transformers.jsはmask tokenのトークナイズがまだ実装されていないようで、[MASK]という入力を与えてもmask tokenとして扱ってくれません。

とりあえず?をmask tokenとして扱うことにします。以下コードで?(344)をmask_token_id(4)に自力で変換しています。

const { input_ids } = await tokenizer(input, {return_tensor: false});
input_ids[input_ids.indexOf(344)] = tokenizer.mask_token_id;

なお、このコードはmask tokenが2つ以上ある場合に対応していません。

run()メソッドへの渡し方

モデルにはinput_ids, attention_maskといった複数の引数を渡す必要があります。普通に配列として並べれば良いようです。

const output = await run(
  env.BIND,
  "id",
  [input_ids, attention_mask, token_type_ids]
);

なお実際には、配列の中身はCloudflare ConstellationのTensorクラスに変換する必要があります。

まとめ

Cloudflare Constellationはまだベータ版であり、ドキュメントが不足していたり動作が不安定だったりします。正式リリースに期待しましょう。特にモデルサイズの制限が緩和されれば、いま流行りのGPTなどのサイズの大きな言語モデルも動かせるようになり、使い道が広がると思います。
またTransformers.jsは、Python版の本家Transformersとはインターフェースが違ったり機能が不足していたりします。これらの実装が進めば、Python版Transformersで言語モデルを扱うのと同じような感覚でCloudflare Constellationを動かしやすくなると思います。