はじめに
最近のAIの進歩は目覚ましく、大規模言語モデルを使用したアプリケーションが次々と登場しています。その多くはサーバーサイドでの実行が主流でしたが、WebGPUとTransformers.jsの登場により、ブラウザ上で完結するAIチャットアプリケーションの実装が可能になりました。
この実装はtransformers.js-examplesのdeepseek-r1-webgpuを参考にtypescriptで実装してました。
大変だった点は以下
ソースコードは下記にあります。
完成イメージ
deepseek-r1の1.5BをWebGPUで動かしてみました。ブラウザですがモデルファイルはローカルにダウンロードして動いてます。ファイルサイズは1.28GB。一番軽量なモデルとはいえ因数分解が解けるぞ… pic.twitter.com/MRF3iJMXhU
— よなか (@yonaka158) February 2, 2025
完成すると、以下のような機能を持つアプリケーションが作成できます：
- ブラウザ上でAIモデルを実行
- リアルタイムのストリーミング応答
- 進捗表示とエラーハンドリング
- 思考プロセスと回答の分離表示
環境要件
- Node.js 18以上
- WebGPU対応ブラウザ（Chrome Canary, Edge Canary等）
- npm 9以上
- GPUメモリ 2GB以上推奨
使用する技術スタック
- React + TypeScript：UIフレームワークと型システム
- Vite：高速な開発環境
- Tailwind CSS：スタイリング
- Transformers.js：AIモデルの実行
- WebGPU：GPU加速による推論
- shadcn/ui：モダンなUIコンポーネント
実装の解説
1. プロジェクトのセットアップ
まず、必要なディレクトリ構造とパッケージをセットアップします：
# プロジェクトの作成
mkdir deepseek-r1-webgpu-ts
cd deepseek-r1-webgpu-ts
npm create vite@latest . -- --template react-ts
# 依存パッケージのインストール
npm install @huggingface/transformers@^3.3.1 \
better-react-mathjax@^2.0.3 \
dompurify@^3.2.3 \
marked@^15.0.5 \
lucide-react
# 開発用パッケージのインストール
npm install -D tailwindcss postcss autoprefixer \
@typescript-eslint/eslint-plugin \
@typescript-eslint/parser \
eslint-config-prettier \
prettier
2. 型定義の実装
TypeScriptの利点を最大限活かすため、まず型定義から始めます。
// src/types.ts
// メッセージの基本型
export interface Message {
role: "user" | "assistant";
content: string;
}
// 状態付きメッセージ型
export interface MessageState extends Message {
timestamp: string;
status: "complete" | "thinking" | "raw";
}
// 処理済みメッセージ型
export interface ProcessedMessage {
thinking: string;
answer: string;
}
// Workerのステータス
export type WorkerStatus = "initializing" | "ready" | "loading" | "error";
// Worker関連の型定義
type WorkerMessageType = "check" | "load" | "generate" | "interrupt" | "reset";
export interface WorkerMessage {
type: WorkerMessageType;
data?: unknown;
}
// ローディング状態の型
export interface LoadingStatus {
status: "loading" | "initiate" | "progress" | "done" | "error" | "ready";
file?: string;
progress?: number;
total?: number;
data?: string;
error?: string;
}
// 生成状態の型
export interface GenerationStatus {
status: "update" | "start" | "complete" | "error";
output?: string | string[];
tps?: number;
numTokens?: number;
state?: "thinking" | "answering";
error?: string;
}
3. WebGPUとTransformers.jsの統合
Web Workerを使用してモデルの実行を行います。これにより、メインスレッドのブロックを防ぎ、UIの応答性を保ちます。
// src/worker.ts（主要な部分のみ抜粋）
class TextGenerationPipeline {
private static model_id = "onnx-community/DeepSeek-R1-Distill-Qwen-1.5B-ONNX";
private static tokenizer: Promise<PreTrainedTokenizer>;
private static model: Promise<PreTrainedModel>;
static async getInstance(): Promise<[PreTrainedTokenizer, PreTrainedModel]> {
try {
this.tokenizer ??= AutoTokenizer.from_pretrained(this.model_id, {
progress_callback: reportProgress,
});
this.model ??= AutoModelForCausalLM.from_pretrained(this.model_id, {
dtype: "q4f16",
device: "webgpu",
progress_callback: reportProgress,
});
return Promise.all([this.tokenizer, this.model]);
} catch (error) {
debugLog("Error in getInstance:", error);
throw error;
}
}
}
4. パフォーマンス最適化
シングルトンパターンの活用
モデルとトークナイザーの管理にシングルトンパターンを採用することで、メモリ使用を最適化し、初期化のオーバーヘッドを削減しています：
// インスタンスの取得（一度だけ初期化される）
const [tokenizer, model] = await TextGenerationPipeline.getInstance();
Web Workerによる非同期処理
UIのブロックを防ぐため、重い処理はWeb Workerで実行します：
// src/App.tsx
useEffect(() => {
const worker = new Worker(new URL("./worker.ts", import.meta.url), {
type: "module",
});
worker.addEventListener("message", handleWorkerMessage);
worker.addEventListener("error", handleWorkerError);
worker.postMessage({ type: "check" });
return () => {
worker.terminate();
};
}, []);
5. エラーハンドリングとデバッグ
開発時のデバッグを容易にするため、詳細なログ機能を実装しています：
// src/worker.ts
const debugLog = (message: string, ...args: any[]): void => {
const timestamp = new Date().toISOString();
console.log(`[Worker Debug ${timestamp}] ${message}`, ...args);
};
6. UI/UXの実装
モダンなUIを実現するため、Tailwind CSSとshadcn/uiを組み合わせています：
// src/components/ProgressBar.tsx
const ProgressBar: React.FC<ProgressBarProps> = ({
progress,
fileName,
fileSize,
}) => {
return (
<div className="p-4 bg-white dark:bg-gray-900">
<div className="flex justify-between mb-1">
<span className="text-sm font-medium text-gray-700 dark:text-gray-200">
{fileName}
</span>
<span className="text-sm text-gray-600 dark:text-gray-400">
{progress}% of {fileSize}GB
</span>
</div>
<div className="w-full bg-gray-200 rounded-full h-2.5 dark:bg-gray-700">
<div
className="bg-blue-600 h-2.5 rounded-full transition-all duration-300"
style={{ width: `${progress}%` }}
/>
</div>
</div>
);
};
実装のポイント
-
型安全性の確保
- 厳格な型定義による開発時のエラー防止
- 型推論を活用した保守性の向上
-
パフォーマンスの最適化
- シングルトンパターンによるメモリ管理
- Web Workerによる非同期処理
- useCallbackとuseRefによるメモ化
-
エラーハンドリング
- 詳細なエラーメッセージ
- デバッグログの実装
- エラー境界の設定
-
ユーザー体験
- ストリーミング応答
- プログレスバー表示
- ダークモード対応
- レスポンシブデザイン
デプロイ時の注意点
-
WebGPU対応の確認
const IS_WEBGPU_AVAILABLE = !!navigator.gpu; -
必要なメモリ容量
- モデルサイズ：約1.5GB
- 推奨メモリ：2GB以上
-
ブラウザの要件
- Chrome/Edge Canary
- WebGPUフラグの有効化
まとめ
本実装では、以下の特徴を持つAIチャットアプリケーションを実現しました：
- TypeScriptによる型安全な実装
- WebGPUによる高速な推論処理
- Web Workerによる非同期処理
- モダンなUIデザイン
特に、エラーハンドリング、デバッグ機能、UI/UXの面で、オリジナルの実装から大きな改善を加えています。