feat(rsvp): use jieba tokenizer to segment words for Chinese books (#3985)

2026-04-29 00:19:05 +08:00
parent 4b0720a3e3
commit 920627ae59
8 changed files with 245 additions and 12 deletions
@@ -33,7 +33,7 @@
    "clippy:check": "cargo clippy -p Readest --no-deps -- -D warnings",
    "format": "pnpm -w format",
    "format:check": "pnpm -w format:check",
-    "prepare-public-vendor": "mkdirp ./public/vendor/pdfjs ./public/vendor/simplecc",
+    "prepare-public-vendor": "mkdirp ./public/vendor/pdfjs ./public/vendor/simplecc ./public/vendor/jieba",
    "copy-pdfjs-js": "cpx \"../../packages/foliate-js/node_modules/pdfjs-dist/legacy/build/{pdf.worker.min.mjs,pdf.min.mjs,pdf.d.mts}\" ./public/vendor/pdfjs",
    "copy-pdfjs-wasm": "cpx \"../../packages/foliate-js/node_modules/pdfjs-dist/wasm/{openjpeg.wasm,qcms_bg.wasm}\" ./public/vendor/pdfjs",
    "copy-pdfjs-fonts": "cpx \"../../packages/foliate-js/node_modules/pdfjs-dist/{cmaps,standard_fonts}/*\" ./public/vendor/pdfjs",
@@ -42,9 +42,11 @@
    "copy-flatten-pdfjs-css": "pnpm copy-flatten-pdfjs-annotation-layer-css && pnpm copy-flatten-pdfjs-text-layer-css",
    "copy-pdfjs": "pnpm copy-pdfjs-js && pnpm copy-pdfjs-wasm && pnpm copy-pdfjs-fonts && pnpm copy-flatten-pdfjs-css",
    "copy-simplecc": "cpx \"../../packages/simplecc-wasm/dist/web/*\" ./public/vendor/simplecc",
+    "copy-jieba": "cpx \"./node_modules/jieba-wasm/pkg/web/*\" ./public/vendor/jieba",
    "setup-pdfjs": "pnpm prepare-public-vendor && pnpm copy-pdfjs",
    "setup-simplecc": "pnpm prepare-public-vendor && pnpm copy-simplecc",
-    "setup-vendors": "pnpm setup-pdfjs && pnpm setup-simplecc",
+    "setup-jieba": "pnpm prepare-public-vendor && pnpm copy-jieba",
+    "setup-vendors": "pnpm setup-pdfjs && pnpm setup-simplecc && pnpm setup-jieba",
    "build-win-x64": "dotenv -e .env.tauri.local -- tauri build --target i686-pc-windows-msvc --bundles nsis",
    "build-win-arm64": "dotenv -e .env.tauri.local -- tauri build --target aarch64-pc-windows-msvc --bundles nsis",
    "build-linux-x64": "dotenv -e .env.tauri.local -- tauri build --target x86_64-unknown-linux-gnu --bundles appimage",
@@ -144,6 +146,7 @@
    "iso-639-2": "^3.0.2",
    "iso-639-3": "^3.0.1",
    "isomorphic-ws": "^5.0.0",
+    "jieba-wasm": "^2.4.0",
    "js-md5": "^0.8.3",
    "jwt-decode": "^4.0.0",
    "lucide-react": "^0.562.0",
@@ -0,0 +1,139 @@
+import { describe, test, expect, beforeAll } from 'vitest';
+// Import the web build directly. Vitest runs in Node, which would otherwise
+// pick the `node` export — that build has no async `init()` and instantiates
+// the WASM via Node FS. In production (Next.js / Tauri) the `browser` export
+// is selected and matches what `src/utils/jieba.ts` uses.
+import init, { cut, cut_all, cut_for_search, tokenize } from 'jieba-wasm/web';
+import { readFile } from 'fs/promises';
+import { join } from 'path';
+
+describe.concurrent('jieba-wasm', () => {
+  beforeAll(async () => {
+    const wasmPath = join(process.cwd(), 'public/vendor/jieba/jieba_rs_wasm_bg.wasm');
+    const wasmBuffer = await readFile(wasmPath);
+    await init({ module_or_path: wasmBuffer });
+  });
+
+  test('cut - canonical README example', () => {
+    expect(cut('我来到北京清华大学', true)).toEqual(['我', '来到', '北京', '清华大学']);
+  });
+
+  test('cut - HMM detects new compounds', () => {
+    // 杭研 isn't in the dict; HMM should infer it as a compound.
+    const tokens = cut('他来到了网易杭研大厦', true);
+    expect(tokens).toContain('网易');
+    expect(tokens).toContain('杭研');
+    expect(tokens).toContain('大厦');
+    expect(tokens.join('')).toBe('他来到了网易杭研大厦');
+  });
+
+  test('cut_all returns all possible word combinations', () => {
+    const tokens = cut_all('我来到北京清华大学');
+    expect(tokens).toContain('清华');
+    expect(tokens).toContain('清华大学');
+    expect(tokens).toContain('华大');
+    expect(tokens).toContain('大学');
+  });
+
+  test('cut_for_search produces finer-grained tokens', () => {
+    const tokens = cut_for_search('小明硕士毕业于中国科学院计算所', true);
+    expect(tokens).toContain('小明');
+    expect(tokens).toContain('硕士');
+    expect(tokens).toContain('毕业');
+    expect(tokens).toContain('中国');
+    expect(tokens).toContain('科学');
+    expect(tokens).toContain('科学院');
+    expect(tokens).toContain('中国科学院');
+    expect(tokens).toContain('计算');
+    expect(tokens).toContain('计算所');
+  });
+
+  test('tokenize returns tokens with start/end offsets', () => {
+    const tokens = tokenize('永和服装饰品有限公司', 'default', true);
+    expect(tokens.length).toBeGreaterThan(0);
+    for (const tok of tokens) {
+      expect(typeof tok.word).toBe('string');
+      expect(typeof tok.start).toBe('number');
+      expect(typeof tok.end).toBe('number');
+      expect(tok.end).toBeGreaterThan(tok.start);
+    }
+    // Recombining tokens by offset reproduces the original string.
+    const combined = tokens
+      .slice()
+      .sort((a, b) => a.start - b.start)
+      .map((t) => t.word)
+      .join('');
+    expect(combined).toBe('永和服装饰品有限公司');
+  });
+
+  describe('long passage from a Chinese book', () => {
+    const sample =
+      '文章采访了一名州警，他从理论上说明这些所谓的“无名车祸”有许多是起因于车内的昆虫：' +
+      '黄蜂、蜜蜂，甚至也可能是蜘蛛或蛾子。驾驶人惊慌了，想要用力拍打虫子，或是摇下车窗让虫子出去。' +
+      '很有可能是虫子蜇了他，也或许驾驶就是失去控制。无论如何轰然一声巨响……一切结束。' +
+      '而那只昆虫，通常安然无恙，快活地嗡嗡叫着飞出冒烟失事的车外，找寻更适合的场所。';
+
+    test('cut preserves the original characters when joined back', () => {
+      const tokens = cut(sample, true);
+      expect(tokens.join('')).toBe(sample);
+    });
+
+    test('cut splits the passage into a reasonable number of tokens', () => {
+      const tokens = cut(sample, true);
+      // Char length is ~170 — token count should be substantially less but
+      // still on the order of dozens, never one token per char.
+      expect(tokens.length).toBeGreaterThan(50);
+      expect(tokens.length).toBeLessThan(sample.length);
+    });
+
+    test('cut recognizes domain-specific words', () => {
+      const tokens = cut(sample, true);
+      const set = new Set(tokens);
+      const expected = [
+        '文章',
+        '采访',
+        '一名',
+        '理论',
+        '所谓',
+        '车祸',
+        '车内',
+        '昆虫',
+        '黄蜂',
+        '蜜蜂',
+        '蜘蛛',
+        '蛾子',
+        '驾驶',
+        '惊慌',
+        '用力',
+        '拍打',
+        '虫子',
+        '车窗',
+        '失去',
+        '控制',
+        '无论如何',
+        '一声',
+        '巨响',
+        '结束',
+        '通常',
+        '安然无恙',
+        '快活地',
+        '嗡嗡叫',
+        '冒烟',
+        '失事',
+        '找寻',
+        '场所',
+      ];
+      for (const word of expected) {
+        expect(set.has(word), `expected token "${word}" in cut output`).toBe(true);
+      }
+    });
+
+    test('cut keeps punctuation as standalone tokens', () => {
+      const tokens = cut(sample, true);
+      expect(tokens).toContain('，');
+      expect(tokens).toContain('。');
+      expect(tokens).toContain('：');
+      expect(tokens).toContain('、');
+    });
+  });
+});
@@ -18,6 +18,7 @@ import { useTranslation } from '@/hooks/useTranslation';
 import { BookNote, PageInfo } from '@/types/book';
 import { TOCItem } from '@/libs/document';
 import { Insets } from '@/types/misc';
+import { initJieba } from '@/utils/jieba';
 import RSVPOverlay from './RSVPOverlay';
 import RSVPStartDialog from './RSVPStartDialog';

@@ -207,15 +208,28 @@ const RSVPControl: React.FC<RSVPControlProps> = ({ bookKey, gridInsets }) => {
        return;
      }

+      const primaryLanguage = bookData.book.primaryLanguage;
+
      // Create controller if not exists
      if (!controllerRef.current) {
-        controllerRef.current = new RSVPController(view, bookKey);
+        controllerRef.current = new RSVPController(view, bookKey, primaryLanguage);
        rsvpSectionRef.current = view.renderer.primaryIndex;
        rsvpChapterHrefRef.current = progress?.sectionHref ?? null;
+      } else {
+        controllerRef.current.setPrimaryLanguage(primaryLanguage);
      }

      const controller = controllerRef.current;

+      // For Chinese books, preload jieba-wasm so that the synchronous word
+      // extractor can use it. Done before requestStart() so the loader has
+      // the dialog's interaction time to fetch ~3.8MB of WASM.
+      if (primaryLanguage?.toLowerCase().startsWith('zh')) {
+        initJieba().catch((e) => {
+          console.warn('Failed to initialize jieba-wasm; falling back to Intl.Segmenter:', e);
+        });
+      }
+
      // Seed localStorage from cloud-synced BookConfig when this device has no local position.
      // This restores RSVP progress saved on another device after a sync.
      const configPos = getConfig(bookKey)?.rsvpPosition;
@@ -242,7 +242,7 @@ const RSVPOverlay: React.FC<RSVPOverlayProps> = ({
  const wordAfter = currentWord ? currentWord.text.substring(currentWord.orpIndex + 1) : '';
  const isCJKWord = currentWord ? containsCJK(currentWord.text) : false;
  const wordLetterSpacing = undefined;
-  const wordSideOffset = isCJKWord ? '0.5em' : '0.3em';
+  const wordSideOffset = isCJKWord ? '0.45em' : '0.3em';

  // Time remaining calculation
  const getTimeRemaining = useCallback((): string | null => {
@@ -20,6 +20,7 @@ export class RSVPController extends EventTarget {
  private view: FoliateView;
  private bookId: string; // Book hash without session suffix, for persistent storage
  private currentCfi: string | null = null;
+  private primaryLanguage: string | undefined;

  private state: RsvpState = {
    active: false,
@@ -39,15 +40,23 @@ export class RSVPController extends EventTarget {
  private countdown: number | null = null;
  private cachedWords: { docIndex: number; doc: Document; words: RsvpWord[] } | null = null;

-  constructor(view: FoliateView, bookKey: string) {
+  constructor(view: FoliateView, bookKey: string, primaryLanguage?: string) {
    super();
    this.view = view;
    // Extract book ID (hash) from bookKey format: "{hash}-{sessionId}"
    // Use only the hash for persistent position storage across sessions
    this.bookId = bookKey.split('-')[0] || bookKey;
+    this.primaryLanguage = primaryLanguage;
    this.loadSettings();
  }

+  setPrimaryLanguage(lang: string | undefined): void {
+    if (this.primaryLanguage === lang) return;
+    this.primaryLanguage = lang;
+    // Language changes invalidate the segmentation result.
+    this.cachedWords = null;
+  }
+
  private loadSettings(): void {
    const savedWpm = this.loadWpmFromStorage();
    if (savedWpm) {
@@ -738,7 +747,7 @@ export class RSVPController extends EventTarget {
    const walk = (node: Node): void => {
      if (node.nodeType === Node.TEXT_NODE) {
        const text = node.textContent || '';
-        const nodeWords = splitTextIntoWords(text);
+        const nodeWords = splitTextIntoWords(text, this.primaryLanguage);

        let offset = 0;
        for (const word of nodeWords) {
@@ -1,6 +1,7 @@
 /**
 * Utility functions for CJK (Chinese, Japanese, Korean) text processing
 */
+import { cutZh, isJiebaReady } from '@/utils/jieba';

 /**
 * Check if a character is a CJK character
@@ -72,9 +73,20 @@ export function getSegmenterLocale(text: string): string | null {
 }

 /**
- * Segment CJK text into words using Intl.Segmenter with punctuation attachment
+ * Segment CJK text into words using Intl.Segmenter with punctuation attachment.
+ * If `language` starts with `zh` and jieba-wasm has been initialized
+ * (see `initJieba` in @/utils/jieba), use it for higher-quality Chinese
+ * segmentation.
 */
-export function segmentCJKText(text: string): string[] {
+export function segmentCJKText(text: string, language?: string): string[] {
+  if (language?.toLowerCase().startsWith('zh') && isJiebaReady()) {
+    try {
+      return segmentWithJieba(text);
+    } catch (error) {
+      console.warn('jieba-wasm failed, falling back to Intl.Segmenter:', error);
+    }
+  }
+
  // Try to use Intl.Segmenter for semantic word segmentation
  if (typeof Intl !== 'undefined' && 'Segmenter' in Intl) {
    try {
@@ -193,10 +205,30 @@ export function getHyphenParts(word: string): string[] {
  return result;
 }

+/**
+ * Segment Chinese text using jieba. Whitespace and standalone punctuation
+ * runs are attached to the previous token so that downstream pause logic
+ * still sees punctuation at word boundaries.
+ */
+function segmentWithJieba(text: string): string[] {
+  const tokens = cutZh(text);
+  const words: string[] = [];
+  for (const token of tokens) {
+    if (!token) continue;
+    if (token.trim() === '') continue;
+    if (isCJKPunctuation(token) && words.length > 0) {
+      words[words.length - 1] = words[words.length - 1] + token;
+      continue;
+    }
+    words.push(token);
+  }
+  return words;
+}
+
 /**
 * Split text into words, handling both CJK and non-CJK text
 */
-export function splitTextIntoWords(text: string): string[] {
+export function splitTextIntoWords(text: string, language?: string): string[] {
  const hasCJK = containsCJK(text);

  if (!hasCJK) {
@@ -238,7 +270,7 @@ export function splitTextIntoWords(text: string): string[] {
      if (currentSegment) {
        if (inCJKSequence) {
          // Segment the CJK text (with any trailing punctuation)
-          words.push(...segmentCJKText(currentSegment));
+          words.push(...segmentCJKText(currentSegment, language));
        } else {
          words.push(currentSegment);
        }
@@ -249,7 +281,7 @@ export function splitTextIntoWords(text: string): string[] {
      // Non-CJK, non-punctuation, non-whitespace character
      if (inCJKSequence && currentSegment) {
        // Segment the CJK text before continuing with non-CJK
-        words.push(...segmentCJKText(currentSegment));
+        words.push(...segmentCJKText(currentSegment, language));
        currentSegment = '';
      }
      currentSegment += char;
@@ -259,7 +291,7 @@ export function splitTextIntoWords(text: string): string[] {

  if (currentSegment) {
    if (inCJKSequence) {
-      words.push(...segmentCJKText(currentSegment));
+      words.push(...segmentCJKText(currentSegment, language));
    } else {
      words.push(currentSegment);
    }
@@ -0,0 +1,28 @@
+import init, { cut } from 'jieba-wasm';
+
+let initialized = false;
+let initPromise: Promise<void> | null = null;
+
+const initJieba = async (): Promise<void> => {
+  if (initialized) return;
+  if (!initPromise) {
+    initPromise = (async () => {
+      try {
+        await init('/vendor/jieba/jieba_rs_wasm_bg.wasm');
+        initialized = true;
+      } catch (e) {
+        initPromise = null;
+        throw e;
+      }
+    })();
+  }
+  return initPromise;
+};
+
+const isJiebaReady = (): boolean => initialized;
+
+const cutZh = (text: string): string[] => {
+  return cut(text, true);
+};
+
+export { initJieba, isJiebaReady, cutZh };
@@ -279,6 +279,9 @@ importers:
      isomorphic-ws:
        specifier: ^5.0.0
        version: 5.0.0(ws@8.19.0)
+      jieba-wasm:
+        specifier: ^2.4.0
+        version: 2.4.0
      js-md5:
        specifier: ^0.8.3
        version: 0.8.3
@@ -6031,6 +6034,9 @@ packages:
    resolution: {integrity: sha512-7vuh85V5cdDofPyxn58nrPjBktZo0u9x1g8WtjQol+jZDaE+fhN+cIvTj11GndBnMnyfrUOG1sZQxCdjKh+DKg==}
    engines: {node: '>= 10.13.0'}

+  jieba-wasm@2.4.0:
+    resolution: {integrity: sha512-ZvQdS+FGifrFXZIXSgOyOgEz+1wdy1P4vSvwe37FVtku9ycSdHTZbHqF5i9tMN1JucoAmeiLBeI6/YaqcGD+KA==}
+
  jiti@1.21.7:
    resolution: {integrity: sha512-/imKNG4EbWNrVjoNC/1H5/9GFy+tqjGBHCaSsN+P2RnPqjsLmv6UD3Ej+Kj8nBWaRAwyk7kK5ZUc+OEatnTR3A==}
    hasBin: true
@@ -15163,6 +15169,8 @@ snapshots:
      merge-stream: 2.0.0
      supports-color: 8.1.1

+  jieba-wasm@2.4.0: {}
+
  jiti@1.21.7: {}

  jiti@2.6.1: {}