From 37a90b734de953989f8c16bd0aeb52fa42259cc3 Mon Sep 17 00:00:00 2001 From: Richard R Date: Mon, 18 May 2026 04:26:39 -0600 Subject: [PATCH] feat(pdf): migrate to PP-DocLayoutV3 ONNX model and update PDF block taxonomy Switch PDF layout parsing to use the PP-DocLayoutV3 ONNX model, replacing the previous Docling-based approach. Update all environment variables, model fetching, and manifest handling for the new model and its artifacts. Refactor block kind taxonomy throughout the codebase, tests, and UI to align with PP-DocLayoutV3 labels, including expanded and renamed block types. Revise document settings, block filtering, and stitching logic to support the new set of block kinds. Update documentation and environment variable references to reflect the model transition. BREAKING CHANGE: PDF parsing now requires PP-DocLayoutV3 ONNX model and updated environment variables; block kind names and settings have changed throughout the system. --- .env.example | 7 +- Dockerfile | 2 +- README.md | 2 +- docs-site/docs/deploy/local-development.md | 2 +- .../docs/reference/environment-variables.md | 31 +- scripts/fetch-models.mjs | 30 +- src/components/documents/DocumentSettings.tsx | 25 +- src/components/views/PDFViewer.tsx | 26 +- src/lib/client/audiobooks/adapters/pdf.ts | 4 +- src/lib/server/pdf-layout/ensureModel.ts | 69 +++-- .../server/pdf-layout/mergeTextWithRegions.ts | 24 +- src/lib/server/pdf-layout/model/LICENSE.txt | 6 +- src/lib/server/pdf-layout/model/manifest.json | 26 +- src/lib/server/pdf-layout/parsePdf.ts | 2 +- src/lib/server/pdf-layout/runLayoutModel.ts | 282 ++++++++++++------ .../pdf-layout/stitchCrossPageBlocks.ts | 21 +- src/lib/shared/document-settings.ts | 18 +- src/types/document-settings.ts | 2 +- src/types/parsed-pdf.ts | 54 +++- tests/unit/pdf-audiobook-adapter.spec.ts | 14 +- tests/unit/pdf-build-page-text.spec.ts | 10 +- .../unit/pdf-merge-text-with-regions.spec.ts | 4 +- .../unit/pdf-stitch-cross-page-blocks.spec.ts | 21 +- 23 files changed, 464 insertions(+), 218 deletions(-) diff --git a/.env.example b/.env.example index 8115116..ec650c1 100644 --- a/.env.example +++ b/.env.example @@ -88,8 +88,11 @@ OPENREADER_COMPUTE_MODE=local # OPENREADER_COMPUTE_WORKER_URL= # OPENREADER_COMPUTE_WORKER_TOKEN= -# Optional override for Docling layout model download URL -# OPENREADER_DOCLING_MODEL_URL=https://huggingface.co/ds4sd/docling-layout-heron/resolve/main/model.onnx +# Optional overrides for PDF layout model artifacts +# OPENREADER_PDF_LAYOUT_MODEL_URL=https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/PP-DocLayoutV3.onnx +# OPENREADER_PDF_LAYOUT_MODEL_DATA_URL=https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/PP-DocLayoutV3.onnx.data +# OPENREADER_PDF_LAYOUT_CONFIG_URL=https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/config.json +# OPENREADER_PDF_LAYOUT_PREPROCESSOR_URL=https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/preprocessor_config.json # (Optional) Override ffmpeg binary path used for audiobook processing FFMPEG_BIN= diff --git a/Dockerfile b/Dockerfile index a4d6c81..764b387 100644 --- a/Dockerfile +++ b/Dockerfile @@ -73,7 +73,7 @@ COPY --from=app-builder /app/THIRD_PARTY_LICENSES /licenses # Include SeaweedFS license text for the copied weed binary. COPY --from=seaweedfs-builder /tmp/SeaweedFS-LICENSE.txt /licenses/SeaweedFS-LICENSE.txt # Include static model notices for runtime-downloaded assets. -COPY src/lib/server/pdf-layout/model/LICENSE.txt /licenses/docling-layout-LICENSE.txt +COPY src/lib/server/pdf-layout/model/LICENSE.txt /licenses/pp-doclayoutv3-LICENSE.txt # Copy the compiled whisper.cpp build output into the runtime image # (includes whisper-cli and its shared libraries, e.g. libwhisper.so, libggml.so) diff --git a/README.md b/README.md index 22e0590..210e215 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ OpenReader is an open source, self-host-friendly text-to-speech document reader - 🎯 **Multi-provider TTS** with OpenAI-compatible endpoints and cloud providers (Kokoro-FastAPI, KittenTTS-FastAPI, Orpheus-FastAPI or OpenAI, Replicate, DeepInfra). - 📖 **Read-along playback** for PDF/EPUB with sentence-aware narration. - ⏱️ **Word-by-word highlighting** via optional `whisper.cpp` timestamps (`OPENREADER_COMPUTE_MODE=local` + `WHISPER_CPP_BIN`). -- 🧱 **Layout-aware PDF parsing** (Docling ONNX) with structured blocks for cleaner TTS/chaptering. +- 🧱 **Layout-aware PDF parsing** (PP-DocLayoutV3 ONNX) with structured blocks for cleaner TTS/chaptering. - 🛜 **Sync + library import** to bring docs across devices and from server-mounted folders. - 🗂️ **Flexible storage** with embedded SeaweedFS or external S3-compatible backends. - 🎧 **Audiobook export** in `m4b`/`mp3` with resumable chapter processing. diff --git a/docs-site/docs/deploy/local-development.md b/docs-site/docs/deploy/local-development.md index 93d144f..875c83f 100644 --- a/docs-site/docs/deploy/local-development.md +++ b/docs-site/docs/deploy/local-development.md @@ -261,7 +261,7 @@ For all environment variables, see [Environment Variables](../reference/environm ::: :::tip Optional model prefetch -To pre-populate the Docling ONNX model cache before first PDF parse, run `pnpm fetch-models`. +To pre-populate the PDF layout ONNX model cache before first PDF parse, run `pnpm fetch-models`. ::: See [Auth](../configure/auth) for app/auth behavior. diff --git a/docs-site/docs/reference/environment-variables.md b/docs-site/docs/reference/environment-variables.md index e023ba8..f9d43cc 100644 --- a/docs-site/docs/reference/environment-variables.md +++ b/docs-site/docs/reference/environment-variables.md @@ -56,7 +56,10 @@ For auth-enabled deployments, use **Settings → Admin** as the primary source o | `OPENREADER_COMPUTE_MODE` | Heavy compute backend | `local` | Set to `none` to disable whisper alignment + PDF layout parsing | | `OPENREADER_COMPUTE_WORKER_URL` | Heavy compute backend | unset | Reserved for future worker backend mode (`worker`) | | `OPENREADER_COMPUTE_WORKER_TOKEN` | Heavy compute backend | unset | Reserved for future worker backend mode (`worker`) | -| `OPENREADER_DOCLING_MODEL_URL` | PDF layout model | Docling HF URL | Override ONNX download URL for `ensureModel()` | +| `OPENREADER_PDF_LAYOUT_MODEL_URL` | PDF layout model | PP-DocLayoutV3 ONNX URL | Override ONNX model URL for `ensureModel()` | +| `OPENREADER_PDF_LAYOUT_MODEL_DATA_URL` | PDF layout model | PP-DocLayoutV3 ONNX data URL | Override ONNX external data URL for `ensureModel()` | +| `OPENREADER_PDF_LAYOUT_CONFIG_URL` | PDF layout model | PP-DocLayoutV3 config URL | Override model config URL for `ensureModel()` | +| `OPENREADER_PDF_LAYOUT_PREPROCESSOR_URL` | PDF layout model | PP-DocLayoutV3 preprocessor URL | Override model preprocessor URL for `ensureModel()` | | `WHISPER_CPP_BIN` | Word timing (local mode) | unset | Set to enable `whisper.cpp` timestamps in `OPENREADER_COMPUTE_MODE=local` | | `FFMPEG_BIN` | Audio runtime | auto-detected (`ffmpeg-static`) | Override ffmpeg binary path | @@ -374,13 +377,31 @@ Reserved bearer token for future external compute worker mode. - Used only when `OPENREADER_COMPUTE_MODE=worker` (not implemented in v1) - Leave unset in v1 -### OPENREADER_DOCLING_MODEL_URL +### OPENREADER_PDF_LAYOUT_MODEL_URL -Override URL for the Docling ONNX layout model downloaded by `ensureModel()`. +Override URL for the PP-DocLayoutV3 ONNX model downloaded by `ensureModel()`. -- Default: `https://huggingface.co/ds4sd/docling-layout-heron/resolve/main/model.onnx` +- Default: `https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/PP-DocLayoutV3.onnx` - Optional for custom mirrors/air-gapped workflows -- You can also pre-populate the model cache via `pnpm fetch-models` + +### OPENREADER_PDF_LAYOUT_MODEL_DATA_URL + +Override URL for the PP-DocLayoutV3 ONNX external data file downloaded by `ensureModel()`. + +- Default: `https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/PP-DocLayoutV3.onnx.data` + +### OPENREADER_PDF_LAYOUT_CONFIG_URL + +Override URL for the PP-DocLayoutV3 `config.json` downloaded by `ensureModel()`. + +- Default: `https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/config.json` + +### OPENREADER_PDF_LAYOUT_PREPROCESSOR_URL + +Override URL for the PP-DocLayoutV3 `preprocessor_config.json` downloaded by `ensureModel()`. + +- Default: `https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/preprocessor_config.json` +- You can pre-populate the model cache via `pnpm fetch-models` ### WHISPER_CPP_BIN diff --git a/scripts/fetch-models.mjs b/scripts/fetch-models.mjs index 24cd5ac..edf56c2 100644 --- a/scripts/fetch-models.mjs +++ b/scripts/fetch-models.mjs @@ -3,13 +3,21 @@ import { mkdir, writeFile } from 'node:fs/promises'; import path from 'node:path'; const modelDir = path.join(process.cwd(), 'docstore', 'model'); -const modelPath = path.join(modelDir, 'docling-layout-heron.onnx'); -const configPath = path.join(modelDir, 'docling-layout-heron.config.json'); -const licensePath = path.join(modelDir, 'docling-layout-heron.LICENSE.txt'); +const modelPath = path.join(modelDir, 'PP-DocLayoutV3.onnx'); +const modelDataPath = path.join(modelDir, 'PP-DocLayoutV3.onnx.data'); +const configPath = path.join(modelDir, 'pp-doclayoutv3.config.json'); +const preprocessorPath = path.join(modelDir, 'pp-doclayoutv3.preprocessor_config.json'); +const licensePath = path.join(modelDir, 'pp-doclayoutv3.LICENSE.txt'); const staticLicensePath = path.join(process.cwd(), 'src', 'lib', 'server', 'pdf-layout', 'model', 'LICENSE.txt'); -const modelUrl = process.env.OPENREADER_DOCLING_MODEL_URL || 'https://huggingface.co/docling-project/docling-layout-heron-onnx/resolve/main/model.onnx'; -const configUrl = process.env.OPENREADER_DOCLING_CONFIG_URL || 'https://huggingface.co/docling-project/docling-layout-heron-onnx/resolve/main/config.json'; +const modelUrl = process.env.OPENREADER_PDF_LAYOUT_MODEL_URL + || 'https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/PP-DocLayoutV3.onnx'; +const modelDataUrl = process.env.OPENREADER_PDF_LAYOUT_MODEL_DATA_URL + || 'https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/PP-DocLayoutV3.onnx.data'; +const configUrl = process.env.OPENREADER_PDF_LAYOUT_CONFIG_URL + || 'https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/config.json'; +const preprocessorUrl = process.env.OPENREADER_PDF_LAYOUT_PREPROCESSOR_URL + || 'https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/preprocessor_config.json'; await mkdir(modelDir, { recursive: true }); @@ -19,12 +27,24 @@ if (!modelRes.ok) { } await writeFile(modelPath, new Uint8Array(await modelRes.arrayBuffer())); +const modelDataRes = await fetch(modelDataUrl); +if (!modelDataRes.ok) { + throw new Error(`Failed to fetch model data: ${modelDataRes.status} ${modelDataRes.statusText}`); +} +await writeFile(modelDataPath, new Uint8Array(await modelDataRes.arrayBuffer())); + const configRes = await fetch(configUrl); if (!configRes.ok) { throw new Error(`Failed to fetch config: ${configRes.status} ${configRes.statusText}`); } await writeFile(configPath, new Uint8Array(await configRes.arrayBuffer())); +const preprocessorRes = await fetch(preprocessorUrl); +if (!preprocessorRes.ok) { + throw new Error(`Failed to fetch preprocessor config: ${preprocessorRes.status} ${preprocessorRes.statusText}`); +} +await writeFile(preprocessorPath, new Uint8Array(await preprocessorRes.arrayBuffer())); + const staticLicense = await import('node:fs/promises').then((m) => m.readFile(staticLicensePath)); await writeFile(licensePath, staticLicense); diff --git a/src/components/documents/DocumentSettings.tsx b/src/components/documents/DocumentSettings.tsx index d63916b..bdb534c 100644 --- a/src/components/documents/DocumentSettings.tsx +++ b/src/components/documents/DocumentSettings.tsx @@ -21,10 +21,27 @@ import { RefreshIcon } from '@/components/icons/Icons'; import type { ParsedPdfBlockKind, PdfParseStatus } from '@/types/parsed-pdf'; const PDF_SKIP_KIND_OPTIONS: Array<{ kind: ParsedPdfBlockKind; label: string }> = [ - { kind: 'page-header', label: 'Page headers' }, - { kind: 'page-footer', label: 'Page footers' }, - { kind: 'footnote', label: 'Footnotes' }, - { kind: 'caption', label: 'Captions' }, + { kind: 'header', label: 'Header' }, + { kind: 'footer', label: 'Footer' }, + { kind: 'footnote', label: 'Footnote' }, + { kind: 'vision_footnote', label: 'Vision footnote' }, + { kind: 'figure_title', label: 'Figure title' }, + { kind: 'doc_title', label: 'Document title' }, + { kind: 'paragraph_title', label: 'Paragraph title' }, + { kind: 'abstract', label: 'Abstract' }, + { kind: 'algorithm', label: 'Algorithm' }, + { kind: 'aside_text', label: 'Aside text' }, + { kind: 'content', label: 'Content' }, + { kind: 'reference', label: 'Reference' }, + { kind: 'reference_content', label: 'Reference content' }, + { kind: 'text', label: 'Text' }, + { kind: 'number', label: 'Number' }, + { kind: 'formula', label: 'Formula' }, + { kind: 'formula_number', label: 'Formula number' }, + { kind: 'table', label: 'Table' }, + { kind: 'chart', label: 'Chart' }, + { kind: 'image', label: 'Image' }, + { kind: 'seal', label: 'Seal' }, ]; const viewTypeTextMapping = [ diff --git a/src/components/views/PDFViewer.tsx b/src/components/views/PDFViewer.tsx index 1cc8e1a..b11447e 100644 --- a/src/components/views/PDFViewer.tsx +++ b/src/components/views/PDFViewer.tsx @@ -354,14 +354,26 @@ export function PDFViewer({ zoomLevel, pdfState }: PDFViewerProps) { const colorForKind = (kind: ParsedPdfBlock['kind']): string => { switch (kind) { - case 'section-header': return 'rgba(34,197,94,0.20)'; - case 'title': return 'rgba(16,185,129,0.20)'; - case 'caption': return 'rgba(245,158,11,0.20)'; + case 'paragraph_title': return 'rgba(34,197,94,0.20)'; + case 'doc_title': return 'rgba(16,185,129,0.20)'; + case 'figure_title': return 'rgba(245,158,11,0.20)'; case 'table': return 'rgba(59,130,246,0.20)'; - case 'picture': return 'rgba(139,92,246,0.20)'; - case 'page-header': - case 'page-footer': - case 'footnote': return 'rgba(239,68,68,0.20)'; + case 'chart': + case 'image': return 'rgba(139,92,246,0.20)'; + case 'header': + case 'footer': + case 'footnote': + case 'vision_footnote': return 'rgba(239,68,68,0.20)'; + case 'formula': + case 'formula_number': return 'rgba(251,146,60,0.22)'; + case 'abstract': + case 'algorithm': + case 'aside_text': + case 'content': + case 'reference': + case 'reference_content': + case 'text': + case 'number': return 'rgba(14,165,233,0.18)'; default: return 'rgba(14,165,233,0.18)'; } }; diff --git a/src/lib/client/audiobooks/adapters/pdf.ts b/src/lib/client/audiobooks/adapters/pdf.ts index 4b52581..6c38631 100644 --- a/src/lib/client/audiobooks/adapters/pdf.ts +++ b/src/lib/client/audiobooks/adapters/pdf.ts @@ -60,6 +60,8 @@ function prepareParsedChapters({ let currentTitle = 'Introduction'; let currentBlocks: ParsedPdfBlock[] = []; + const chapterBoundaryKinds = new Set(['paragraph_title', 'doc_title']); + const flush = () => { if (!currentBlocks.length) return; const text = chapterTextFromBlocks(currentBlocks, smartSentenceSplitting, maxBlockLength); @@ -74,7 +76,7 @@ function prepareParsedChapters({ }; for (const block of allBlocks) { - if (block.kind === 'section-header') { + if (chapterBoundaryKinds.has(block.kind)) { flush(); currentTitle = block.text.trim() || `Chapter ${chapters.length + 1}`; currentBlocks.push(block); diff --git a/src/lib/server/pdf-layout/ensureModel.ts b/src/lib/server/pdf-layout/ensureModel.ts index 70d5c73..e39dd62 100644 --- a/src/lib/server/pdf-layout/ensureModel.ts +++ b/src/lib/server/pdf-layout/ensureModel.ts @@ -4,12 +4,16 @@ import { access, mkdir, rename, writeFile, readFile, unlink, copyFile } from 'fs import { DOCSTORE_DIR } from '@/lib/server/storage/library-mount'; import manifest from '@/lib/server/pdf-layout/model/manifest.json'; -const DEFAULT_MODEL_URL = 'https://huggingface.co/docling-project/docling-layout-heron-onnx/resolve/main/model.onnx'; -const DEFAULT_CONFIG_URL = 'https://huggingface.co/docling-project/docling-layout-heron-onnx/resolve/main/config.json'; +const DEFAULT_MODEL_URL = 'https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/PP-DocLayoutV3.onnx'; +const DEFAULT_MODEL_DATA_URL = 'https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/PP-DocLayoutV3.onnx.data'; +const DEFAULT_CONFIG_URL = 'https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/config.json'; +const DEFAULT_PREPROCESSOR_URL = 'https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main/preprocessor_config.json'; const MODEL_DIR = path.join(DOCSTORE_DIR, 'model'); -const MODEL_PATH = path.join(MODEL_DIR, 'docling-layout-heron.onnx'); -const CONFIG_PATH = path.join(MODEL_DIR, 'docling-layout-heron.config.json'); -const LICENSE_PATH = path.join(MODEL_DIR, 'docling-layout-heron.LICENSE.txt'); +export const MODEL_PATH = path.join(MODEL_DIR, 'PP-DocLayoutV3.onnx'); +export const MODEL_DATA_PATH = path.join(MODEL_DIR, 'PP-DocLayoutV3.onnx.data'); +export const MODEL_CONFIG_PATH = path.join(MODEL_DIR, 'pp-doclayoutv3.config.json'); +export const MODEL_PREPROCESSOR_PATH = path.join(MODEL_DIR, 'pp-doclayoutv3.preprocessor_config.json'); +const LICENSE_PATH = path.join(MODEL_DIR, 'pp-doclayoutv3.LICENSE.txt'); const STATIC_LICENSE_PATH = path.join(process.cwd(), 'src/lib/server/pdf-layout/model/LICENSE.txt'); let inflight: Promise | null = null; @@ -51,15 +55,22 @@ async function verifyFile(pathToFile: string, manifestPath: string): Promise { await copyFile(STATIC_LICENSE_PATH, LICENSE_PATH); if (!(await verifyFile(LICENSE_PATH, 'LICENSE.txt'))) { - throw new Error('Docling model license checksum verification failed'); + throw new Error('PDF layout model license checksum verification failed'); } } async function ensureModelInternal(): Promise { try { await access(MODEL_PATH); - await access(CONFIG_PATH); - if (await verifyFile(MODEL_PATH, 'model.onnx') && await verifyFile(CONFIG_PATH, 'config.json')) { + await access(MODEL_DATA_PATH); + await access(MODEL_CONFIG_PATH); + await access(MODEL_PREPROCESSOR_PATH); + if ( + await verifyFile(MODEL_PATH, 'model.onnx') + && await verifyFile(MODEL_DATA_PATH, 'model.onnx.data') + && await verifyFile(MODEL_CONFIG_PATH, 'config.json') + && await verifyFile(MODEL_PREPROCESSOR_PATH, 'preprocessor_config.json') + ) { await ensureLicense(); return MODEL_PATH; } @@ -68,24 +79,44 @@ async function ensureModelInternal(): Promise { } await mkdir(MODEL_DIR, { recursive: true }); - const tmpPath = `${MODEL_PATH}.tmp`; - const configTmpPath = `${CONFIG_PATH}.tmp`; - const modelUrl = process.env.OPENREADER_DOCLING_MODEL_URL?.trim() || DEFAULT_MODEL_URL; - const configUrl = process.env.OPENREADER_DOCLING_CONFIG_URL?.trim() || DEFAULT_CONFIG_URL; + const modelTmpPath = `${MODEL_PATH}.tmp`; + const modelDataTmpPath = `${MODEL_DATA_PATH}.tmp`; + const configTmpPath = `${MODEL_CONFIG_PATH}.tmp`; + const preprocessorTmpPath = `${MODEL_PREPROCESSOR_PATH}.tmp`; + const modelUrl = process.env.OPENREADER_PDF_LAYOUT_MODEL_URL?.trim() + || DEFAULT_MODEL_URL; + const modelDataUrl = process.env.OPENREADER_PDF_LAYOUT_MODEL_DATA_URL?.trim() + || DEFAULT_MODEL_DATA_URL; + const configUrl = process.env.OPENREADER_PDF_LAYOUT_CONFIG_URL?.trim() + || DEFAULT_CONFIG_URL; + const preprocessorUrl = process.env.OPENREADER_PDF_LAYOUT_PREPROCESSOR_URL?.trim() + || DEFAULT_PREPROCESSOR_URL; - await downloadToFile(modelUrl, tmpPath); - if (!(await verifyFile(tmpPath, 'model.onnx'))) { - await unlink(tmpPath).catch(() => undefined); - throw new Error('Docling model checksum verification failed'); + await downloadToFile(modelUrl, modelTmpPath); + if (!(await verifyFile(modelTmpPath, 'model.onnx'))) { + await unlink(modelTmpPath).catch(() => undefined); + throw new Error('PDF layout model checksum verification failed'); + } + await downloadToFile(modelDataUrl, modelDataTmpPath); + if (!(await verifyFile(modelDataTmpPath, 'model.onnx.data'))) { + await unlink(modelDataTmpPath).catch(() => undefined); + throw new Error('PDF layout model external data checksum verification failed'); } await downloadToFile(configUrl, configTmpPath); if (!(await verifyFile(configTmpPath, 'config.json'))) { await unlink(configTmpPath).catch(() => undefined); - throw new Error('Docling model config checksum verification failed'); + throw new Error('PDF layout model config checksum verification failed'); + } + await downloadToFile(preprocessorUrl, preprocessorTmpPath); + if (!(await verifyFile(preprocessorTmpPath, 'preprocessor_config.json'))) { + await unlink(preprocessorTmpPath).catch(() => undefined); + throw new Error('PDF layout model preprocessor checksum verification failed'); } - await rename(tmpPath, MODEL_PATH); - await rename(configTmpPath, CONFIG_PATH); + await rename(modelTmpPath, MODEL_PATH); + await rename(modelDataTmpPath, MODEL_DATA_PATH); + await rename(configTmpPath, MODEL_CONFIG_PATH); + await rename(preprocessorTmpPath, MODEL_PREPROCESSOR_PATH); await ensureLicense(); return MODEL_PATH; } diff --git a/src/lib/server/pdf-layout/mergeTextWithRegions.ts b/src/lib/server/pdf-layout/mergeTextWithRegions.ts index 8b56834..08e0743 100644 --- a/src/lib/server/pdf-layout/mergeTextWithRegions.ts +++ b/src/lib/server/pdf-layout/mergeTextWithRegions.ts @@ -1,15 +1,23 @@ import type { LayoutRegion, PdfTextItem } from '@/lib/server/pdf-layout/types'; -const NON_TEXT_REGION_LABELS = new Set(['picture', 'table']); +const NON_TEXT_REGION_LABELS = new Set(['chart', 'image', 'table', 'seal']); const TEXT_ASSIGNABLE_LABELS = new Set([ - 'title', - 'section-header', - 'paragraph', - 'list-item', - 'caption', - 'page-header', - 'page-footer', + 'abstract', + 'algorithm', + 'aside_text', + 'content', + 'doc_title', + 'figure_title', + 'footer', 'footnote', + 'formula_number', + 'header', + 'number', + 'paragraph_title', + 'reference', + 'reference_content', + 'text', + 'vision_footnote', 'formula', ]); diff --git a/src/lib/server/pdf-layout/model/LICENSE.txt b/src/lib/server/pdf-layout/model/LICENSE.txt index d2b2034..cb60fe1 100644 --- a/src/lib/server/pdf-layout/model/LICENSE.txt +++ b/src/lib/server/pdf-layout/model/LICENSE.txt @@ -1,3 +1,3 @@ -Docling layout model assets are distributed under Apache-2.0 / CDLA terms by the model publisher. -See the upstream release for the authoritative license text: -https://huggingface.co/ds4sd/docling-layout-heron +PP-DocLayoutV3 ONNX model assets are distributed by the model publisher under Apache-2.0. +See upstream for the authoritative license and terms: +https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX diff --git a/src/lib/server/pdf-layout/model/manifest.json b/src/lib/server/pdf-layout/model/manifest.json index c688e28..9f4807d 100644 --- a/src/lib/server/pdf-layout/model/manifest.json +++ b/src/lib/server/pdf-layout/model/manifest.json @@ -1,21 +1,31 @@ { - "name": "docling-layout-heron", - "version": "docling-project/docling-layout-heron-onnx@main", + "name": "pp-doclayoutv3", + "version": "Bei0001/PP-DocLayoutV3-ONNX@main", "files": [ { "path": "model.onnx", - "sha256": "59c81a3a2923042d85034ffc487f8f47e4854117e879aef89b2b9f728fb4922a", - "size": 171220471 + "sha256": "c0721928ff08741bb208ebed539c77170db5234a68cb7e546e6cc9bc172a695b", + "size": 5088167 + }, + { + "path": "model.onnx.data", + "sha256": "34df3e4b79d7bbbf82abce1b4f3cde3d540fa57ad42ec8905c352b97c408d437", + "size": 136774480 }, { "path": "config.json", - "sha256": "c6e67b6cdf64fa245779d0409bb994aa96e8c1790e15ec07a65843628e232180", - "size": 878 + "sha256": "3cf834b91d23a756b1519bce4db42c09e852f3e35c35092dd5a3e253a50c071a", + "size": 2460 + }, + { + "path": "preprocessor_config.json", + "sha256": "519fe0187a43a1ca429e3ad8317bab8700f0d5e8fb3a6e3a0a413ffac078ba42", + "size": 575 }, { "path": "LICENSE.txt", - "sha256": "8f5030e085c59609746fce1010230f12b4a565abb6bdc19a2c2d3735ba1710d9", - "size": 209 + "sha256": "578a6ba6f86b0692a8f719843f575a3eebf4705768ac5c37d149f441208f601f", + "size": 195 } ] } diff --git a/src/lib/server/pdf-layout/parsePdf.ts b/src/lib/server/pdf-layout/parsePdf.ts index e95cc35..84d1703 100644 --- a/src/lib/server/pdf-layout/parsePdf.ts +++ b/src/lib/server/pdf-layout/parsePdf.ts @@ -140,7 +140,7 @@ export async function parsePdf(input: ParsePdfInput): Promise const doc: ParsedPdfDocument = { schemaVersion: 1, documentId: input.documentId, - parserVersion: 'docling-layout-heron@v1+pdfjs@4.8.69', + parserVersion: 'pp-doclayoutv3-onnx@800+pdfjs@4.8.69', parsedAt: Date.now(), pages, }; diff --git a/src/lib/server/pdf-layout/runLayoutModel.ts b/src/lib/server/pdf-layout/runLayoutModel.ts index 608de35..7407ea0 100644 --- a/src/lib/server/pdf-layout/runLayoutModel.ts +++ b/src/lib/server/pdf-layout/runLayoutModel.ts @@ -1,8 +1,7 @@ import * as ort from 'onnxruntime-node'; import { readFile } from 'fs/promises'; -import path from 'path'; import type { LayoutRegion, PdfTextItem } from '@/lib/server/pdf-layout/types'; -import { ensureModel } from '@/lib/server/pdf-layout/ensureModel'; +import { ensureModel, MODEL_CONFIG_PATH, MODEL_PREPROCESSOR_PATH } from '@/lib/server/pdf-layout/ensureModel'; interface RunLayoutInput { pageWidth: number; @@ -11,42 +10,74 @@ interface RunLayoutInput { pagePng: Buffer; } -const INPUT_SIZE = 640; -const MIN_SCORE = 0.6; +const DEFAULT_INPUT_SIZE = 800; +const MIN_SCORE = 0.5; +const CLASS_MIN_SCORE: Partial> = { + header: 0.4, + footer: 0.4, + figure_title: 0.45, + footnote: 0.45, + vision_footnote: 0.45, +}; const LABEL_MAP: Record = { - caption: 'caption', + // PP-DocLayoutV3 labels + abstract: 'abstract', + algorithm: 'algorithm', + aside_text: 'aside_text', + chart: 'chart', + content: 'content', + display_formula: 'formula', + doc_title: 'doc_title', + figure_title: 'figure_title', + footer: 'footer', + footer_image: 'footer', footnote: 'footnote', - formula: 'formula', - list_item: 'list-item', - page_footer: 'page-footer', - page_header: 'page-header', - picture: 'picture', - section_header: 'section-header', + formula_number: 'formula_number', + header: 'header', + header_image: 'header', + image: 'image', + inline_formula: 'formula', + number: 'number', + paragraph_title: 'paragraph_title', + reference: 'reference', + reference_content: 'reference_content', + seal: 'seal', table: 'table', - text: 'paragraph', - title: 'title', - document_index: null, - code: null, - checkbox_selected: null, - checkbox_unselected: null, - form: null, - key_value_region: null, + text: 'text', + vertical_text: 'text', + vision_footnote: 'vision_footnote', }; const MIN_REGION_SIZE: Partial> = { - paragraph: { minWidth: 24, minHeight: 14 }, - 'section-header': { minWidth: 24, minHeight: 14 }, - title: { minWidth: 24, minHeight: 14 }, - 'list-item': { minWidth: 18, minHeight: 12 }, - caption: { minWidth: 18, minHeight: 10 }, + abstract: { minWidth: 24, minHeight: 14 }, + algorithm: { minWidth: 24, minHeight: 14 }, + aside_text: { minWidth: 24, minHeight: 14 }, + content: { minWidth: 24, minHeight: 14 }, + text: { minWidth: 24, minHeight: 14 }, + reference: { minWidth: 24, minHeight: 14 }, + reference_content: { minWidth: 24, minHeight: 14 }, + paragraph_title: { minWidth: 24, minHeight: 14 }, + doc_title: { minWidth: 24, minHeight: 14 }, + number: { minWidth: 18, minHeight: 12 }, + figure_title: { minWidth: 18, minHeight: 10 }, footnote: { minWidth: 18, minHeight: 10 }, - 'page-header': { minWidth: 18, minHeight: 10 }, - 'page-footer': { minWidth: 18, minHeight: 10 }, + vision_footnote: { minWidth: 18, minHeight: 10 }, + header: { minWidth: 18, minHeight: 10 }, + footer: { minWidth: 18, minHeight: 10 }, }; +interface ModelPreprocessor { + inputWidth: number; + inputHeight: number; + rescaleFactor: number; + mean: [number, number, number]; + std: [number, number, number]; +} + let sessionPromise: Promise | null = null; let idToLabelPromise: Promise> | null = null; +let preprocessorPromise: Promise | null = null; let canvasFnsPromise: Promise<{ createCanvasFn: (width: number, height: number) => { getContext: (kind: '2d') => CanvasRenderingContext2D }; loadImageFn: (src: Buffer) => Promise<{ width: number; height: number } & CanvasImageSource>; @@ -96,14 +127,13 @@ async function getSession(): Promise { async function getIdToLabel(): Promise> { if (!idToLabelPromise) { idToLabelPromise = (async () => { - const modelPath = await ensureModel(); - const configPath = path.join(path.dirname(modelPath), 'docling-layout-heron.config.json'); - const raw = await readFile(configPath, 'utf8'); + await ensureModel(); + const raw = await readFile(MODEL_CONFIG_PATH, 'utf8'); const parsed = JSON.parse(raw) as { id2label?: Record }; const out: Record = {}; for (const [key, value] of Object.entries(parsed.id2label ?? {})) { const n = Number(key); - if (Number.isFinite(n)) out[n] = value; + if (Number.isFinite(n)) out[n] = String(value ?? '').trim(); } return out; })(); @@ -111,49 +141,74 @@ async function getIdToLabel(): Promise> { return idToLabelPromise; } -function preprocessLetterboxed( - image: CanvasImageSource, - createCanvasFn: (width: number, height: number) => { getContext: (kind: '2d') => CanvasRenderingContext2D }, -): { - tensor: ort.Tensor; - scale: number; - padX: number; - padY: number; -} { - const sourceWidth = Math.max(1, Number((image as { width?: number }).width ?? INPUT_SIZE)); - const sourceHeight = Math.max(1, Number((image as { height?: number }).height ?? INPUT_SIZE)); - const scale = Math.min(INPUT_SIZE / sourceWidth, INPUT_SIZE / sourceHeight); - const drawWidth = Math.max(1, Math.round(sourceWidth * scale)); - const drawHeight = Math.max(1, Math.round(sourceHeight * scale)); - const padX = Math.floor((INPUT_SIZE - drawWidth) / 2); - const padY = Math.floor((INPUT_SIZE - drawHeight) / 2); +async function getPreprocessor(): Promise { + if (!preprocessorPromise) { + preprocessorPromise = (async () => { + await ensureModel(); + const raw = await readFile(MODEL_PREPROCESSOR_PATH, 'utf8'); + const parsed = JSON.parse(raw) as { + size?: { width?: number; height?: number }; + rescale_factor?: number; + image_mean?: number[]; + image_std?: number[]; + }; - const canvas = createCanvasFn(INPUT_SIZE, INPUT_SIZE); + const inputWidth = Math.max(1, Number(parsed.size?.width ?? DEFAULT_INPUT_SIZE)); + const inputHeight = Math.max(1, Number(parsed.size?.height ?? DEFAULT_INPUT_SIZE)); + const rescaleFactor = Number.isFinite(parsed.rescale_factor) ? Number(parsed.rescale_factor) : (1 / 255); + const mean = [ + Number(parsed.image_mean?.[0] ?? 0), + Number(parsed.image_mean?.[1] ?? 0), + Number(parsed.image_mean?.[2] ?? 0), + ] as [number, number, number]; + const std = [ + Number(parsed.image_std?.[0] ?? 1), + Number(parsed.image_std?.[1] ?? 1), + Number(parsed.image_std?.[2] ?? 1), + ] as [number, number, number]; + + return { + inputWidth, + inputHeight, + rescaleFactor, + mean, + std, + }; + })(); + } + return preprocessorPromise; +} + +function preprocessResized( + image: CanvasImageSource, + preprocessor: ModelPreprocessor, + createCanvasFn: (width: number, height: number) => { getContext: (kind: '2d') => CanvasRenderingContext2D }, +): ort.Tensor { + const canvas = createCanvasFn(preprocessor.inputWidth, preprocessor.inputHeight); const ctx = canvas.getContext('2d'); ctx.fillStyle = '#ffffff'; - ctx.fillRect(0, 0, INPUT_SIZE, INPUT_SIZE); + ctx.fillRect(0, 0, preprocessor.inputWidth, preprocessor.inputHeight); ctx.imageSmoothingEnabled = true; ctx.imageSmoothingQuality = 'high'; - ctx.drawImage(image, padX, padY, drawWidth, drawHeight); + ctx.drawImage(image, 0, 0, preprocessor.inputWidth, preprocessor.inputHeight); - const imageData = ctx.getImageData(0, 0, INPUT_SIZE, INPUT_SIZE); - const data = new Uint8Array(1 * 3 * INPUT_SIZE * INPUT_SIZE); - for (let y = 0; y < INPUT_SIZE; y += 1) { - for (let x = 0; x < INPUT_SIZE; x += 1) { - const pixelIndex = (y * INPUT_SIZE + x) * 4; - const chwIndex = y * INPUT_SIZE + x; - data[0 * INPUT_SIZE * INPUT_SIZE + chwIndex] = imageData.data[pixelIndex]; - data[1 * INPUT_SIZE * INPUT_SIZE + chwIndex] = imageData.data[pixelIndex + 1]; - data[2 * INPUT_SIZE * INPUT_SIZE + chwIndex] = imageData.data[pixelIndex + 2]; + const imageData = ctx.getImageData(0, 0, preprocessor.inputWidth, preprocessor.inputHeight); + const chw = new Float32Array(1 * 3 * preprocessor.inputWidth * preprocessor.inputHeight); + const channelSize = preprocessor.inputWidth * preprocessor.inputHeight; + for (let y = 0; y < preprocessor.inputHeight; y += 1) { + for (let x = 0; x < preprocessor.inputWidth; x += 1) { + const pixelIndex = (y * preprocessor.inputWidth + x) * 4; + const idx = y * preprocessor.inputWidth + x; + const r = imageData.data[pixelIndex] * preprocessor.rescaleFactor; + const g = imageData.data[pixelIndex + 1] * preprocessor.rescaleFactor; + const b = imageData.data[pixelIndex + 2] * preprocessor.rescaleFactor; + + chw[idx] = (r - preprocessor.mean[0]) / Math.max(1e-8, preprocessor.std[0]); + chw[channelSize + idx] = (g - preprocessor.mean[1]) / Math.max(1e-8, preprocessor.std[1]); + chw[channelSize * 2 + idx] = (b - preprocessor.mean[2]) / Math.max(1e-8, preprocessor.std[2]); } } - - return { - tensor: new ort.Tensor('uint8', data, [1, 3, INPUT_SIZE, INPUT_SIZE]), - scale, - padX, - padY, - }; + return new ort.Tensor('float32', chw, [1, 3, preprocessor.inputHeight, preprocessor.inputWidth]); } function clampBox( @@ -169,6 +224,36 @@ function clampBox( return [x0, y0, x1, y1]; } +function softmaxMax(logits: Float32Array, offset: number, count: number): { index: number; score: number } { + let maxLogit = Number.NEGATIVE_INFINITY; + let maxIndex = 0; + for (let i = 0; i < count; i += 1) { + const value = logits[offset + i]; + if (value > maxLogit) { + maxLogit = value; + maxIndex = i; + } + } + + let sum = 0; + for (let i = 0; i < count; i += 1) { + sum += Math.exp(logits[offset + i] - maxLogit); + } + + const score = sum > 0 ? (1 / sum) : 0; + return { index: maxIndex, score }; +} + +function normalizeModelLabel(rawLabel: string): string { + const normalized = rawLabel.trim().toLowerCase().replace(/[\s-]+/g, '_'); + if (normalized.endsWith('_image')) { + const base = normalized.slice(0, -'_image'.length); + if (base === 'header' || base === 'footer') return normalized; + } + // Some exports suffix duplicate classes (e.g. header_1, footer_1, text_1). + return normalized.replace(/_\d+$/g, ''); +} + export async function runLayoutModel(input: RunLayoutInput): Promise { const { pageWidth, pageHeight, textItems, pagePng } = input; if (!textItems.length) return []; @@ -177,44 +262,49 @@ export async function runLayoutModel(input: RunLayoutInput): Promise([ + 'text', + 'content', + 'reference_content', + 'aside_text', + 'abstract', + 'algorithm', + 'reference', +]); + function stripTrailingClosers(text: string): string { return text.trim().replace(/[\"'”’\]\)]+$/g, ''); } @@ -9,7 +19,7 @@ function isSentenceTerminal(text: string): boolean { } function canStitch(a: ParsedPdfBlock, b: ParsedPdfBlock): boolean { - if (!['paragraph', 'list-item'].includes(a.kind)) return false; + if (!STITCHABLE_KINDS.has(a.kind)) return false; if (a.kind !== b.kind) return false; if (isSentenceTerminal(a.text)) return false; const next = b.text.trim(); @@ -18,13 +28,16 @@ function canStitch(a: ParsedPdfBlock, b: ParsedPdfBlock): boolean { return true; } -const HARD_BOUNDARY_KINDS = new Set(['section-header', 'title']); +const HARD_BOUNDARY_KINDS = new Set([ + 'paragraph_title', + 'doc_title', +]); function findTailCandidateIndex(blocks: ParsedPdfBlock[]): number { for (let i = blocks.length - 1; i >= 0; i -= 1) { const block = blocks[i]; if (!block || !block.text.trim()) continue; - if (block.kind === 'paragraph' || block.kind === 'list-item') return i; + if (STITCHABLE_KINDS.has(block.kind)) return i; } return -1; } @@ -33,7 +46,7 @@ function findHeadCandidateIndex(blocks: ParsedPdfBlock[]): number { for (let i = 0; i < blocks.length; i += 1) { const block = blocks[i]; if (!block || !block.text.trim()) continue; - if (block.kind === 'paragraph' || block.kind === 'list-item') return i; + if (STITCHABLE_KINDS.has(block.kind)) return i; } return -1; } diff --git a/src/lib/shared/document-settings.ts b/src/lib/shared/document-settings.ts index d22ef6c..537038a 100644 --- a/src/lib/shared/document-settings.ts +++ b/src/lib/shared/document-settings.ts @@ -2,25 +2,11 @@ import { DEFAULT_DOCUMENT_SETTINGS, type DocumentSettings, } from '@/types/document-settings'; -import type { ParsedPdfBlockKind } from '@/types/parsed-pdf'; - -const PDF_BLOCK_KINDS: ParsedPdfBlockKind[] = [ - 'title', - 'section-header', - 'paragraph', - 'list-item', - 'caption', - 'table', - 'picture', - 'page-header', - 'page-footer', - 'footnote', - 'formula', -]; +import { PARSED_PDF_BLOCK_KINDS, type ParsedPdfBlockKind } from '@/types/parsed-pdf'; function normalizeSkipKinds(value: unknown): ParsedPdfBlockKind[] { if (!Array.isArray(value)) return [...(DEFAULT_DOCUMENT_SETTINGS.pdf?.skipBlockKinds ?? [])]; - const allow = new Set(PDF_BLOCK_KINDS); + const allow = new Set(PARSED_PDF_BLOCK_KINDS); const out: ParsedPdfBlockKind[] = []; for (const entry of value) { if (typeof entry !== 'string') continue; diff --git a/src/types/document-settings.ts b/src/types/document-settings.ts index bb591ad..0262ed5 100644 --- a/src/types/document-settings.ts +++ b/src/types/document-settings.ts @@ -12,7 +12,7 @@ export interface DocumentSettings { export const DEFAULT_DOCUMENT_SETTINGS: DocumentSettings = { schemaVersion: 1, pdf: { - skipBlockKinds: ['page-header', 'page-footer', 'footnote'], + skipBlockKinds: ['header', 'footer', 'footnote', 'vision_footnote'], chaptersFromSections: true, }, }; diff --git a/src/types/parsed-pdf.ts b/src/types/parsed-pdf.ts index f64f126..c1c1daa 100644 --- a/src/types/parsed-pdf.ts +++ b/src/types/parsed-pdf.ts @@ -1,15 +1,49 @@ export type ParsedPdfBlockKind = - | 'title' - | 'section-header' - | 'paragraph' - | 'list-item' - | 'caption' - | 'table' - | 'picture' - | 'page-header' - | 'page-footer' + | 'abstract' + | 'algorithm' + | 'aside_text' + | 'chart' + | 'content' + | 'formula' + | 'doc_title' + | 'figure_title' + | 'footer' | 'footnote' - | 'formula'; + | 'formula_number' + | 'header' + | 'image' + | 'number' + | 'paragraph_title' + | 'reference' + | 'reference_content' + | 'seal' + | 'table' + | 'text' + | 'vision_footnote'; + +export const PARSED_PDF_BLOCK_KINDS: ParsedPdfBlockKind[] = [ + 'abstract', + 'algorithm', + 'aside_text', + 'chart', + 'content', + 'formula', + 'doc_title', + 'figure_title', + 'footer', + 'footnote', + 'formula_number', + 'header', + 'image', + 'number', + 'paragraph_title', + 'reference', + 'reference_content', + 'seal', + 'table', + 'text', + 'vision_footnote', +]; export interface ParsedPdfBlockFragment { page: number; diff --git a/tests/unit/pdf-audiobook-adapter.spec.ts b/tests/unit/pdf-audiobook-adapter.spec.ts index d430401..cf357cf 100644 --- a/tests/unit/pdf-audiobook-adapter.spec.ts +++ b/tests/unit/pdf-audiobook-adapter.spec.ts @@ -4,7 +4,7 @@ import type { ParsedPdfDocument } from '../../src/types/parsed-pdf'; import type { DocumentSettings } from '../../src/types/document-settings'; test.describe('pdf audiobook adapter', () => { - test('builds chapters from section headers and filters skipped kinds', async () => { + test('builds chapters from paragraph titles and filters skipped kinds', async () => { const parsed: ParsedPdfDocument = { schemaVersion: 1, documentId: 'doc-1', @@ -18,31 +18,31 @@ test.describe('pdf audiobook adapter', () => { blocks: [ { id: 'b1', - kind: 'section-header', + kind: 'paragraph_title', text: 'Intro', fragments: [{ page: 1, bbox: [0, 80, 100, 90], text: 'Intro', readingOrder: 0 }], }, { id: 'b2', - kind: 'paragraph', + kind: 'text', text: 'Welcome text.', fragments: [{ page: 1, bbox: [0, 60, 100, 79], text: 'Welcome text.', readingOrder: 1 }], }, { id: 'b3', - kind: 'page-header', + kind: 'header', text: 'Header line', fragments: [{ page: 1, bbox: [0, 95, 100, 100], text: 'Header line', readingOrder: 2 }], }, { id: 'b4', - kind: 'section-header', + kind: 'paragraph_title', text: 'Second', fragments: [{ page: 1, bbox: [0, 40, 100, 50], text: 'Second', readingOrder: 3 }], }, { id: 'b5', - kind: 'paragraph', + kind: 'text', text: 'More body.', fragments: [{ page: 1, bbox: [0, 20, 100, 39], text: 'More body.', readingOrder: 4 }], }, @@ -54,7 +54,7 @@ test.describe('pdf audiobook adapter', () => { const settings: DocumentSettings = { schemaVersion: 1, pdf: { - skipBlockKinds: ['page-header'], + skipBlockKinds: ['header'], chaptersFromSections: true, }, }; diff --git a/tests/unit/pdf-build-page-text.spec.ts b/tests/unit/pdf-build-page-text.spec.ts index e5b8ac2..247507c 100644 --- a/tests/unit/pdf-build-page-text.spec.ts +++ b/tests/unit/pdf-build-page-text.spec.ts @@ -11,26 +11,26 @@ test.describe('buildPageTextFromBlocks', () => { blocks: [ { id: 'b2', - kind: 'page-header', + kind: 'header', text: 'Copyright Header', fragments: [{ page: 1, bbox: [0, 90, 100, 100], text: 'Copyright Header', readingOrder: 0 }], }, { id: 'b1', - kind: 'paragraph', + kind: 'text', text: 'Body text', fragments: [{ page: 1, bbox: [0, 20, 100, 80], text: 'Body text', readingOrder: 1 }], }, { id: 'b3', - kind: 'caption', + kind: 'figure_title', text: 'Figure caption', fragments: [{ page: 1, bbox: [0, 5, 100, 19], text: 'Figure caption', readingOrder: 2 }], }, ], }; - expect(buildPageTextFromBlocks(page, ['page-header'])).toBe('Body text Figure caption'); - expect(buildPageTextFromBlocks(page, ['page-header', 'caption'])).toBe('Body text'); + expect(buildPageTextFromBlocks(page, ['header'])).toBe('Body text Figure caption'); + expect(buildPageTextFromBlocks(page, ['header', 'figure_title'])).toBe('Body text'); }); }); diff --git a/tests/unit/pdf-merge-text-with-regions.spec.ts b/tests/unit/pdf-merge-text-with-regions.spec.ts index 09be6dc..f0f8bb7 100644 --- a/tests/unit/pdf-merge-text-with-regions.spec.ts +++ b/tests/unit/pdf-merge-text-with-regions.spec.ts @@ -4,8 +4,8 @@ import { mergeTextWithRegions } from '../../src/lib/server/pdf-layout/mergeTextW test.describe('mergeTextWithRegions', () => { test('assigns text items to containing regions by centroid and joins in reading order', () => { const regions = [ - { bbox: [0, 0, 100, 50] as [number, number, number, number], label: 'paragraph' as const }, - { bbox: [0, 50, 100, 100] as [number, number, number, number], label: 'caption' as const }, + { bbox: [0, 0, 100, 50] as [number, number, number, number], label: 'text' as const }, + { bbox: [0, 50, 100, 100] as [number, number, number, number], label: 'figure_title' as const }, ]; const textItems = [ diff --git a/tests/unit/pdf-stitch-cross-page-blocks.spec.ts b/tests/unit/pdf-stitch-cross-page-blocks.spec.ts index 2a0ff5d..cdc3ac2 100644 --- a/tests/unit/pdf-stitch-cross-page-blocks.spec.ts +++ b/tests/unit/pdf-stitch-cross-page-blocks.spec.ts @@ -39,12 +39,12 @@ test.describe('stitchCrossPageBlocks', () => { test('stitches paragraph continuation across footer/header noise', () => { const doc = makeDoc( [ - makeBlock('b1', 'paragraph', 'This sentence continues', 1, 0), - makeBlock('b2', 'page-footer', 'Footer text', 1, 1), + makeBlock('b1', 'text', 'This sentence continues', 1, 0), + makeBlock('b2', 'footer', 'Footer text', 1, 1), ], [ - makeBlock('b3', 'page-header', 'Header text', 2, 0), - makeBlock('b4', 'paragraph', 'into the next page.', 2, 1), + makeBlock('b3', 'header', 'Header text', 2, 0), + makeBlock('b4', 'text', 'into the next page.', 2, 1), ], ); @@ -57,14 +57,14 @@ test.describe('stitchCrossPageBlocks', () => { expect(page2?.blocks.map((b) => b.id)).toEqual(['b3']); }); - test('does not stitch across section-header boundary', () => { + test('does not stitch across paragraph-title boundary', () => { const doc = makeDoc( [ - makeBlock('b1', 'paragraph', 'This sentence continues', 1, 0), + makeBlock('b1', 'text', 'This sentence continues', 1, 0), ], [ - makeBlock('b2', 'section-header', '2 New Section', 2, 0), - makeBlock('b3', 'paragraph', 'into the next page.', 2, 1), + makeBlock('b2', 'paragraph_title', '2 New Section', 2, 0), + makeBlock('b3', 'text', 'into the next page.', 2, 1), ], ); @@ -76,10 +76,10 @@ test.describe('stitchCrossPageBlocks', () => { test('does not stitch when tail has sentence terminal', () => { const doc = makeDoc( [ - makeBlock('b1', 'paragraph', 'This sentence is complete.', 1, 0), + makeBlock('b1', 'text', 'This sentence is complete.', 1, 0), ], [ - makeBlock('b2', 'paragraph', 'next sentence starts here', 2, 0), + makeBlock('b2', 'text', 'next sentence starts here', 2, 0), ], ); @@ -88,4 +88,3 @@ test.describe('stitchCrossPageBlocks', () => { expect(stitched.pages[1]?.blocks.map((b) => b.id)).toEqual(['b2']); }); }); -