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

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

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

## 環境
- cloudflare/constellation: 0.0.12
- xenova/transformers: 2.3.0

## 実装
[2023/05/15の公式日本語ブログ記事](https://blog.cloudflare.com/ja-jp/introducing-constellation-ja-jp/)に記載の構築手順は一部足りない情報やtomlのパースエラーが含まれていたりするので参考にせず、[公式ドキュメント](https://developers.cloudflare.com/constellation/get-started/first-constellation-worker/)の方を参考にして準備していきます。

Cloudflare Constellationで使えるモデルの大きさには現状厳しい制限がありますが、[6月のアップデート](https://blog.cloudflare.com/globally-distributed-ai-and-a-constellation-update/)で50MBまでのモデルがアップロードできるようになりました。[京都大学の公開しているBERTモデルであるku\-nlp/deberta\-v2\-tiny\-japanese\-char\-wwm](https://huggingface.co/ku-nlp/deberta-v2-tiny-japanese-char-wwm)であればサイズが小さいので、載せてみました。

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

```typescript: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`で起動して、適当な入力をしてみたのですが

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/50391/95beb5fa-2444-9db4-c1ee-c7f17535343f.png)

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

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/50391/b553c33e-c94c-bb3e-63ba-a5eac57c6730.png)

## 動作するまでにつまづいた点
### モデルのアップロード
Cloudflare Constellationは現状ONNXモデルをアップロードする必要があります。
[Hugging Face Optimumでrinna GPTモデルをONNXに変換する](https://qiita.com/suzuki_sh/items/7c0fb56f24948ea8e6ed)などを参考に変換してからアップロードします。`$ wrangler constellation model upload`でアップロードできます。

### トークナイザーの初期化
[Transfoermers.jsはFastTokenizer形式のみを受け入れる](https://github.com/xenova/transformers.js/issues/93#issuecomment-1512990056)ようです。[convert_slow_tokenizer](https://github.com/huggingface/transformers/blob/main/src/transformers/convert_slow_tokenizer.py)で変換して、`tokenizer.json`を入手します。

```python:適当な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`を配置して、読み込みます。

```ts: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のトークナイズがまだ実装されていない](https://github.com/xenova/transformers.js/blob/aceab9bf3d1be9dadf6601b8f775e519376258e3/src/tokenizers.js#L1980)ようで、`[MASK]`という入力を与えてもmask tokenとして扱ってくれません。

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

```ts
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`といった複数の引数を渡す必要があります。普通に配列として並べれば良いようです。

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

なお実際には、配列の中身は[Cloudflare ConstellationのTensorクラス](https://developers.cloudflare.com/constellation/platform/client-api/#tensor-class)に変換する必要があります。

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