ブラウザでDeepSeek-R1を動作させる【TypeScript & WebGPU実装】

## はじめに

最近のAIの進歩は目覚ましく、大規模言語モデルを使用したアプリケーションが次々と登場しています。その多くはサーバーサイドでの実行が主流でしたが、WebGPUとTransformers.jsの登場により、**ブラウザ上で完結するAIチャットアプリケーション**の実装が可能になりました。

この実装は[transformers.js-examples](https://github.com/huggingface/transformers.js-examples/)の[deepseek-r1-webgpu](https://github.com/huggingface/transformers.js-examples/tree/main/deepseek-r1-webgpu)を参考にtypescriptで実装してました。

大変だった点は以下

https://qiita.com/yonaka15/items/a49d0737de39541c1ce1

ソースコードは下記にあります。

https://github.com/yonaka15/deepseek-r1-webgpu-ts


### 完成イメージ

<blockquote class="twitter-tweet" data-media-max-width="560"><p lang="ja" dir="ltr">deepseek-r1の1.5BをWebGPUで動かしてみました。ブラウザですがモデルファイルはローカルにダウンロードして動いてます。ファイルサイズは1.28GB。一番軽量なモデルとはいえ因数分解が解けるぞ… <a href="https://t.co/MRF3iJMXhU">pic.twitter.com/MRF3iJMXhU</a></p>&mdash; よなか (@yonaka158) <a href="https://twitter.com/yonaka158/status/1885922054647754806?ref_src=twsrc%5Etfw">February 2, 2025</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

完成すると、以下のような機能を持つアプリケーションが作成できます：

- ブラウザ上で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. プロジェクトのセットアップ

まず、必要なディレクトリ構造とパッケージをセットアップします：

```bash
# プロジェクトの作成
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の利点を最大限活かすため、まず型定義から始めます。

```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の応答性を保ちます。

```typescript
// 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. パフォーマンス最適化

#### シングルトンパターンの活用

モデルとトークナイザーの管理にシングルトンパターンを採用することで、メモリ使用を最適化し、初期化のオーバーヘッドを削減しています：

```typescript
// インスタンスの取得（一度だけ初期化される）
const [tokenizer, model] = await TextGenerationPipeline.getInstance();
```

#### Web Workerによる非同期処理

UIのブロックを防ぐため、重い処理はWeb Workerで実行します：

```typescript
// 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. エラーハンドリングとデバッグ

開発時のデバッグを容易にするため、詳細なログ機能を実装しています：

```typescript
// 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を組み合わせています：

```tsx
// 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>
  );
};
```

## 実装のポイント

1. **型安全性の確保**
   - 厳格な型定義による開発時のエラー防止
   - 型推論を活用した保守性の向上

2. **パフォーマンスの最適化**
   - シングルトンパターンによるメモリ管理
   - Web Workerによる非同期処理
   - useCallbackとuseRefによるメモ化

3. **エラーハンドリング**
   - 詳細なエラーメッセージ
   - デバッグログの実装
   - エラー境界の設定

4. **ユーザー体験**
   - ストリーミング応答
   - プログレスバー表示
   - ダークモード対応
   - レスポンシブデザイン

## デプロイ時の注意点

1. **WebGPU対応の確認**
   ```typescript
   const IS_WEBGPU_AVAILABLE = !!navigator.gpu;
   ```

2. **必要なメモリ容量**
   - モデルサイズ：約1.5GB
   - 推奨メモリ：2GB以上

3. **ブラウザの要件**
   - Chrome/Edge Canary
   - WebGPUフラグの有効化

## まとめ

本実装では、以下の特徴を持つAIチャットアプリケーションを実現しました：

- TypeScriptによる型安全な実装
- WebGPUによる高速な推論処理
- Web Workerによる非同期処理
- モダンなUIデザイン

特に、エラーハンドリング、デバッグ機能、UI/UXの面で、オリジナルの実装から大きな改善を加えています。

## 参考リンク

- [実装レポジトリ](https://github.com/yonaka15/deepseek-r1-webgpu-ts)
- [Transformers.js Examples](https://github.com/huggingface/transformers.js-examples/)
- [DeepSeek-R1 WebGPU Example](https://github.com/huggingface/transformers.js-examples/tree/main/deepseek-r1-webgpu)
- [ONNX DeepSeek-R1 Model](https://huggingface.co/onnx-community/DeepSeek-R1-Distill-Qwen-1.5B-ONNX)