Transformers.jsで日本語言語モデルが動作するWeb APIを作る

Transformers.jsを使うと、機械学習モデルをJavaScript環境で動作させることができます。本稿ではFastifyを使って、Transformers.jsが動作するWeb APIを作ってみます。

この記事は、以下の記事と同じものをNode.jsで再現を試みたものになります。もし実際にTransformersが動作するWeb APIを開発したいという場合は、Node.jsではなく素直にPythonで本家Transformersを使うほうが良いでしょう。

環境

• Node.js 20
• @xenova/transformers 2.3.1
• fastify 4.19.2

また環境構築にはDockerを用いていて、WORKDIR /appしています。Dockerfile等も含めたすべてのプログラムはGitHubにあります。

実装

モデルの準備

京都大学の公開しているBERTモデルであるku-nlp/deberta-v2-tiny-japanese-char-wwmをONNXに変換して、適当なディレクトリに格納します。ここでは/app/my_fast_onnx_bert/とします。

$ tree my_fast_onnx_bert/
my_fast_onnx_bert/
├── config.json
├── onnx
│   └── model_quantized.onnx
├── tokenizer.json
└── tokenizer_config.json

fastifyプログラムの実装

FastifyのQuick Startを参考にプログラムを実装します。

『Cloudflare Constellationで日本語言語モデルを動かしたかった』で述べたように、Transformers.jsはmask tokenのトークナイズがまだ実装されていないようなので、ID=4にあたる?をID=344にあたるmask tokenと変換して扱うことにしています。このプログラムはmaskが2つ以上ある入力には対応していません。

server.js

import Fastify from 'fastify'
import { env } from '@xenova/transformers';
import { BertForMaskedLM, BertTokenizer, Tensor } from '@xenova/transformers';

import tokenizer_json from './my_fast_onnx_bert/tokenizer.json' assert { type: "json" };
import tokenizer_config from './my_fast_onnx_bert/tokenizer_config.json' assert { type: "json" };

env.remoteModels = false;
env.localModelPath = '/app/';

const tokenizer = await new BertTokenizer(tokenizer_json, tokenizer_config);
const model = await BertForMaskedLM.from_pretrained('./my_fast_onnx_bert');

const fastify = Fastify({
  logger: true
});

fastify.get('/', async function handler (request, reply) {
  const input = request.query.input;
  const inputs = await tokenizer(input);

  const input_ids = inputs.input_ids[0].tolist();
  const attention_mask = inputs.attention_mask[0].tolist();
  const token_type_ids = inputs.token_type_ids[0].tolist();

  // convert ? (344) to tokenizer.mask_token_id (4)
  // https://github.com/xenova/transformers.js/blob/aceab9bf3d1be9dadf6601b8f775e519376258e3/src/tokenizers.js#L1980
  input_ids[input_ids.indexOf(BigInt(344))] = BigInt(tokenizer.mask_token_id);

  const toTensor = ids => new Tensor("int64", ids, [1, ids.length]);

  const { logits } = await model({
      input_ids: toTensor(input_ids),
      attention_mask: toTensor(attention_mask),
      token_type_ids: toTensor(token_type_ids),
  });

  const mask_token_index = input_ids.indexOf(BigInt(tokenizer.mask_token_id));

  const predicteds = logits[0][mask_token_index].tolist();

  // max_id
  const predicted_token_id = predicteds.indexOf(Math.max(...predicteds));
  const answer = tokenizer.decode([predicted_token_id]);

  return { input: input, answer: answer };
});

try {
  await fastify.listen({ port: 3000, host: "0.0.0.0" });
} catch (err) {
  fastify.log.error(err);
  process.exit(1);
}

実行

$ node serverで実行します。DockerであればDockerfileにCMD node serverしておいて、3000番ポートを空けて起動します。

http://localhost:3000/?input=吾輩は?である。名前はまだ無い にアクセスすると、以下のようなレスポンスが返って来ます。

response

{"input":"吾輩は?である。名前はまだ無い","answer":"名"}

それっぽい回答ではあります。これにてNode.jsのAPIサーバで、日本語言語モデルが動作しました。