はじめに

最近の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の面で、オリジナルの実装から大きな改善を加えています。

参考リンク