見出し画像

Transformers.js によるブラウザ上で動く「日→英」音声翻訳の実装

kawamou

Transformers.js の検証がてらローカルのブラウザで完結する音声翻訳を実装した。下記は実行結果であり、動画上マイク音声が入っていないが「日本語音声入力 → テキスト化 → 翻訳 → 英語音声出力」ができている ↓


Transformers.js

Hugging Face 所属の @xenova 筆頭に開発が進む Transformers.js。

Python の Transformers ライブラリと同等の機能を持ち、ブラウザ上で機械学習モデルを実行できる JavaScript ライブラリである。

精力的なアップデートは v3 アップデートに向けた PR を見ると分かりやすい。v3 に向け例えば WebGPU 対応も着々と進行している。

closes: WebGPU Support

実装編

概要

今回の音声翻訳は：

の 3 段構成にした。各モデルをそれぞれロードし利用する。マルチモーダルに音声入力 to 音声出力できる変態モデル (e.g. GPT-4o) が跋扈する時代ではナイーブに過ぎるがまずは素朴に組んだ。@xenova が公開しているモデルを今回は採用したが、基本的には ONNX に変換したモデルであれば実行可能なので独自選定も容易い。

またアプリケーションフレームワークとしては Remix を採用した。

ソースコード

かなり雑に書いているので後ほど直す。

実装上の注意

実装上の注意点は Building a React application を確認するのが良い。個人的注意点は ① Web Worker を利用すること ② シングルトンにすること だった。

まず Web Worker だが、ML モデルのロードや推論等々は処理が重いので UI スレッドとは別で処理したい。今回は Shopify の react-web-worker を利用した。Building a React application では Web Worker とのメッセージが入り組んでいるが、react-web-worker は useWorker を利用し直感的に書ける。

続いてシングルトンについて。ロードされたモデルはデフォルトでブラウザの Cache API を利用し保存される。

DevTools > Applicaton > Cache storage にいる

Okay yes, it looks like you are creating multiple versions of the model, leading to an out of memory error. The current Caching API caches the model, but does not force you to only create one instance of the model.

https://github.com/xenova/transformers.js/issues/8#issuecomment-1458867413

Cache API によってモデルはキャッシュされるがインスタンスが複数個できる場合があり、メモリ不足に陥るのでシングルトンにせよとのこと。

export const getSttPipeline =
	async (): Promise<AutomaticSpeechRecognitionPipeline> => {
		if (sttPipelineInstance === null) {
			const pipelineInstance = await pipeline(
				"automatic-speech-recognition",
				"Xenova/whisper-small",
			);
			sttPipelineInstance =
				pipelineInstance as AutomaticSpeechRecognitionPipeline;
		}
		return sttPipelineInstance;
	};

export const stt = async (audio: Float32Array) => {
	const stt = await getSTTPipelineSingleton(() => {});
	const result = await stt(audio, {
		language: "ja",
		task: "transcribe",
	});
	return result;
};

感想

これだけ直感的かつ簡単に ML モデルがブラウザ上で動くのはとても良い。@xenova がデモとともにコードを定期的に公開してるのでキャッチアップやりやすいのも好印象。

Future Work

まだ検証でハローワールドしただけなのでいくつか追加で検証ポイントがある。

メモリ関連

複数のモデルを動かすと途中幾度となくメモリ関連のエラーに遭遇した。良く見たのは RangeError: offset is out of bounds とか。色々 Issue 眺めたが結局各モデルを小さくしたら解消した。当初は翻訳にリッチなモデル Xenova/nllb-200-distilled-600M を利用していたが Issue#280 を参考にマルチ翻訳モデルは止めたり。今回は検証だから問題ないが精度目指す場合大きなモデルも必要になるので動かす方法を調査する必要がある。

実行時間

お気付きの方もいるかもだが冒頭の動画は推論部分を倍速にしている。当方 M2 Max 96GB だがそれでも推論に少し時間がかかる。ただ @xenova のデモを見る限りかなり高速動作に見え v3 で WebGPU フル対応になると更に高速化可能だろうか？少なくとも Whisper についてはこのブランチの対応を v3 使えば試せそう（やってみる）。

いいなと思ったら応援しよう！