feat(compute): add external compute worker backend and integration

Introduce support for external compute worker mode (`COMPUTE_MODE=worker`) using a new `WorkerComputeBackend`. This enables offloading heavy ONNX Whisper alignment and PDF layout parsing to a standalone worker service (Redis + BullMQ), improving scalability and compatibility with serverless/limited environments.

- Add `@openreader/compute-core` as a shared package for ONNX inference and PDF parsing logic.
- Implement `WorkerComputeBackend` and worker contract/types for remote job execution.
- Update compute backend selection logic and remove previous worker mode guards.
- Extend `WhisperAlignInput` and `PdfLayoutInput` types to support object keys for remote data access.
- Refactor local compute backend to use `@openreader/compute-core` and support both buffer and object key inputs.
- Update job runner, TTS segment alignment, and PDF layout parsing flows to use new compute backend APIs.
- Add scripts, Docker workflow, and documentation for deploying and running the compute worker.
- Update environment variable docs and examples for worker mode, including storage requirements and configuration.
- Document published images and stack changes to reflect the new compute worker architecture.

BREAKING CHANGE: `COMPUTE_MODE=worker` now requires an external compute worker service and S3-compatible object storage. Embedded SeaweedFS (`weed mini`) is not supported in worker mode. See the new documentation for deployment and configuration details.
This commit is contained in:
Richard R 2026-05-19 15:21:25 -06:00
parent 3a21f2a5f5
commit f1aa1c3e3b
53 changed files with 4880 additions and 96 deletions

View file

@ -80,10 +80,13 @@ IMPORT_LIBRARY_DIRS=
# Heavy compute backend mode for ONNX whisper alignment + PDF layout parsing.
# local = run compute in-process (default)
# none = disable both capabilities (good for preview/serverless)
# worker = reserved for future external worker mode (not implemented in v1)
# worker = external durable worker mode (requires Redis-backed compute-worker service)
COMPUTE_MODE=local
# COMPUTE_WORKER_URL=
# COMPUTE_WORKER_TOKEN=
# Required when COMPUTE_MODE=worker
# COMPUTE_WORKER_URL=http://localhost:8081
# COMPUTE_WORKER_TOKEN=local-compute-token
# Worker mode requires worker-reachable shared object storage.
# Non-exposed embedded weed mini is not supported in worker mode.
# Optional Whisper ONNX base URL override (must contain all expected files)
# WHISPER_MODEL_BASE_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main

View file

@ -22,30 +22,17 @@ jobs:
prepare:
runs-on: ubuntu-24.04
outputs:
image_name: ${{ steps.image-name.outputs.image_name }}
legacy_image_name: ${{ steps.image-name.outputs.legacy_image_name }}
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
web_image_name: ${{ steps.image-name.outputs.web_image_name }}
web_legacy_image_name: ${{ steps.image-name.outputs.web_legacy_image_name }}
compute_worker_image_name: ${{ steps.image-name.outputs.compute_worker_image_name }}
steps:
- name: Compute Docker image names
id: image-name
run: |
owner="${GITHUB_REPOSITORY_OWNER,,}"
echo "image_name=${owner}/openreader" >> "$GITHUB_OUTPUT"
echo "legacy_image_name=${owner}/openreader-webui" >> "$GITHUB_OUTPUT"
- name: Extract metadata (tags, labels) for Docker
id: meta
uses: docker/metadata-action@v6
with:
images: |
${{ env.REGISTRY }}/${{ steps.image-name.outputs.image_name }}
${{ env.REGISTRY }}/${{ steps.image-name.outputs.legacy_image_name }}
tags: |
type=raw,value=latest,enable=${{ (github.event_name == 'push' && !contains(github.ref, '-pre')) || (github.event_name == 'workflow_dispatch' && inputs.use_latest_tag == true) }},priority=1000
type=semver,pattern={{version}}
type=ref,event=tag
type=ref,event=branch,enable=${{ github.event_name == 'workflow_dispatch' && inputs.use_latest_tag != true }}
echo "web_image_name=${owner}/openreader" >> "$GITHUB_OUTPUT"
echo "web_legacy_image_name=${owner}/openreader-webui" >> "$GITHUB_OUTPUT"
echo "compute_worker_image_name=${owner}/openreader-compute-worker" >> "$GITHUB_OUTPUT"
build:
needs: prepare
@ -58,12 +45,30 @@ jobs:
fail-fast: false
matrix:
include:
- arch: amd64
- image_target: web
arch: amd64
platform: linux/amd64
runner: ubuntu-24.04
- arch: arm64
context: .
dockerfile: ./Dockerfile
- image_target: web
arch: arm64
platform: linux/arm64
runner: ubuntu-24.04-arm
context: .
dockerfile: ./Dockerfile
- image_target: compute-worker
arch: amd64
platform: linux/amd64
runner: ubuntu-24.04
context: .
dockerfile: ./compute/worker/Dockerfile
- image_target: compute-worker
arch: arm64
platform: linux/arm64
runner: ubuntu-24.04-arm
context: .
dockerfile: ./compute/worker/Dockerfile
steps:
- name: Checkout repository
@ -82,46 +87,63 @@ jobs:
username: ${{ github.actor }}
password: ${{ secrets.GITHUB_TOKEN }}
- name: Build and push Docker image (${{ matrix.arch }})
- name: Select image name
id: image
run: |
if [ "${{ matrix.image_target }}" = "web" ]; then
echo "image_name=${{ needs.prepare.outputs.web_image_name }}" >> "$GITHUB_OUTPUT"
else
echo "image_name=${{ needs.prepare.outputs.compute_worker_image_name }}" >> "$GITHUB_OUTPUT"
fi
- name: Build and push Docker image (${{ matrix.image_target }} / ${{ matrix.arch }})
id: build
uses: docker/build-push-action@v7
with:
context: .
context: ${{ matrix.context }}
file: ${{ matrix.dockerfile }}
platforms: ${{ matrix.platform }}
labels: ${{ needs.prepare.outputs.labels }}
cache-from: type=gha,scope=${{ matrix.arch }}
cache-to: type=gha,mode=max,scope=${{ matrix.arch }}
cache-from: type=gha,scope=${{ matrix.image_target }}-${{ matrix.arch }}
cache-to: type=gha,mode=max,scope=${{ matrix.image_target }}-${{ matrix.arch }}
provenance: false
outputs: type=image,name=${{ env.REGISTRY }}/${{ needs.prepare.outputs.image_name }},push-by-digest=true,name-canonical=true,push=true
outputs: type=image,name=${{ env.REGISTRY }}/${{ steps.image.outputs.image_name }},push-by-digest=true,name-canonical=true,push=true
- name: Export digest
run: |
mkdir -p /tmp/digests
mkdir -p /tmp/digests/${{ matrix.image_target }}
digest="${{ steps.build.outputs.digest }}"
touch "/tmp/digests/${digest#sha256:}"
touch "/tmp/digests/${{ matrix.image_target }}/${digest#sha256:}"
- name: Upload digest
uses: actions/upload-artifact@v7
with:
name: digests-${{ matrix.arch }}
path: /tmp/digests
name: digests-${{ matrix.image_target }}-${{ matrix.arch }}
path: /tmp/digests/${{ matrix.image_target }}
if-no-files-found: error
retention-days: 1
merge:
needs: [prepare, build]
runs-on: ubuntu-24.04
runs-on: ${{ matrix.runner }}
permissions:
actions: read
contents: read
packages: write
strategy:
fail-fast: false
matrix:
include:
- image_target: web
runner: ubuntu-24.04
- image_target: compute-worker
runner: ubuntu-24.04
steps:
- name: Download digests
uses: actions/download-artifact@v8
with:
path: /tmp/digests
pattern: digests-*
path: /tmp/digests/${{ matrix.image_target }}
pattern: digests-${{ matrix.image_target }}-*
merge-multiple: true
- name: Set up Docker Buildx
@ -134,20 +156,46 @@ jobs:
username: ${{ github.actor }}
password: ${{ secrets.GITHUB_TOKEN }}
- name: Select image names
id: image
run: |
if [ "${{ matrix.image_target }}" = "web" ]; then
echo "image_name=${{ needs.prepare.outputs.web_image_name }}" >> "$GITHUB_OUTPUT"
echo "metadata_images<<EOF" >> "$GITHUB_OUTPUT"
echo "${{ env.REGISTRY }}/${{ needs.prepare.outputs.web_image_name }}" >> "$GITHUB_OUTPUT"
echo "${{ env.REGISTRY }}/${{ needs.prepare.outputs.web_legacy_image_name }}" >> "$GITHUB_OUTPUT"
echo "EOF" >> "$GITHUB_OUTPUT"
else
echo "image_name=${{ needs.prepare.outputs.compute_worker_image_name }}" >> "$GITHUB_OUTPUT"
echo "metadata_images=${{ env.REGISTRY }}/${{ needs.prepare.outputs.compute_worker_image_name }}" >> "$GITHUB_OUTPUT"
fi
- name: Extract metadata (tags, labels) for Docker
id: meta
uses: docker/metadata-action@v6
with:
images: ${{ steps.image.outputs.metadata_images }}
tags: |
type=raw,value=latest,enable=${{ (github.event_name == 'push' && !contains(github.ref, '-pre')) || (github.event_name == 'workflow_dispatch' && inputs.use_latest_tag == true) }},priority=1000
type=semver,pattern={{version}}
type=ref,event=tag
type=ref,event=branch,enable=${{ github.event_name == 'workflow_dispatch' && inputs.use_latest_tag != true }}
- name: Create manifest list and push
working-directory: /tmp/digests
working-directory: /tmp/digests/${{ matrix.image_target }}
run: |
docker buildx imagetools create \
$(echo "$TAGS" | xargs -I {} echo -t {}) \
$(printf '${{ env.REGISTRY }}/${{ needs.prepare.outputs.image_name }}@sha256:%s ' *)
$(printf '${{ env.REGISTRY }}/${{ steps.image.outputs.image_name }}@sha256:%s ' *)
env:
TAGS: ${{ needs.prepare.outputs.tags }}
TAGS: ${{ steps.meta.outputs.tags }}
- name: Output build information
run: |
echo "✅ Docker images built and pushed successfully!"
echo "✅ Docker image built and pushed successfully!"
echo "🐋 Target: ${{ matrix.image_target }}"
echo "🐋 Images:"
echo '${{ needs.prepare.outputs.tags }}' | sed 's/^/ - /'
echo '${{ steps.meta.outputs.tags }}' | sed 's/^/ - /'
echo "📝 Event: ${{ github.event_name }}"
if [ "${{ github.event_name }}" == "workflow_dispatch" ]; then
echo "📝 Triggered by manual workflow dispatch on branch: ${{ github.ref_name }}"

View file

@ -20,7 +20,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 built-in ONNX Whisper alignment in local compute mode (`COMPUTE_MODE=local`).
- ⏱️ **Word-by-word highlighting** via ONNX Whisper alignment in local mode (`COMPUTE_MODE=local`) or external worker mode (`COMPUTE_MODE=worker`).
- 🧱 **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.
@ -33,6 +33,7 @@ OpenReader is an open source, self-host-friendly text-to-speech document reader
| --- | --- |
| Run with Docker | [Docker Quick Start](https://docs.openreader.richardr.dev/docker-quick-start) |
| Deploy on Vercel | [Vercel Deployment](https://docs.openreader.richardr.dev/deploy/vercel-deployment) |
| Deploy external compute worker | [Compute Worker (Redis + BullMQ)](https://docs.openreader.richardr.dev/deploy/compute-worker) |
| Develop locally | [Local Development](https://docs.openreader.richardr.dev/deploy/local-development) |
| Configure auth | [Auth](https://docs.openreader.richardr.dev/configure/auth) |
| Configure SQL database | [Database and Migrations](https://docs.openreader.richardr.dev/configure/database) |

17
compute/core/package.json Normal file
View file

@ -0,0 +1,17 @@
{
"name": "@openreader/compute-core",
"version": "0.0.0",
"private": true,
"type": "module",
"dependencies": {
"@huggingface/tokenizers": "^0.1.3",
"@napi-rs/canvas": "^0.1.100",
"ffmpeg-static": "^5.3.0",
"jszip": "^3.10.1",
"onnxruntime-node": "^1.26.0",
"pdfjs-dist": "4.8.69"
},
"exports": {
".": "./src/index.ts"
}
}

87
compute/core/src/index.ts Normal file
View file

@ -0,0 +1,87 @@
import type { TTSSentenceAlignment } from './types/tts';
import type { ParsedPdfDocument } from './types/parsed-pdf';
import { ensureWhisperModel } from './whisper/ensureModel';
import { alignAudioWithText } from './whisper/alignment';
import { ensureModel as ensurePdfLayoutModel } from './pdf-layout/ensureModel';
import { parsePdf } from './pdf-layout/parsePdf';
export type {
TTSAudioBuffer,
TTSAudioBytes,
TTSSentenceAlignment,
TTSSentenceWord,
} from './types/tts';
export type {
ParsedPdfBlockKind,
ParsedPdfBlockFragment,
ParsedPdfBlock,
ParsedPdfPage,
ParsedPdfDocument,
} from './types/parsed-pdf';
export const ALIGN_QUEUE_NAME = 'whisper-align';
export const PDF_LAYOUT_QUEUE_NAME = 'pdf-layout';
export interface WhisperAlignJobRequest {
text: string;
lang?: string;
cacheKey?: string;
audioObjectKey: string;
}
export interface WhisperAlignJobResult {
alignments: TTSSentenceAlignment[];
}
export interface PdfLayoutJobRequest {
documentId: string;
namespace: string | null;
documentObjectKey: string;
}
export interface PdfLayoutJobResult {
parsed: ParsedPdfDocument;
}
export type WorkerJobState = 'queued' | 'running' | 'succeeded' | 'failed';
export interface WorkerJobErrorShape {
message: string;
code?: string;
}
export interface WorkerJobStatusResponse<Result> {
status: WorkerJobState;
result?: Result;
error?: WorkerJobErrorShape;
}
export async function ensureComputeModels(): Promise<void> {
await Promise.all([ensureWhisperModel(), ensurePdfLayoutModel()]);
}
export async function runWhisperAlignmentFromAudioBuffer(input: {
audioBuffer: ArrayBuffer;
text: string;
cacheKey?: string;
lang?: string;
}): Promise<WhisperAlignJobResult> {
const alignments = await alignAudioWithText(
input.audioBuffer,
input.text,
input.cacheKey,
{ lang: input.lang },
);
return { alignments };
}
export async function runPdfLayoutFromPdfBuffer(input: {
documentId: string;
pdfBytes: ArrayBuffer;
}): Promise<PdfLayoutJobResult> {
const parsed = await parsePdf({
documentId: input.documentId,
pdfBytes: input.pdfBytes,
});
return { parsed };
}

View file

@ -0,0 +1,132 @@
import path from 'path';
import { fileURLToPath } from 'url';
import { createHash } from 'crypto';
import { access, mkdir, rename, writeFile, readFile, unlink, copyFile } from 'fs/promises';
import { DOCSTORE_DIR } from '../runtime/docstore';
import manifest from './model/manifest.json';
const DEFAULT_MODEL_BASE_URL = 'https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main';
const PDF_LAYOUT_MODEL_BASE_URL_ENV = 'PDF_LAYOUT_MODEL_BASE_URL';
const MODEL_DIR = path.join(DOCSTORE_DIR, 'model');
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 MODULE_DIR = path.dirname(fileURLToPath(import.meta.url));
const STATIC_LICENSE_PATH = path.join(MODULE_DIR, 'model', 'LICENSE.txt');
let inflight: Promise<string> | null = null;
async function sha256Hex(filePath: string): Promise<string> {
const bytes = await readFile(filePath);
return createHash('sha256').update(bytes).digest('hex');
}
async function downloadToFile(url: string, outPath: string): Promise<void> {
const res = await fetch(url);
if (!res.ok) {
throw new Error(`Download failed for ${url}: ${res.status} ${res.statusText}`);
}
const bytes = new Uint8Array(await res.arrayBuffer());
await writeFile(outPath, bytes);
}
function joinModelUrl(baseUrl: string, relativePath: string): string {
return `${baseUrl.replace(/\/+$/, '')}/${relativePath}`;
}
function manifestEntry(filePath: string): { sha256: string; size: number } | null {
const found = manifest.files.find((entry) => entry.path === filePath);
if (!found || !found.sha256) return null;
return {
sha256: found.sha256.toLowerCase(),
size: Number(found.size),
};
}
async function verifyFile(pathToFile: string, manifestPath: string): Promise<boolean> {
const expected = manifestEntry(manifestPath);
if (!expected) return true;
const bytes = await readFile(pathToFile);
if (Number.isFinite(expected.size) && expected.size > 0 && bytes.byteLength !== expected.size) {
return false;
}
const actual = await sha256Hex(pathToFile);
return actual === expected.sha256;
}
async function ensureLicense(): Promise<void> {
await copyFile(STATIC_LICENSE_PATH, LICENSE_PATH);
if (!(await verifyFile(LICENSE_PATH, 'LICENSE.txt'))) {
throw new Error('PDF layout model license checksum verification failed');
}
}
async function ensureModelInternal(): Promise<string> {
try {
await access(MODEL_PATH);
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;
}
} catch {
// continue
}
await mkdir(MODEL_DIR, { recursive: true });
const modelTmpPath = `${MODEL_PATH}.tmp`;
const modelDataTmpPath = `${MODEL_DATA_PATH}.tmp`;
const configTmpPath = `${MODEL_CONFIG_PATH}.tmp`;
const preprocessorTmpPath = `${MODEL_PREPROCESSOR_PATH}.tmp`;
const modelBaseUrl = process.env[PDF_LAYOUT_MODEL_BASE_URL_ENV]?.trim()
|| DEFAULT_MODEL_BASE_URL;
const modelUrl = joinModelUrl(modelBaseUrl, 'PP-DocLayoutV3.onnx');
const modelDataUrl = joinModelUrl(modelBaseUrl, 'PP-DocLayoutV3.onnx.data');
const configUrl = joinModelUrl(modelBaseUrl, 'config.json');
const preprocessorUrl = joinModelUrl(modelBaseUrl, 'preprocessor_config.json');
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('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(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;
}
export async function ensureModel(): Promise<string> {
if (inflight) return inflight;
inflight = ensureModelInternal().finally(() => {
inflight = null;
});
return inflight;
}

View file

@ -0,0 +1,153 @@
import type { LayoutRegion, PdfTextItem } from './types';
const NON_TEXT_REGION_LABELS = new Set<LayoutRegion['label']>(['chart', 'image', 'table', 'seal']);
const TEXT_ASSIGNABLE_LABELS = new Set<LayoutRegion['label']>([
'abstract',
'algorithm',
'aside_text',
'content',
'doc_title',
'figure_title',
'footer',
'footnote',
'formula_number',
'header',
'number',
'paragraph_title',
'reference',
'reference_content',
'text',
'vision_footnote',
'formula',
]);
function centroid(item: PdfTextItem): { x: number; y: number } {
return {
x: item.x + item.width / 2,
y: item.y + item.height / 2,
};
}
function contains(region: LayoutRegion, item: PdfTextItem): boolean {
const c = centroid(item);
return c.x >= region.bbox[0] && c.x <= region.bbox[2] && c.y >= region.bbox[1] && c.y <= region.bbox[3];
}
function sortReadingOrder(items: PdfTextItem[]): PdfTextItem[] {
const tolerance = 2;
return [...items].sort((a, b) => {
if (Math.abs(a.y - b.y) <= tolerance) return a.x - b.x;
return a.y - b.y;
});
}
function joinText(items: PdfTextItem[]): string {
let out = '';
let prev: PdfTextItem | null = null;
for (const item of items) {
if (!prev) {
out += item.text;
prev = item;
continue;
}
const prevEndX = prev.x + prev.width;
const gap = item.x - prevEndX;
const lineJump = item.y - prev.y;
const lineBreak = lineJump > Math.max(2, Math.min(prev.height, item.height) * 0.6);
const avgCharWidth = item.width / Math.max(1, item.text.length);
const needsSpace = lineBreak || gap > Math.max(avgCharWidth * 0.3, 2);
out += needsSpace ? ` ${item.text}` : item.text;
prev = item;
}
return out.replace(/\s+/g, ' ').trim();
}
function regionArea(region: LayoutRegion): number {
return Math.max(1, (region.bbox[2] - region.bbox[0]) * (region.bbox[3] - region.bbox[1]));
}
function regionScore(region: LayoutRegion): number {
return Number.isFinite(region.confidence) ? Number(region.confidence) : 0;
}
export interface RegionTextBlock {
region: LayoutRegion;
text: string;
items: PdfTextItem[];
sourceOrder: number;
}
export function mergeTextWithRegions(regions: LayoutRegion[], textItems: PdfTextItem[]): RegionTextBlock[] {
const sourceIndex = new Map<PdfTextItem, number>();
for (let i = 0; i < textItems.length; i += 1) {
sourceIndex.set(textItems[i]!, i);
}
const chunkSourceOrder = (items: PdfTextItem[]): number => {
let min = Number.POSITIVE_INFINITY;
for (const item of items) {
const index = sourceIndex.get(item);
if (typeof index === 'number' && index < min) min = index;
}
return Number.isFinite(min) ? min : Number.MAX_SAFE_INTEGER;
};
const assignableRegions = regions
.map((region, index) => ({ region, index }))
.filter(({ region }) => TEXT_ASSIGNABLE_LABELS.has(region.label));
const assignedByRegion = new Map<number, PdfTextItem[]>();
for (const item of textItems) {
const candidates = assignableRegions.filter(({ region }) => contains(region, item));
if (candidates.length === 0) continue;
candidates.sort((a, b) => {
const scoreDelta = regionScore(b.region) - regionScore(a.region);
if (Math.abs(scoreDelta) > 1e-6) return scoreDelta;
return regionArea(a.region) - regionArea(b.region);
});
const winner = candidates[0];
const list = assignedByRegion.get(winner.index) ?? [];
list.push(item);
assignedByRegion.set(winner.index, list);
}
const out: RegionTextBlock[] = [];
for (const [regionIndex, assignedItems] of assignedByRegion.entries()) {
const region = regions[regionIndex];
if (!region) continue;
if (assignedItems.length === 0) continue;
const ordered = sortReadingOrder(assignedItems);
const text = joinText(ordered);
if (!text) continue;
out.push({
region,
text,
items: ordered,
sourceOrder: chunkSourceOrder(ordered),
});
}
for (const region of regions) {
if (!NON_TEXT_REGION_LABELS.has(region.label)) continue;
out.push({
region,
text: '',
items: [],
sourceOrder: Number.MAX_SAFE_INTEGER,
});
}
out.sort((a, b) => {
if (a.sourceOrder !== b.sourceOrder) return a.sourceOrder - b.sourceOrder;
const ay = a.region.bbox[1];
const by = b.region.bbox[1];
if (Math.abs(ay - by) <= 2) return a.region.bbox[0] - b.region.bbox[0];
return ay - by;
});
return out;
}

View file

@ -0,0 +1,3 @@
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

View file

@ -0,0 +1,31 @@
{
"name": "pp-doclayoutv3",
"version": "Bei0001/PP-DocLayoutV3-ONNX@main",
"files": [
{
"path": "model.onnx",
"sha256": "c0721928ff08741bb208ebed539c77170db5234a68cb7e546e6cc9bc172a695b",
"size": 5088167
},
{
"path": "model.onnx.data",
"sha256": "34df3e4b79d7bbbf82abce1b4f3cde3d540fa57ad42ec8905c352b97c408d437",
"size": 136774480
},
{
"path": "config.json",
"sha256": "3cf834b91d23a756b1519bce4db42c09e852f3e35c35092dd5a3e253a50c071a",
"size": 2460
},
{
"path": "preprocessor_config.json",
"sha256": "519fe0187a43a1ca429e3ad8317bab8700f0d5e8fb3a6e3a0a413ffac078ba42",
"size": 575
},
{
"path": "LICENSE.txt",
"sha256": "578a6ba6f86b0692a8f719843f575a3eebf4705768ac5c37d149f441208f601f",
"size": 195
}
]
}

View file

@ -0,0 +1,165 @@
import path from 'path';
import type { TextItem } from 'pdfjs-dist/types/src/display/api';
import type { ParsedPdfDocument, ParsedPdfPage } from '../types/parsed-pdf';
import type { PdfTextItem } from './types';
import { ensureModel } from './ensureModel';
import { runLayoutModel } from './runLayoutModel';
import { mergeTextWithRegions } from './mergeTextWithRegions';
import { stitchCrossPageBlocks } from './stitchCrossPageBlocks';
import { renderPage } from './renderPage';
interface ParsePdfInput {
documentId: string;
pdfBytes: ArrayBuffer;
}
const LAYOUT_RENDER_SCALE = 1.5;
export function normalizeTextItemsForLayout(items: TextItem[], pageHeight: number): PdfTextItem[] {
return items
.filter((item) => {
if (!(typeof item.str === 'string' && item.str.trim().length > 0)) return false;
const transform = item.transform;
if (!Array.isArray(transform) || transform.length < 6) return false;
// Reject heavily skewed/rotated text runs (e.g. vertical margin labels
// such as arXiv metadata) so they do not get merged into body blocks.
const skewX = Number(transform[1] ?? 0);
const skewY = Number(transform[2] ?? 0);
if (Math.abs(skewX) > 0.5 || Math.abs(skewY) > 0.5) return false;
return true;
})
.map((item) => {
const x = Number(item.transform[4] ?? 0);
const width = Math.max(0, Number(item.width ?? 0));
const height = Math.max(1, Math.abs(Number(item.transform[3] ?? 1)));
const baselineY = Number(item.transform[5] ?? 0);
// pdf.js text transforms are in PDF user-space (origin bottom-left).
// Normalize into top-left page coordinates to match rendered image/model boxes.
const y = Math.max(0, pageHeight - baselineY - height);
return {
text: item.str,
x,
y,
width,
height,
};
});
}
export async function parsePdf(input: ParsePdfInput): Promise<ParsedPdfDocument> {
await ensureModel();
// Keep independent byte copies for text extraction and page rendering. pdf.js
// can detach buffers passed to getDocument().
const pdfBytesForText = new Uint8Array(input.pdfBytes).slice();
const pdfBytesForRender = new Uint8Array(input.pdfBytes).slice();
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-ignore
const pdfjs = await import('pdfjs-dist/legacy/build/pdf.mjs');
if (pdfjs.GlobalWorkerOptions) {
pdfjs.GlobalWorkerOptions.workerSrc = 'pdfjs-dist/legacy/build/pdf.worker.mjs';
pdfjs.GlobalWorkerOptions.workerPort = null;
}
const standardFontDir = path.join(process.cwd(), 'node_modules', 'pdfjs-dist', 'standard_fonts');
const standardFontDataUrl = `${standardFontDir.replace(/\/?$/, '/')}`;
const loadingTask = pdfjs.getDocument({
data: pdfBytesForText,
useWorkerFetch: false,
standardFontDataUrl,
isEvalSupported: false,
});
const pdf = await loadingTask.promise;
try {
const pages: ParsedPdfPage[] = [];
let nextBlockId = 1;
let sawText = false;
for (let pageNumber = 1; pageNumber <= pdf.numPages; pageNumber += 1) {
const page = await pdf.getPage(pageNumber);
const viewport = page.getViewport({ scale: 1.0 });
const textContent = await page.getTextContent();
const textItems = normalizeTextItemsForLayout(
textContent.items.filter((item): item is TextItem => 'str' in item && 'transform' in item),
viewport.height,
);
if (textItems.length > 0) sawText = true;
const rendered = await renderPage({
pdfBytes: pdfBytesForRender.buffer.slice(
pdfBytesForRender.byteOffset,
pdfBytesForRender.byteOffset + pdfBytesForRender.byteLength,
),
pageNumber,
scale: LAYOUT_RENDER_SCALE,
});
const scaleX = rendered.width / Math.max(1, viewport.width);
const scaleY = rendered.height / Math.max(1, viewport.height);
const layoutTextItems = textItems.map((item) => ({
...item,
x: item.x * scaleX,
y: item.y * scaleY,
width: item.width * scaleX,
height: item.height * scaleY,
}));
const regions = await runLayoutModel({
pageWidth: rendered.width,
pageHeight: rendered.height,
textItems: layoutTextItems,
pageImage: rendered.image,
});
const merged = mergeTextWithRegions(regions, layoutTextItems);
if (textItems.length > 0 && merged.length === 0) {
throw new Error(`layout-merge-empty: page=${pageNumber} regions=${regions.length}`);
}
const blocks = merged
.map((entry, readingOrder) => ({
id: `b${String(nextBlockId++).padStart(4, '0')}`,
kind: entry.region.label,
fragments: [{
page: pageNumber,
bbox: [
entry.region.bbox[0] / scaleX,
entry.region.bbox[1] / scaleY,
entry.region.bbox[2] / scaleX,
entry.region.bbox[3] / scaleY,
] as [number, number, number, number],
text: entry.text,
readingOrder,
...(typeof entry.region.confidence === 'number' ? { modelConfidence: entry.region.confidence } : {}),
}],
text: entry.text,
}));
pages.push({
pageNumber,
width: viewport.width,
height: viewport.height,
blocks,
});
}
if (!sawText) {
throw new Error('no-text-layer');
}
const doc: ParsedPdfDocument = {
schemaVersion: 1,
documentId: input.documentId,
parserVersion: 'pp-doclayoutv3-onnx@800+pdfjs@4.8.69',
parsedAt: Date.now(),
pages,
};
return stitchCrossPageBlocks(doc);
} finally {
await pdf.destroy().catch(() => undefined);
await loadingTask.destroy().catch(() => undefined);
}
}

View file

@ -0,0 +1,162 @@
import path from 'path';
import type { Canvas } from '@napi-rs/canvas';
type CanvasRuntime = {
DOMMatrixCtor: unknown;
Path2DCtor: unknown;
createCanvasFn: (width: number, height: number) => Canvas;
};
let canvasRuntimePromise: Promise<CanvasRuntime> | null = null;
async function loadCanvasRuntime(): Promise<CanvasRuntime> {
if (!canvasRuntimePromise) {
canvasRuntimePromise = (async () => {
const mod = await import('@napi-rs/canvas');
const namespace = mod as Record<string, unknown>;
const fallback = (namespace.default ?? {}) as Record<string, unknown>;
const createCanvasFn = (namespace.createCanvas ?? fallback.createCanvas) as
| ((width: number, height: number) => Canvas)
| undefined;
const DOMMatrixCtor = namespace.DOMMatrix ?? fallback.DOMMatrix;
const Path2DCtor = namespace.Path2D ?? fallback.Path2D;
if (typeof createCanvasFn !== 'function') {
throw new Error(
`Canvas runtime missing createCanvas export (keys=${Object.keys(namespace).join(',')}; defaultKeys=${Object.keys(fallback).join(',')})`,
);
}
if (!DOMMatrixCtor || !Path2DCtor) {
throw new Error(
`Canvas runtime missing DOMMatrix/Path2D exports (keys=${Object.keys(namespace).join(',')}; defaultKeys=${Object.keys(fallback).join(',')})`,
);
}
return {
DOMMatrixCtor,
Path2DCtor,
createCanvasFn,
};
})();
}
return canvasRuntimePromise;
}
function ensureNodeCanvasGlobals(runtime: CanvasRuntime): void {
const g = globalThis as Record<string, unknown>;
if (typeof g.DOMMatrix === 'undefined') g.DOMMatrix = runtime.DOMMatrixCtor;
if (typeof g.Path2D === 'undefined') g.Path2D = runtime.Path2DCtor;
}
interface RenderInput {
pdfBytes: ArrayBuffer;
pageNumber: number;
scale?: number;
targetWidth?: number;
format?: 'png' | 'jpeg';
jpegQuality?: number;
}
function createPdfjsCanvasFactory(runtime: CanvasRuntime) {
return class OpenReaderCanvasFactory {
create(width: number, height: number) {
const canvas = runtime.createCanvasFn(Math.max(1, Math.floor(width)), Math.max(1, Math.floor(height)));
return {
canvas,
context: canvas.getContext('2d') as unknown as CanvasRenderingContext2D,
};
}
reset(target: { canvas: Canvas; context: CanvasRenderingContext2D }, width: number, height: number): void {
target.canvas.width = Math.max(1, Math.floor(width));
target.canvas.height = Math.max(1, Math.floor(height));
}
destroy(target: { canvas: Canvas; context: CanvasRenderingContext2D }): void {
target.canvas.width = 0;
target.canvas.height = 0;
// @ts-expect-error pdf.js expects these nulled on destroy
target.canvas = null;
// @ts-expect-error pdf.js expects these nulled on destroy
target.context = null;
}
};
}
export async function renderPage({
pdfBytes,
pageNumber,
scale = 1.5,
targetWidth,
format = 'png',
jpegQuality = 82,
}: RenderInput): Promise<{
width: number;
height: number;
image: Buffer;
contentType: 'image/png' | 'image/jpeg';
}> {
// pdf.js may detach the provided ArrayBuffer. Work with an isolated copy so
// callers can safely reuse their original bytes across pages/calls.
const isolatedBytes = new Uint8Array(pdfBytes).slice();
const canvasRuntime = await loadCanvasRuntime();
ensureNodeCanvasGlobals(canvasRuntime);
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-ignore
const pdfjs = await import('pdfjs-dist/legacy/build/pdf.mjs');
if (pdfjs.GlobalWorkerOptions) {
pdfjs.GlobalWorkerOptions.workerSrc = 'pdfjs-dist/legacy/build/pdf.worker.mjs';
pdfjs.GlobalWorkerOptions.workerPort = null;
}
const standardFontDir = path.join(process.cwd(), 'node_modules', 'pdfjs-dist', 'standard_fonts');
const standardFontDataUrl = `${standardFontDir.replace(/\/?$/, '/')}`;
const loadingTask = pdfjs.getDocument({
data: isolatedBytes,
useWorkerFetch: false,
standardFontDataUrl,
isEvalSupported: false,
// Ensure pdf.js transport uses our canvas backend in Node/Next runtime.
CanvasFactory: createPdfjsCanvasFactory(canvasRuntime),
});
const pdf = await loadingTask.promise;
try {
const page = await pdf.getPage(pageNumber);
const baseViewport = page.getViewport({ scale: 1.0 });
const effectiveScale = typeof targetWidth === 'number' && Number.isFinite(targetWidth) && targetWidth > 0
? (Math.max(1, Math.round(targetWidth)) / Math.max(1, baseViewport.width))
: scale;
const viewport = page.getViewport({ scale: effectiveScale });
const width = Math.max(1, Math.floor(viewport.width));
const height = Math.max(1, Math.floor(viewport.height));
const canvas = canvasRuntime.createCanvasFn(width, height);
const ctx = canvas.getContext('2d');
ctx.fillStyle = '#ffffff';
ctx.fillRect(0, 0, width, height);
const renderTask = page.render({
canvasContext: ctx as unknown as CanvasRenderingContext2D,
viewport,
intent: 'display',
});
await renderTask.promise;
const contentType = format === 'jpeg' ? 'image/jpeg' : 'image/png';
const image = format === 'jpeg'
? canvas.toBuffer('image/jpeg', jpegQuality)
: canvas.toBuffer('image/png');
return {
width,
height,
image,
contentType,
};
} finally {
await pdf.destroy().catch(() => undefined);
await loadingTask.destroy().catch(() => undefined);
}
}

View file

@ -0,0 +1,331 @@
import * as ort from 'onnxruntime-node';
import { readFile } from 'fs/promises';
import type { LayoutRegion, PdfTextItem } from './types';
import { ensureModel, MODEL_CONFIG_PATH, MODEL_PREPROCESSOR_PATH } from './ensureModel';
interface RunLayoutInput {
pageWidth: number;
pageHeight: number;
textItems: PdfTextItem[];
pageImage: Buffer;
}
const DEFAULT_INPUT_SIZE = 800;
const MIN_SCORE = 0.5;
const CLASS_MIN_SCORE: Partial<Record<LayoutRegion['label'], number>> = {
header: 0.4,
footer: 0.4,
figure_title: 0.45,
footnote: 0.45,
vision_footnote: 0.45,
};
const LABEL_MAP: Record<string, LayoutRegion['label'] | null> = {
// 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_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: 'text',
vertical_text: 'text',
vision_footnote: 'vision_footnote',
};
const MIN_REGION_SIZE: Partial<Record<LayoutRegion['label'], { minWidth: number; minHeight: number }>> = {
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 },
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<ort.InferenceSession> | null = null;
let idToLabelPromise: Promise<Record<number, string>> | null = null;
let preprocessorPromise: Promise<ModelPreprocessor> | null = null;
let canvasFnsPromise: Promise<{
createCanvasFn: (width: number, height: number) => { getContext: (kind: '2d') => CanvasRenderingContext2D };
loadImageFn: (src: Buffer) => Promise<{ width: number; height: number } & CanvasImageSource>;
}> | null = null;
async function getCanvasFns(): Promise<{
createCanvasFn: (width: number, height: number) => { getContext: (kind: '2d') => CanvasRenderingContext2D };
loadImageFn: (src: Buffer) => Promise<{ width: number; height: number } & CanvasImageSource>;
}> {
if (!canvasFnsPromise) {
canvasFnsPromise = (async () => {
const mod = await import('@napi-rs/canvas');
const namespace = mod as Record<string, unknown>;
const fallback = (namespace.default ?? {}) as Record<string, unknown>;
const createCanvasFn = (namespace.createCanvas ?? fallback.createCanvas) as
| ((width: number, height: number) => { getContext: (kind: '2d') => CanvasRenderingContext2D })
| undefined;
const loadImageFn = (namespace.loadImage ?? fallback.loadImage) as
| ((src: Buffer) => Promise<{ width: number; height: number } & CanvasImageSource>)
| undefined;
if (typeof createCanvasFn !== 'function' || typeof loadImageFn !== 'function') {
throw new Error(
`Canvas runtime missing createCanvas/loadImage exports (keys=${Object.keys(namespace).join(',')}; defaultKeys=${Object.keys(fallback).join(',')})`,
);
}
return { createCanvasFn, loadImageFn };
})();
}
return canvasFnsPromise;
}
async function getSession(): Promise<ort.InferenceSession> {
if (!sessionPromise) {
sessionPromise = (async () => {
const modelPath = await ensureModel();
return ort.InferenceSession.create(modelPath, {
executionProviders: ['cpu'],
graphOptimizationLevel: 'all',
});
})();
}
return sessionPromise;
}
async function getIdToLabel(): Promise<Record<number, string>> {
if (!idToLabelPromise) {
idToLabelPromise = (async () => {
await ensureModel();
const raw = await readFile(MODEL_CONFIG_PATH, 'utf8');
const parsed = JSON.parse(raw) as { id2label?: Record<string, string> };
const out: Record<number, string> = {};
for (const [key, value] of Object.entries(parsed.id2label ?? {})) {
const n = Number(key);
if (Number.isFinite(n)) out[n] = String(value ?? '').trim();
}
return out;
})();
}
return idToLabelPromise;
}
async function getPreprocessor(): Promise<ModelPreprocessor> {
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 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, preprocessor.inputWidth, preprocessor.inputHeight);
ctx.imageSmoothingEnabled = true;
ctx.imageSmoothingQuality = 'high';
ctx.drawImage(image, 0, 0, preprocessor.inputWidth, preprocessor.inputHeight);
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 new ort.Tensor('float32', chw, [1, 3, preprocessor.inputHeight, preprocessor.inputWidth]);
}
function clampBox(
bbox: [number, number, number, number],
pageWidth: number,
pageHeight: number,
): [number, number, number, number] | null {
const x0 = Math.max(0, Math.min(pageWidth, bbox[0]));
const y0 = Math.max(0, Math.min(pageHeight, bbox[1]));
const x1 = Math.max(0, Math.min(pageWidth, bbox[2]));
const y1 = Math.max(0, Math.min(pageHeight, bbox[3]));
if (x1 <= x0 || y1 <= y0) return null;
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<LayoutRegion[]> {
const { pageWidth, pageHeight, textItems, pageImage } = input;
if (!textItems.length) return [];
if (!pageImage || pageImage.byteLength === 0) {
throw new Error('layout-render-missing-page-image');
}
try {
const [session, idToLabel, preprocessor, canvasFns] = await Promise.all([
getSession(),
getIdToLabel(),
getPreprocessor(),
getCanvasFns(),
]);
const decodedPageImage = await canvasFns.loadImageFn(pageImage);
const pixelValues = preprocessResized(decodedPageImage, preprocessor, canvasFns.createCanvasFn);
const output = await session.run({ pixel_values: pixelValues });
const logits = output.logits?.data as Float32Array | undefined;
const predBoxes = output.pred_boxes?.data as Float32Array | undefined;
if (!logits || !predBoxes) return [];
const numQueries = Math.floor(predBoxes.length / 4);
if (numQueries <= 0) return [];
const classCount = Math.floor(logits.length / numQueries);
if (classCount <= 0) return [];
const regions: LayoutRegion[] = [];
for (let queryIdx = 0; queryIdx < numQueries; queryIdx += 1) {
const cls = softmaxMax(logits, queryIdx * classCount, classCount);
const rawLabel = idToLabel[cls.index];
if (!rawLabel) continue;
const mapped = LABEL_MAP[normalizeModelLabel(rawLabel)];
if (!mapped) continue;
const minScore = CLASS_MIN_SCORE[mapped] ?? MIN_SCORE;
if (!Number.isFinite(cls.score) || cls.score < minScore) continue;
const cx = predBoxes[queryIdx * 4 + 0] * pageWidth;
const cy = predBoxes[queryIdx * 4 + 1] * pageHeight;
const w = predBoxes[queryIdx * 4 + 2] * pageWidth;
const h = predBoxes[queryIdx * 4 + 3] * pageHeight;
const rawBox: [number, number, number, number] = [
cx - w / 2,
cy - h / 2,
cx + w / 2,
cy + h / 2,
];
const clamped = clampBox(rawBox, pageWidth, pageHeight);
if (!clamped) continue;
const sizeRule = MIN_REGION_SIZE[mapped];
if (sizeRule) {
const width = clamped[2] - clamped[0];
const height = clamped[3] - clamped[1];
if (width < sizeRule.minWidth || height < sizeRule.minHeight) continue;
}
regions.push({
bbox: clamped,
label: mapped,
confidence: cls.score,
});
}
return regions.sort((a, b) => (b.confidence ?? 0) - (a.confidence ?? 0));
} catch (error) {
throw new Error(
`layout-model-inference-failed: ${error instanceof Error ? error.message : String(error)}`,
);
}
}

View file

@ -0,0 +1,144 @@
import type { ParsedPdfDocument, ParsedPdfBlock } from '../types/parsed-pdf';
const STITCHABLE_KINDS = new Set<ParsedPdfBlock['kind']>([
'text',
'content',
'reference_content',
'aside_text',
'abstract',
'algorithm',
'reference',
]);
function stripTrailingClosers(text: string): string {
return text.trim().replace(/[\"'”’\]\)]+$/g, '');
}
function isSentenceTerminal(text: string): boolean {
return /[.!?]$/.test(stripTrailingClosers(text));
}
function canStitch(a: ParsedPdfBlock, b: ParsedPdfBlock): boolean {
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();
if (!next) return false;
if (/^[A-Z]/.test(next)) return false;
return true;
}
function splitHeadContinuation(text: string): { continuation: string; remainder: string } {
const normalized = text.replace(/\s+/g, ' ').trim();
if (!normalized) return { continuation: '', remainder: '' };
const CLOSERS = new Set(['"', "'", '”', '', ')', ']', '}']);
const isTerminal = (ch: string): boolean => ch === '.' || ch === '!' || ch === '?';
for (let i = 0; i < normalized.length; i += 1) {
const ch = normalized[i];
if (!isTerminal(ch)) continue;
const prev = i > 0 ? normalized[i - 1] : '';
const next = i + 1 < normalized.length ? normalized[i + 1] : '';
if (ch === '.' && /\d/.test(prev) && /\d/.test(next)) continue;
let cut = i + 1;
while (cut < normalized.length && CLOSERS.has(normalized[cut])) cut += 1;
const after = cut < normalized.length ? normalized[cut] : '';
if (!after || /\s/.test(after) || /[A-Z]/.test(after)) {
return {
continuation: normalized.slice(0, cut).trim(),
remainder: normalized.slice(cut).trim(),
};
}
}
return {
continuation: normalized,
remainder: '',
};
}
const HARD_BOUNDARY_KINDS = new Set<ParsedPdfBlock['kind']>([
'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 (STITCHABLE_KINDS.has(block.kind)) return i;
}
return -1;
}
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 (STITCHABLE_KINDS.has(block.kind)) return i;
}
return -1;
}
function hasHardBoundaryBetween(
pageBlocks: ParsedPdfBlock[],
startInclusive: number,
endExclusive: number,
): boolean {
for (let i = startInclusive; i < endExclusive; i += 1) {
const block = pageBlocks[i];
if (block && HARD_BOUNDARY_KINDS.has(block.kind)) return true;
}
return false;
}
export function stitchCrossPageBlocks(doc: ParsedPdfDocument): ParsedPdfDocument {
const pages = doc.pages.map((page) => ({ ...page, blocks: page.blocks.map((b) => ({ ...b, fragments: b.fragments.map((f) => ({ ...f })) })) }));
for (let i = 0; i < pages.length - 1; i += 1) {
const page = pages[i];
const next = pages[i + 1];
const tailIndex = findTailCandidateIndex(page.blocks);
const headIndex = findHeadCandidateIndex(next.blocks);
if (tailIndex < 0 || headIndex < 0) continue;
if (hasHardBoundaryBetween(page.blocks, tailIndex + 1, page.blocks.length)) continue;
if (hasHardBoundaryBetween(next.blocks, 0, headIndex)) continue;
const tail = page.blocks[tailIndex];
const head = next.blocks[headIndex];
if (!tail || !head) continue;
if (!canStitch(tail, head)) continue;
const { continuation, remainder } = splitHeadContinuation(head.text);
if (!continuation) continue;
const continuationFragment = head.fragments[0]
? { ...head.fragments[0], text: continuation }
: null;
if (continuationFragment) {
tail.fragments.push(continuationFragment);
}
tail.text = `${tail.text} ${continuation}`.replace(/\s+/g, ' ').trim();
if (!remainder) {
next.blocks.splice(headIndex, 1);
continue;
}
head.text = remainder;
if (head.fragments[0]) {
head.fragments[0].text = remainder;
}
}
return {
...doc,
pages,
};
}

View file

@ -0,0 +1,15 @@
import type { ParsedPdfBlockKind } from '../types/parsed-pdf';
export interface PdfTextItem {
text: string;
x: number;
y: number;
width: number;
height: number;
}
export interface LayoutRegion {
bbox: [number, number, number, number];
label: ParsedPdfBlockKind;
confidence?: number;
}

View file

@ -0,0 +1,7 @@
import path from 'path';
export const DOCSTORE_DIR = process.env.COMPUTE_DOCSTORE_DIR?.trim() || path.join(process.cwd(), 'docstore');
export function getDocstoreDir(): string {
return DOCSTORE_DIR;
}

View file

@ -0,0 +1,36 @@
import { existsSync } from 'fs';
import ffmpegStatic from 'ffmpeg-static';
function normalizePath(value: unknown): string | null {
if (typeof value !== 'string') return null;
const trimmed = value.trim();
return trimmed.length > 0 ? trimmed : null;
}
function resolveBinary(envValue: string | null, bundledValue: string | null, envVarName: string, packageName: string): string {
if (envValue) {
if ((envValue.includes('/') || envValue.includes('\\')) && !existsSync(envValue)) {
throw new Error(`${envVarName} points to a missing binary: ${envValue}`);
}
return envValue;
}
if (!bundledValue) {
throw new Error(`${packageName} binary is unavailable on this platform. Set ${envVarName} to an installed binary path.`);
}
if ((bundledValue.includes('/') || bundledValue.includes('\\')) && !existsSync(bundledValue)) {
throw new Error(`${packageName} resolved to a missing binary path: ${bundledValue}`);
}
return bundledValue;
}
export function getFFmpegPath(): string {
return resolveBinary(
normalizePath(process.env.FFMPEG_BIN),
normalizePath(ffmpegStatic),
'FFMPEG_BIN',
'ffmpeg-static',
);
}

View file

@ -0,0 +1,53 @@
export type 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;
bbox: [number, number, number, number];
text: string;
readingOrder: number;
modelConfidence?: number;
}
export interface ParsedPdfBlock {
id: string;
kind: ParsedPdfBlockKind;
fragments: ParsedPdfBlockFragment[];
text: string;
parentSectionId?: string;
}
export interface ParsedPdfPage {
pageNumber: number;
width: number;
height: number;
blocks: ParsedPdfBlock[];
}
export interface ParsedPdfDocument {
schemaVersion: 1;
documentId: string;
parserVersion: string;
parsedAt: number;
pages: ParsedPdfPage[];
}

View file

@ -0,0 +1,16 @@
export type TTSAudioBuffer = ArrayBuffer;
export type TTSAudioBytes = number[];
export interface TTSSentenceWord {
text: string;
startSec: number;
endSec: number;
charStart: number;
charEnd: number;
}
export interface TTSSentenceAlignment {
sentence: string;
sentenceIndex: number;
words: TTSSentenceWord[];
}

View file

@ -0,0 +1,54 @@
import type { TTSSentenceAlignment, TTSSentenceWord } from '../types/tts';
function preprocessSentenceForAudio(text: string): string {
return text
.replace(/\S*(?:https?:\/\/|www\.)([^\/\s]+)(?:\/\S*)?/gi, '- (link to $1) -')
.replace(/(\w+)-\s+(\w+)/g, '$1$2')
.replace(/\*/g, '')
.replace(/\s+/g, ' ')
.trim();
}
export interface WhisperWord {
start: number;
end: number;
word: string;
}
export function mapWordsToSentenceOffsets(sentence: string, words: WhisperWord[]): TTSSentenceAlignment {
const normalizedSentence = preprocessSentenceForAudio(sentence);
const lowerSentence = normalizedSentence.toLowerCase();
let cursor = 0;
const alignedWords: TTSSentenceWord[] = words.map((w) => {
const token = w.word.trim();
if (!token) {
return {
text: '',
startSec: w.start,
endSec: w.end,
charStart: cursor,
charEnd: cursor,
};
}
const idx = lowerSentence.indexOf(token.toLowerCase(), cursor);
const start = idx >= 0 ? idx : cursor;
const end = Math.min(normalizedSentence.length, start + token.length);
cursor = Math.max(cursor, end);
return {
text: token,
startSec: w.start,
endSec: w.end,
charStart: start,
charEnd: end,
};
}).filter((word) => word.text.length > 0);
return {
sentence,
sentenceIndex: 0,
words: alignedWords,
};
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,242 @@
import path from 'path';
import { fileURLToPath } from 'url';
import { createHash } from 'crypto';
import { access, copyFile, mkdir, readFile, rename, unlink, writeFile } from 'fs/promises';
import { DOCSTORE_DIR } from '../runtime/docstore';
import manifest from './model/manifest.json';
const MODULE_DIR = path.dirname(fileURLToPath(import.meta.url));
const MODEL_DIR = path.join(DOCSTORE_DIR, 'model', 'whisper-base_timestamped');
const STATIC_LICENSE_PATH = path.join(MODULE_DIR, 'model', 'LICENSE.txt');
export const WHISPER_CONFIG_PATH = path.join(MODEL_DIR, 'config.json');
export const WHISPER_GENERATION_CONFIG_PATH = path.join(MODEL_DIR, 'generation_config.json');
export const WHISPER_TOKENIZER_PATH = path.join(MODEL_DIR, 'tokenizer.json');
export const WHISPER_TOKENIZER_CONFIG_PATH = path.join(MODEL_DIR, 'tokenizer_config.json');
export const WHISPER_ENCODER_MODEL_PATH = path.join(MODEL_DIR, 'onnx', 'encoder_model_int8.onnx');
export const WHISPER_DECODER_MERGED_MODEL_PATH = path.join(MODEL_DIR, 'onnx', 'decoder_model_merged_int8.onnx');
export const WHISPER_DECODER_WITH_PAST_MODEL_PATH = path.join(MODEL_DIR, 'onnx', 'decoder_with_past_model_int8.onnx');
const BASE_MODEL_URL = 'https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main';
const WHISPER_MODEL_BASE_URL_ENV = 'WHISPER_MODEL_BASE_URL';
const MODEL_RELATIVE_PATHS: string[] = [
'config.json',
'generation_config.json',
'tokenizer.json',
'tokenizer_config.json',
'merges.txt',
'vocab.json',
'normalizer.json',
'added_tokens.json',
'preprocessor_config.json',
'special_tokens_map.json',
'onnx/encoder_model_int8.onnx',
'onnx/decoder_model_merged_int8.onnx',
'onnx/decoder_with_past_model_int8.onnx',
];
const DEFAULT_URLS: Record<string, string> = {
'config.json': `${BASE_MODEL_URL}/config.json`,
'generation_config.json': `${BASE_MODEL_URL}/generation_config.json`,
'tokenizer.json': `${BASE_MODEL_URL}/tokenizer.json`,
'tokenizer_config.json': `${BASE_MODEL_URL}/tokenizer_config.json`,
'merges.txt': `${BASE_MODEL_URL}/merges.txt`,
'vocab.json': `${BASE_MODEL_URL}/vocab.json`,
'normalizer.json': `${BASE_MODEL_URL}/normalizer.json`,
'added_tokens.json': `${BASE_MODEL_URL}/added_tokens.json`,
'preprocessor_config.json': `${BASE_MODEL_URL}/preprocessor_config.json`,
'special_tokens_map.json': `${BASE_MODEL_URL}/special_tokens_map.json`,
'onnx/encoder_model_int8.onnx': `${BASE_MODEL_URL}/onnx/encoder_model_int8.onnx`,
'onnx/decoder_model_merged_int8.onnx': `${BASE_MODEL_URL}/onnx/decoder_model_merged_int8.onnx`,
'onnx/decoder_with_past_model_int8.onnx': `${BASE_MODEL_URL}/onnx/decoder_with_past_model_int8.onnx`,
};
type ManifestEntry = { path: string; sha256?: string; size?: number };
export interface WhisperArtifactSpec {
path: string;
sha256?: string;
size?: number;
url: string;
}
export interface WhisperStaticArtifactSpec {
path: string;
sha256?: string;
size?: number;
sourcePath: string;
}
export type WhisperFetch = (input: RequestInfo | URL, init?: RequestInit) => Promise<Response>;
const MANIFEST_FILES = manifest.files as ManifestEntry[];
const MODEL_FILES = MANIFEST_FILES.filter((entry) => entry.path !== 'LICENSE.txt');
const LICENSE_FILE = MANIFEST_FILES.find((entry) => entry.path === 'LICENSE.txt');
function normalizeExpected(entry: { sha256?: string; size?: number }): { sha256: string | null; size: number } {
return {
sha256: typeof entry.sha256 === 'string' ? entry.sha256.toLowerCase() : null,
size: Number(entry.size ?? 0),
};
}
function resolvePath(relativePath: string, modelDir: string): string {
return path.join(modelDir, relativePath);
}
function joinModelUrl(baseUrl: string, relativePath: string): string {
return `${baseUrl.replace(/\/+$/, '')}/${relativePath}`;
}
function resolveUrl(relativePath: string): string {
const overrideBase = process.env[WHISPER_MODEL_BASE_URL_ENV]?.trim();
if (overrideBase) {
return joinModelUrl(overrideBase, relativePath);
}
const fallback = DEFAULT_URLS[relativePath];
if (!fallback) {
throw new Error(`No default URL configured for Whisper model artifact: ${relativePath}`);
}
return fallback;
}
function sha256OfBytes(bytes: Uint8Array): string {
return createHash('sha256').update(bytes).digest('hex');
}
function verifyBytes(bytes: Uint8Array, expected: { sha256?: string; size?: number }): boolean {
const normalized = normalizeExpected(expected);
if (Number.isFinite(normalized.size) && normalized.size > 0 && bytes.byteLength !== normalized.size) {
return false;
}
if (!normalized.sha256) return true;
return sha256OfBytes(bytes) === normalized.sha256;
}
async function verifyFile(filePath: string, expected: { sha256?: string; size?: number }): Promise<boolean> {
const bytes = await readFile(filePath);
return verifyBytes(bytes, expected);
}
async function downloadToFile(fetchImpl: WhisperFetch, url: string, outPath: string): Promise<void> {
const res = await fetchImpl(url);
if (!res.ok) {
throw new Error(`Download failed for ${url}: ${res.status} ${res.statusText}`);
}
const bytes = new Uint8Array(await res.arrayBuffer());
await writeFile(outPath, bytes);
}
export async function ensureWhisperArtifacts(options: {
modelDir: string;
artifacts: WhisperArtifactSpec[];
staticArtifacts?: WhisperStaticArtifactSpec[];
fetchImpl?: WhisperFetch;
}): Promise<void> {
const {
modelDir,
artifacts,
staticArtifacts = [],
fetchImpl = fetch,
} = options;
try {
await Promise.all(artifacts.map(async (artifact) => {
const target = resolvePath(artifact.path, modelDir);
await access(target);
const valid = await verifyFile(target, artifact);
if (!valid) {
throw new Error(`Checksum mismatch for existing Whisper artifact: ${artifact.path}`);
}
}));
await Promise.all(staticArtifacts.map(async (artifact) => {
const target = resolvePath(artifact.path, modelDir);
await access(target);
const valid = await verifyFile(target, artifact);
if (!valid) {
throw new Error(`Checksum mismatch for existing Whisper static artifact: ${artifact.path}`);
}
}));
return;
} catch {
// Continue to repair/download.
}
for (const artifact of artifacts) {
const target = resolvePath(artifact.path, modelDir);
const targetDir = path.dirname(target);
const tmp = `${target}.tmp`;
await mkdir(targetDir, { recursive: true });
await downloadToFile(fetchImpl, artifact.url, tmp);
if (!(await verifyFile(tmp, artifact))) {
await unlink(tmp).catch(() => undefined);
throw new Error(`Whisper artifact checksum verification failed: ${artifact.path}`);
}
await rename(tmp, target);
}
for (const artifact of staticArtifacts) {
const target = resolvePath(artifact.path, modelDir);
const targetDir = path.dirname(target);
await mkdir(targetDir, { recursive: true });
await copyFile(artifact.sourcePath, target);
if (!(await verifyFile(target, artifact))) {
throw new Error(`Whisper static artifact checksum verification failed: ${artifact.path}`);
}
}
}
export function createSingleflightRunner<T>(work: () => Promise<T>): () => Promise<T> {
let inflight: Promise<T> | null = null;
return async () => {
if (inflight) return inflight;
inflight = work().finally(() => {
inflight = null;
});
return inflight;
};
}
async function ensureModelInternal(): Promise<string> {
if (process.env[WHISPER_MODEL_BASE_URL_ENV]?.trim()) {
for (const relativePath of MODEL_RELATIVE_PATHS) {
if (!(relativePath in DEFAULT_URLS)) {
throw new Error(`Missing default URL path mapping for Whisper artifact: ${relativePath}`);
}
}
}
const artifacts: WhisperArtifactSpec[] = MODEL_FILES.map((entry) => ({
path: entry.path,
sha256: entry.sha256,
size: entry.size,
url: resolveUrl(entry.path),
}));
const staticArtifacts: WhisperStaticArtifactSpec[] = LICENSE_FILE
? [{
path: LICENSE_FILE.path,
sha256: LICENSE_FILE.sha256,
size: LICENSE_FILE.size,
sourcePath: STATIC_LICENSE_PATH,
}]
: [];
await ensureWhisperArtifacts({
modelDir: MODEL_DIR,
artifacts,
staticArtifacts,
});
return WHISPER_ENCODER_MODEL_PATH;
}
const ensureWhisperModelSingleflight = createSingleflightRunner(ensureModelInternal);
export async function ensureWhisperModel(): Promise<string> {
return ensureWhisperModelSingleflight();
}

View file

@ -0,0 +1,21 @@
MIT License
Copyright (c) 2022 OpenAI
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

View file

@ -0,0 +1,76 @@
{
"name": "whisper-base_timestamped-int8",
"version": "onnx-community/whisper-base_timestamped@608c49e61301901684bc36cac8f74b95ff6b5a8e",
"files": [
{
"path": "config.json",
"sha256": "f4d0608f7d918166da7edb3e188de5ef1bfe70d9802e785d271fd88111e9cf4b",
"size": 2243
},
{
"path": "generation_config.json",
"sha256": "61070cf8de25b1e9256e8e102ded49d8d24a8369ed36ef84fdf21549e68125a0",
"size": 3832
},
{
"path": "tokenizer.json",
"sha256": "27fc476bfe7f17299480be2273fc0608e4d5a99aba2ab5dec5374b4482d1a566",
"size": 2480466
},
{
"path": "tokenizer_config.json",
"sha256": "2e036e4dbacfdeb7242c7d4ec4149f4a16e86026048f94d1637e3a8ee9c6a573",
"size": 282682
},
{
"path": "merges.txt",
"sha256": "2df2990a395e35e8dfbc7511e08c12d56018d8d04691e0133e5d63b21e154dc6",
"size": 493869
},
{
"path": "vocab.json",
"sha256": "50d6a919f0a0601d56a04eb583c780d18553aa388254ba3158eb6a00f13e2c1a",
"size": 1036584
},
{
"path": "normalizer.json",
"sha256": "bf1c507dc8724ca9cf9903640dacfb69dae2f00edee4f21ceba106a7392f26dd",
"size": 52666
},
{
"path": "added_tokens.json",
"sha256": "9715fd2243b6f06a5858b5e32950d2853f73dd5bc201aafcf76f5082a2d8acd1",
"size": 34604
},
{
"path": "preprocessor_config.json",
"sha256": "a6a76d28c93edb273669eb9e0b0636a2bddbb1272c3261e47b7ca6dfdbac1b8d",
"size": 339
},
{
"path": "special_tokens_map.json",
"sha256": "e67ae3a0aaa99abcd9f187138e12db1f65c16a14761c50ef10eef2c174a7a691",
"size": 2194
},
{
"path": "onnx/encoder_model_int8.onnx",
"sha256": "152da96dd8ff3f28f3fadabc2e8960405a277846453ff94ed411fe935a72917f",
"size": 23159150
},
{
"path": "onnx/decoder_model_merged_int8.onnx",
"sha256": "cf9a8d5bcddc0917a0078135b484cedcaf44f28909cd91910abd29dced9171db",
"size": 53712708
},
{
"path": "onnx/decoder_with_past_model_int8.onnx",
"sha256": "bdd92860d0ed7dff2aca623963378cbba1b617bfae127356db1c8aa8baa930ef",
"size": 50131672
},
{
"path": "LICENSE.txt",
"sha256": "b5d65a59060e68c4ff940e1eddfa6f94b2d68fdf58ed7f4dd57721c997e35e9d",
"size": 1063
}
]
}

Binary file not shown.

View file

@ -0,0 +1,21 @@
export function buildGoertzelCoefficients(freqBins: number, fftSize: number): Float64Array {
const coeffs = new Float64Array(freqBins);
for (let k = 0; k < freqBins; k += 1) {
coeffs[k] = 2 * Math.cos((2 * Math.PI * k) / fftSize);
}
return coeffs;
}
export function goertzelPower(samples: Float32Array, coeff: number): number {
let s1 = 0;
let s2 = 0;
for (let i = 0; i < samples.length; i += 1) {
const s0 = samples[i] + (coeff * s1) - s2;
s2 = s1;
s1 = s0;
}
const power = (s1 * s1) + (s2 * s2) - (coeff * s1 * s2);
if (!Number.isFinite(power) || power < 0) return 0;
return power;
}

View file

@ -0,0 +1,449 @@
import type { Tokenizer } from '@huggingface/tokenizers';
import type * as ort from 'onnxruntime-node';
const PUNCTUATION_REGEX = '\\p{P}\\u0021-\\u002F\\u003A-\\u0040\\u005B-\\u0060\\u007B-\\u007E';
const PUNCTUATION_ONLY_REGEX = new RegExp(`^[${PUNCTUATION_REGEX}]+$`, 'gu');
type TokenTimestamp = [start: number, end: number];
export interface WhisperWordTiming {
word: string;
startSec: number;
endSec: number;
}
function medianFilter(data: Float32Array, windowSize: number): Float32Array {
if (windowSize % 2 === 0 || windowSize <= 0) {
throw new Error('Window size must be a positive odd number');
}
const output = new Float32Array(data.length);
const buffer = new Float32Array(windowSize);
const halfWindow = Math.floor(windowSize / 2);
for (let i = 0; i < data.length; i += 1) {
let valuesIndex = 0;
for (let j = -halfWindow; j <= halfWindow; j += 1) {
let index = i + j;
if (index < 0) {
index = Math.abs(index);
} else if (index >= data.length) {
index = (2 * (data.length - 1)) - index;
}
buffer[valuesIndex] = data[index];
valuesIndex += 1;
}
const sortable = Array.from(buffer);
sortable.sort((a, b) => a - b);
output[i] = sortable[halfWindow] ?? 0;
}
return output;
}
function dynamicTimeWarping(matrix: Float32Array[], rows: number, cols: number): [number[], number[]] {
const cost: number[][] = Array.from({ length: rows + 1 }, () => Array(cols + 1).fill(Number.POSITIVE_INFINITY));
const trace: number[][] = Array.from({ length: rows + 1 }, () => Array(cols + 1).fill(-1));
cost[0][0] = 0;
for (let j = 1; j <= cols; j += 1) {
for (let i = 1; i <= rows; i += 1) {
const c0 = cost[i - 1][j - 1];
const c1 = cost[i - 1][j];
const c2 = cost[i][j - 1];
let c: number;
let t: number;
if (c0 < c1 && c0 < c2) {
c = c0;
t = 0;
} else if (c1 < c0 && c1 < c2) {
c = c1;
t = 1;
} else {
c = c2;
t = 2;
}
cost[i][j] = matrix[i - 1][j - 1] + c;
trace[i][j] = t;
}
}
for (let i = 0; i <= cols; i += 1) trace[0][i] = 2;
for (let i = 0; i <= rows; i += 1) trace[i][0] = 1;
let i = rows;
let j = cols;
const textIndices: number[] = [];
const timeIndices: number[] = [];
while (i > 0 || j > 0) {
textIndices.push(i - 1);
timeIndices.push(j - 1);
const step = trace[i][j];
if (step === 0) {
i -= 1;
j -= 1;
} else if (step === 1) {
i -= 1;
} else if (step === 2) {
j -= 1;
} else {
throw new Error(`Unexpected DTW trace state at [${i}, ${j}]`);
}
}
textIndices.reverse();
timeIndices.reverse();
return [textIndices, timeIndices];
}
function round2(value: number): number {
return Math.round(value * 100) / 100;
}
function decodeTokens(tokenizer: Pick<Tokenizer, 'decode'>, tokens: number[]): string {
return tokenizer.decode(tokens, { skip_special_tokens: false });
}
function splitTokensOnUnicode(
tokenizer: Pick<Tokenizer, 'decode'>,
tokens: number[],
): [string[], number[][], number[][]] {
const decodedFull = decodeTokens(tokenizer, tokens);
const replacementChar = '\uFFFD';
const words: string[] = [];
const wordTokens: number[][] = [];
const tokenIndices: number[][] = [];
let currentTokens: number[] = [];
let currentIndices: number[] = [];
let unicodeOffset = 0;
for (let i = 0; i < tokens.length; i += 1) {
currentTokens.push(tokens[i]);
currentIndices.push(i);
const decoded = decodeTokens(tokenizer, currentTokens);
if (
!decoded.includes(replacementChar)
|| decodedFull[unicodeOffset + decoded.indexOf(replacementChar)] === replacementChar
) {
words.push(decoded);
wordTokens.push(currentTokens);
tokenIndices.push(currentIndices);
currentTokens = [];
currentIndices = [];
unicodeOffset += decoded.length;
}
}
return [words, wordTokens, tokenIndices];
}
function splitTokensOnSpaces(
tokenizer: Pick<Tokenizer, 'decode'>,
tokens: number[],
eosTokenId: number,
): [string[], number[][], number[][]] {
const [subwords, subwordTokens, subwordIndices] = splitTokensOnUnicode(tokenizer, tokens);
const words: string[] = [];
const wordTokens: number[][] = [];
const tokenIndices: number[][] = [];
for (let i = 0; i < subwords.length; i += 1) {
const subword = subwords[i];
const tokenList = subwordTokens[i];
const indices = subwordIndices[i];
const special = tokenList[0] >= eosTokenId;
const withSpace = subword.startsWith(' ');
const trimmed = subword.trim();
const punctuation = PUNCTUATION_ONLY_REGEX.test(trimmed);
if (special || withSpace || punctuation || words.length === 0) {
words.push(subword);
wordTokens.push([...tokenList]);
tokenIndices.push([...indices]);
} else {
const ix = words.length - 1;
words[ix] += subword;
wordTokens[ix].push(...tokenList);
tokenIndices[ix].push(...indices);
}
}
return [words, wordTokens, tokenIndices];
}
function mergePunctuations(
words: string[],
tokens: number[][],
indices: number[][],
prependPunctuations = '"\'“¡¿([{-',
appendPunctuations = '"\'.。,!?::”)]}、',
): [string[], number[][], number[][]] {
const newWords = words.map((w) => `${w}`);
const newTokens = tokens.map((t) => [...t]);
const newIndices = indices.map((idx) => [...idx]);
let i = newWords.length - 2;
let j = newWords.length - 1;
while (i >= 0) {
if (newWords[i].startsWith(' ') && prependPunctuations.includes(newWords[i].trim())) {
newWords[j] = newWords[i] + newWords[j];
newTokens[j] = [...newTokens[i], ...newTokens[j]];
newIndices[j] = [...newIndices[i], ...newIndices[j]];
newWords[i] = '';
newTokens[i] = [];
newIndices[i] = [];
} else {
j = i;
}
i -= 1;
}
i = 0;
j = 1;
while (j < newWords.length) {
if (!newWords[i].endsWith(' ') && appendPunctuations.includes(newWords[j])) {
newWords[i] += newWords[j];
newTokens[i] = [...newTokens[i], ...newTokens[j]];
newIndices[i] = [...newIndices[i], ...newIndices[j]];
newWords[j] = '';
newTokens[j] = [];
newIndices[j] = [];
} else {
i = j;
}
j += 1;
}
return [
newWords.filter((w) => w.length > 0),
newTokens.filter((t) => t.length > 0),
newIndices.filter((t) => t.length > 0),
];
}
function combineTokensIntoWords(
tokenizer: Pick<Tokenizer, 'decode'>,
tokens: number[],
eosTokenId: number,
language = 'english',
): [string[], number[][], number[][]] {
let words: string[];
let wordTokens: number[][];
let tokenIndices: number[][];
if (['chinese', 'japanese', 'thai', 'lao', 'myanmar', 'zh', 'ja', 'th', 'lo', 'my'].includes(language)) {
[words, wordTokens, tokenIndices] = splitTokensOnUnicode(tokenizer, tokens);
} else {
[words, wordTokens, tokenIndices] = splitTokensOnSpaces(tokenizer, tokens, eosTokenId);
}
return mergePunctuations(words, wordTokens, tokenIndices);
}
export function extractTokenStartTimestamps(input: {
crossAttentions: Record<string, ort.Tensor>;
decoderLayers: number;
alignmentHeads: Array<[number, number]>;
numFrames: number;
numInputIds: number;
timePrecision?: number;
sequenceLength: number;
}): number[] {
const {
crossAttentions,
decoderLayers,
alignmentHeads,
numFrames,
numInputIds,
timePrecision = 0.02,
sequenceLength,
} = input;
const frameCount = Math.max(1, numFrames);
const perLayer: Float32Array[] = [];
for (let layer = 0; layer < decoderLayers; layer += 1) {
const key = `cross_attentions.${layer}`;
const tensor = crossAttentions[key];
if (!tensor) continue;
perLayer[layer] = tensor.data as Float32Array;
}
const selected: Float32Array[] = [];
let seqLen = 0;
let attnFrames = 0;
for (const [layer, head] of alignmentHeads) {
const flat = perLayer[layer];
if (!flat) continue;
const layerTensor = crossAttentions[`cross_attentions.${layer}`];
if (!layerTensor || layerTensor.dims.length < 4) continue;
const [, numHeads, currentSeqLen, currentFrames] = layerTensor.dims;
if (head >= numHeads) continue;
seqLen = currentSeqLen;
attnFrames = Math.min(currentFrames, frameCount);
const headSlice = new Float32Array(seqLen * attnFrames);
for (let s = 0; s < seqLen; s += 1) {
for (let f = 0; f < attnFrames; f += 1) {
const flatIndex = (((head * currentSeqLen) + s) * currentFrames) + f;
headSlice[(s * attnFrames) + f] = flat[flatIndex] ?? 0;
}
}
selected.push(headSlice);
}
if (!selected.length || seqLen === 0 || attnFrames === 0) {
return new Array(sequenceLength).fill(0);
}
const normalizedHeads = selected.map((headData) => {
const means = new Float32Array(attnFrames);
const stds = new Float32Array(attnFrames);
for (let f = 0; f < attnFrames; f += 1) {
let sum = 0;
for (let s = 0; s < seqLen; s += 1) sum += headData[(s * attnFrames) + f];
const mean = sum / seqLen;
means[f] = mean;
let varSum = 0;
for (let s = 0; s < seqLen; s += 1) {
const d = headData[(s * attnFrames) + f] - mean;
varSum += d * d;
}
stds[f] = Math.sqrt(varSum / seqLen) || 1;
}
const out = new Float32Array(headData.length);
for (let s = 0; s < seqLen; s += 1) {
const row = new Float32Array(attnFrames);
for (let f = 0; f < attnFrames; f += 1) {
row[f] = (headData[(s * attnFrames) + f] - means[f]) / stds[f];
}
const filtered = medianFilter(row, 7);
out.set(filtered, s * attnFrames);
}
return out;
});
const croppedRows = Math.max(0, seqLen - numInputIds);
if (croppedRows === 0) return new Array(sequenceLength).fill(0);
const matrix: Float32Array[] = Array.from({ length: croppedRows }, () => new Float32Array(attnFrames));
for (const headData of normalizedHeads) {
for (let r = 0; r < croppedRows; r += 1) {
const srcRow = r + numInputIds;
for (let f = 0; f < attnFrames; f += 1) {
matrix[r][f] += headData[(srcRow * attnFrames) + f];
}
}
}
const scale = 1 / normalizedHeads.length;
for (let r = 0; r < croppedRows; r += 1) {
for (let f = 0; f < attnFrames; f += 1) {
matrix[r][f] = -matrix[r][f] * scale;
}
}
const [textIndices, timeIndices] = dynamicTimeWarping(matrix, croppedRows, attnFrames);
const jumps = new Array(textIndices.length).fill(false);
for (let i = 0; i < textIndices.length; i += 1) {
jumps[i] = i === 0 ? true : textIndices[i] !== textIndices[i - 1];
}
const jumpTimes: number[] = [];
for (let i = 0; i < jumps.length; i += 1) {
if (jumps[i]) jumpTimes.push(timeIndices[i] * timePrecision);
}
const timestamps = new Array(sequenceLength).fill(0);
for (let i = 0; i < numInputIds && i < timestamps.length; i += 1) timestamps[i] = 0;
for (let i = 0; i < jumpTimes.length && (numInputIds + i) < timestamps.length; i += 1) {
timestamps[numInputIds + i] = jumpTimes[i];
}
if (timestamps.length > 0 && jumpTimes.length > 0) {
timestamps[timestamps.length - 1] = jumpTimes[jumpTimes.length - 1];
}
return timestamps;
}
export function buildWordsFromTimestampedTokens(input: {
tokens: number[];
tokenStartTimestamps: number[];
tokenizer: Pick<Tokenizer, 'decode'>;
eosTokenId: number;
promptLength: number;
timestampBeginTokenId: number;
timePrecision?: number;
language?: string;
}): WhisperWordTiming[] {
const {
tokens,
tokenStartTimestamps,
tokenizer,
eosTokenId,
promptLength,
timestampBeginTokenId,
timePrecision = 0.02,
language = 'english',
} = input;
const limit = Math.min(tokens.length, tokenStartTimestamps.length);
const tokenRanges: TokenTimestamp[] = [];
for (let i = 0; i < limit; i += 1) {
const start = tokenStartTimestamps[i] ?? 0;
const end = i + 1 < limit ? (tokenStartTimestamps[i + 1] ?? (start + timePrecision)) : (start + timePrecision);
tokenRanges.push([start, Math.max(start, end)]);
}
const words: WhisperWordTiming[] = [];
let segmentStart: number | null = null;
let textTokens: number[] = [];
let textRanges: TokenTimestamp[] = [];
const flushSegment = (segmentEnd: number | null) => {
if (!textTokens.length) return;
const [wordTexts, , tokenIndices] = combineTokensIntoWords(tokenizer, textTokens, eosTokenId, language);
for (let i = 0; i < wordTexts.length; i += 1) {
const indices = tokenIndices[i];
if (!indices.length) continue;
const start = textRanges[indices[0]]?.[0] ?? segmentStart ?? 0;
const end = textRanges[indices[indices.length - 1]]?.[1] ?? segmentEnd ?? start;
const clampedStart = segmentStart == null ? start : Math.max(segmentStart, start);
const clampedEndBase = segmentEnd == null ? end : Math.min(segmentEnd, end);
const clampedEnd = Math.max(
clampedStart + (clampedEndBase <= clampedStart ? timePrecision : 0),
clampedEndBase,
);
words.push({
word: wordTexts[i].trim(),
startSec: round2(clampedStart),
endSec: round2(clampedEnd),
});
}
textTokens = [];
textRanges = [];
};
for (let i = promptLength; i < limit; i += 1) {
const token = tokens[i];
if (token === eosTokenId) break;
if (token >= timestampBeginTokenId) {
const ts = (token - timestampBeginTokenId) * timePrecision;
if (segmentStart == null) {
segmentStart = ts;
} else {
flushSegment(ts);
segmentStart = ts;
}
continue;
}
textTokens.push(token);
textRanges.push(tokenRanges[i]);
}
flushSegment(null);
return words.filter((w) => w.word.length > 0);
}

View file

@ -0,0 +1,7 @@
{
"extends": "../../tsconfig.json",
"compilerOptions": {
"noEmit": true
},
"include": ["src/**/*.ts"]
}

View file

@ -0,0 +1,25 @@
# Compute worker bind
COMPUTE_WORKER_HOST=0.0.0.0
COMPUTE_WORKER_PORT=8081
COMPUTE_LOG_FORMAT=pretty
# COMPUTE_LOG_LEVEL=info
# App <-> worker auth
COMPUTE_WORKER_TOKEN=local-compute-token
# Redis/BullMQ
REDIS_URL=redis://redis:6379
# Shared object storage (must be reachable from worker)
S3_BUCKET=openreader-documents
S3_REGION=us-east-1
S3_ACCESS_KEY_ID=devkey
S3_SECRET_ACCESS_KEY=devsecret
S3_PREFIX=openreader
# Optional for non-AWS S3-compatible endpoints:
S3_ENDPOINT=http://host.docker.internal:8333
S3_FORCE_PATH_STYLE=true
# Queue + execution tuning
COMPUTE_QUEUE_MAX_DEPTH=64
COMPUTE_PREWARM_MODELS=true

20
compute/worker/Dockerfile Normal file
View file

@ -0,0 +1,20 @@
FROM node:lts
RUN npm install -g pnpm@11.1.2
WORKDIR /workspace
COPY package.json pnpm-lock.yaml pnpm-workspace.yaml ./
COPY compute/worker/package.json compute/worker/package.json
COPY compute/core/package.json compute/core/package.json
RUN pnpm install --frozen-lockfile
COPY . .
ENV COMPUTE_WORKER_HOST=0.0.0.0
ENV COMPUTE_WORKER_PORT=8081
EXPOSE 8081
CMD ["pnpm", "--filter", "@openreader/compute-worker", "dev"]

View file

@ -0,0 +1,41 @@
services:
redis:
image: redis:7-alpine
container_name: openreader-compute-redis
ports:
- "6379:6379"
compute-worker:
build:
context: ../..
dockerfile: compute/worker/Dockerfile
container_name: openreader-compute-worker
depends_on:
- redis
env_file:
- ./.env
environment:
REDIS_URL: ${REDIS_URL:-redis://redis:6379}
COMPUTE_WORKER_HOST: ${COMPUTE_WORKER_HOST:-0.0.0.0}
COMPUTE_WORKER_PORT: ${COMPUTE_WORKER_PORT:-8081}
COMPUTE_WORKER_TOKEN: ${COMPUTE_WORKER_TOKEN:-local-compute-token}
S3_PREFIX: ${S3_PREFIX:-openreader}
COMPUTE_PREWARM_MODELS: ${COMPUTE_PREWARM_MODELS:-true}
ports:
- "8081:8081"
develop:
watch:
- action: sync+restart
path: .
target: /workspace/compute/worker
- action: sync+restart
path: ../../compute/core
target: /workspace/compute/core
- action: rebuild
path: ./package.json
- action: rebuild
path: ../../compute/core/package.json
- action: rebuild
path: ../../pnpm-lock.yaml
- action: rebuild
path: ../../pnpm-workspace.yaml

View file

@ -0,0 +1,23 @@
{
"name": "@openreader/compute-worker",
"version": "0.0.0",
"private": true,
"type": "module",
"scripts": {
"dev": "tsx watch src/server.ts",
"start": "tsx src/server.ts"
},
"dependencies": {
"@aws-sdk/client-s3": "^3.1045.0",
"@openreader/compute-core": "workspace:*",
"bullmq": "^5.61.2",
"fastify": "^5.6.2",
"ioredis": "^5.8.2",
"pino": "^9.14.0",
"pino-pretty": "^13.1.2",
"zod": "^4.1.12"
},
"devDependencies": {
"tsx": "^4.20.6"
}
}

View file

@ -0,0 +1,432 @@
import Fastify, { type FastifyReply, type FastifyRequest } from 'fastify';
import { Queue, Worker, type Job, type JobsOptions } from 'bullmq';
import IORedis from 'ioredis';
import { z } from 'zod';
import {
ALIGN_QUEUE_NAME,
PDF_LAYOUT_QUEUE_NAME,
ensureComputeModels,
runPdfLayoutFromPdfBuffer,
runWhisperAlignmentFromAudioBuffer,
type PdfLayoutJobRequest,
type PdfLayoutJobResult,
type WhisperAlignJobRequest,
type WhisperAlignJobResult,
type WorkerJobStatusResponse,
} from '@openreader/compute-core';
import { GetObjectCommand, S3Client } from '@aws-sdk/client-s3';
function requireEnv(name: string): string {
const value = process.env[name]?.trim();
if (!value) throw new Error(`${name} is required`);
return value;
}
function readIntEnv(name: string, fallback: number): number {
const raw = process.env[name]?.trim();
if (!raw) return fallback;
const parsed = Number(raw);
if (!Number.isFinite(parsed) || parsed <= 0) return fallback;
return Math.floor(parsed);
}
function parseBoolEnv(name: string, fallback: boolean): boolean {
const raw = process.env[name]?.trim();
if (!raw) return fallback;
const normalized = raw.toLowerCase();
return normalized === '1' || normalized === 'true' || normalized === 'yes' || normalized === 'on';
}
function buildLoggerConfig(): boolean | Record<string, unknown> {
const format = (process.env.COMPUTE_LOG_FORMAT?.trim().toLowerCase() || 'pretty');
if (format === 'json') return true;
return {
level: process.env.COMPUTE_LOG_LEVEL?.trim() || 'info',
transport: {
target: 'pino-pretty',
options: {
colorize: true,
translateTime: 'SYS:standard',
ignore: 'pid,hostname',
},
},
};
}
function normalizeS3Prefix(prefix: string | undefined): string {
const value = (prefix || 'openreader').trim();
return value ? value.replace(/^\/+|\/+$/g, '') : 'openreader';
}
function buildS3Client(): S3Client {
const bucket = requireEnv('S3_BUCKET');
const region = requireEnv('S3_REGION');
const accessKeyId = requireEnv('S3_ACCESS_KEY_ID');
const secretAccessKey = requireEnv('S3_SECRET_ACCESS_KEY');
const endpoint = process.env.S3_ENDPOINT?.trim() || undefined;
const forcePathStyle = parseBoolEnv('S3_FORCE_PATH_STYLE', false);
void bucket;
return new S3Client({
region,
endpoint,
forcePathStyle,
requestChecksumCalculation: 'WHEN_REQUIRED',
responseChecksumValidation: 'WHEN_REQUIRED',
credentials: {
accessKeyId,
secretAccessKey,
},
});
}
async function bodyToBuffer(body: unknown): Promise<Buffer> {
if (!body) return Buffer.alloc(0);
if (body instanceof Uint8Array) return Buffer.from(body);
if (ArrayBuffer.isView(body)) return Buffer.from(body.buffer, body.byteOffset, body.byteLength);
if (body instanceof ArrayBuffer) return Buffer.from(body);
if (typeof body === 'object' && body !== null && 'transformToByteArray' in body) {
const maybe = body as { transformToByteArray?: () => Promise<Uint8Array> };
if (typeof maybe.transformToByteArray === 'function') {
return Buffer.from(await maybe.transformToByteArray());
}
}
if (typeof body === 'object' && body !== null && 'on' in body) {
const stream = body as NodeJS.ReadableStream;
const chunks: Buffer[] = [];
for await (const chunk of stream) {
if (Buffer.isBuffer(chunk)) chunks.push(chunk);
else if (typeof chunk === 'string') chunks.push(Buffer.from(chunk));
else chunks.push(Buffer.from(chunk as Uint8Array));
}
return Buffer.concat(chunks);
}
throw new Error('Unsupported S3 response body type');
}
function toArrayBuffer(bytes: Uint8Array): ArrayBuffer {
const copy = new Uint8Array(bytes.byteLength);
copy.set(bytes);
return copy.buffer;
}
async function withTimeout<T>(promise: Promise<T>, timeoutMs: number, label: string): Promise<T> {
let timer: NodeJS.Timeout | null = null;
try {
return await Promise.race([
promise,
new Promise<T>((_, reject) => {
timer = setTimeout(() => reject(new Error(`${label} timed out after ${timeoutMs}ms`)), timeoutMs);
}),
]);
} finally {
if (timer) clearTimeout(timer);
}
}
function parseRetryAfterSeconds(raw: string | undefined): number {
const parsed = Number(raw ?? '');
if (!Number.isFinite(parsed) || parsed <= 0) return 2;
return Math.max(1, Math.floor(parsed));
}
function isAuthed(request: FastifyRequest, expectedToken: string): boolean {
const auth = request.headers.authorization;
if (!auth?.startsWith('Bearer ')) return false;
const token = auth.slice('Bearer '.length).trim();
return token === expectedToken;
}
const alignSchema = z.object({
text: z.string().trim().min(1),
lang: z.string().trim().min(1).max(16).optional(),
cacheKey: z.string().trim().min(1).max(256).optional(),
audioObjectKey: z.string().trim().min(1).max(2048),
});
const layoutSchema = z.object({
documentId: z.string().trim().min(1),
namespace: z.string().trim().min(1).max(128).nullable(),
documentObjectKey: z.string().trim().min(1).max(2048),
});
function mapJobState<Result>(job: Job): WorkerJobStatusResponse<Result> {
if (job.failedReason) {
return {
status: 'failed',
error: {
message: job.failedReason || 'Worker job failed',
},
};
}
if (typeof job.returnvalue !== 'undefined' && job.finishedOn) {
return {
status: 'succeeded',
result: job.returnvalue as Result,
};
}
if (job.processedOn) return { status: 'running' };
return { status: 'queued' };
}
async function getQueueDepth(queue: Queue): Promise<number> {
const counts = await queue.getJobCounts('waiting', 'active', 'prioritized', 'delayed');
return (counts.waiting ?? 0) + (counts.active ?? 0) + (counts.prioritized ?? 0) + (counts.delayed ?? 0);
}
async function main(): Promise<void> {
const port = readIntEnv('COMPUTE_WORKER_PORT', 8081);
const host = process.env.COMPUTE_WORKER_HOST?.trim() || '0.0.0.0';
const workerToken = requireEnv('COMPUTE_WORKER_TOKEN');
const redisUrl = requireEnv('REDIS_URL');
const queueMaxDepth = readIntEnv('COMPUTE_QUEUE_MAX_DEPTH', 64);
const retryAfterSec = parseRetryAfterSeconds(process.env.COMPUTE_QUEUE_RETRY_AFTER_SEC);
const whisperConcurrency = readIntEnv('COMPUTE_WHISPER_CONCURRENCY', 1);
const pdfConcurrency = readIntEnv('COMPUTE_PDF_CONCURRENCY', 2);
const whisperTimeoutMs = readIntEnv('COMPUTE_WHISPER_TIMEOUT_MS', 30_000);
const pdfTimeoutMs = readIntEnv('COMPUTE_PDF_TIMEOUT_MS', 90_000);
const attempts = readIntEnv('COMPUTE_JOB_ATTEMPTS', 2);
const prewarmModels = parseBoolEnv('COMPUTE_PREWARM_MODELS', true);
const redis = new IORedis(redisUrl, {
maxRetriesPerRequest: null,
enableReadyCheck: true,
});
const queueDefaults: JobsOptions = {
attempts,
backoff: {
type: 'exponential',
delay: 500,
},
removeOnComplete: {
age: 60 * 60,
count: 1000,
},
removeOnFail: {
age: 24 * 60 * 60,
count: 5000,
},
};
const alignQueue = new Queue<WhisperAlignJobRequest, WhisperAlignJobResult>(ALIGN_QUEUE_NAME, {
connection: redis,
defaultJobOptions: queueDefaults,
});
const layoutQueue = new Queue<PdfLayoutJobRequest, PdfLayoutJobResult>(PDF_LAYOUT_QUEUE_NAME, {
connection: redis,
defaultJobOptions: queueDefaults,
});
const s3 = buildS3Client();
const s3Bucket = requireEnv('S3_BUCKET');
const s3Prefix = normalizeS3Prefix(process.env.S3_PREFIX);
const ensureSafeKey = (key: string): string => {
const trimmed = key.trim();
if (!trimmed.startsWith(`${s3Prefix}/`)) {
throw new Error('Object key prefix mismatch');
}
return trimmed;
};
const readObjectByKey = async (key: string): Promise<ArrayBuffer> => {
const safeKey = ensureSafeKey(key);
const response = await s3.send(new GetObjectCommand({
Bucket: s3Bucket,
Key: safeKey,
}));
const bytes = await bodyToBuffer(response.Body);
return toArrayBuffer(new Uint8Array(bytes));
};
const alignWorker = new Worker<WhisperAlignJobRequest, WhisperAlignJobResult>(
ALIGN_QUEUE_NAME,
async (job) => {
const parsed = alignSchema.parse(job.data);
const audioBuffer = await readObjectByKey(parsed.audioObjectKey);
return withTimeout(
runWhisperAlignmentFromAudioBuffer({
audioBuffer,
text: parsed.text,
cacheKey: parsed.cacheKey,
lang: parsed.lang,
}),
whisperTimeoutMs,
'whisper alignment job',
);
},
{
connection: redis,
concurrency: whisperConcurrency,
},
);
const layoutWorker = new Worker<PdfLayoutJobRequest, PdfLayoutJobResult>(
PDF_LAYOUT_QUEUE_NAME,
async (job) => {
const parsed = layoutSchema.parse(job.data);
const pdfBytes = await readObjectByKey(parsed.documentObjectKey);
return withTimeout(
runPdfLayoutFromPdfBuffer({
documentId: parsed.documentId,
pdfBytes,
}),
pdfTimeoutMs,
'pdf layout job',
);
},
{
connection: redis,
concurrency: pdfConcurrency,
},
);
alignWorker.on('failed', (job, err) => {
console.error('[compute-worker] align job failed', {
jobId: job?.id,
error: err.message,
});
});
layoutWorker.on('failed', (job, err) => {
console.error('[compute-worker] layout job failed', {
jobId: job?.id,
error: err.message,
});
});
if (prewarmModels) {
await ensureComputeModels();
}
const app = Fastify({
logger: buildLoggerConfig(),
});
app.addHook('onRequest', async (request, reply) => {
const path = request.url.split('?')[0] ?? request.url;
if (path === '/health/live' || path === '/health/ready') return;
if (!isAuthed(request, workerToken)) {
return reply.code(401).send({ error: 'Unauthorized' });
}
return;
});
app.get('/health/live', async () => ({ ok: true }));
app.get('/health/ready', async (_request, reply) => {
try {
await redis.ping();
return { ok: true };
} catch (error) {
reply.code(503);
return {
ok: false,
error: error instanceof Error ? error.message : String(error),
};
}
});
const rejectIfSaturated = async (queue: Queue, reply: FastifyReply): Promise<boolean> => {
const depth = await getQueueDepth(queue);
if (depth < queueMaxDepth) return false;
reply.header('Retry-After', String(retryAfterSec));
reply.code(429).send({
error: 'Queue is saturated',
retryAfterSeconds: retryAfterSec,
queueDepth: depth,
queueMaxDepth,
});
return true;
};
app.post('/align/whisper/jobs', async (request, reply) => {
const parsed = alignSchema.safeParse(request.body);
if (!parsed.success) {
reply.code(400);
return {
error: 'Invalid request body',
issues: parsed.error.issues,
};
}
if (await rejectIfSaturated(alignQueue, reply)) return;
const job = await alignQueue.add('align', parsed.data);
reply.code(202);
return { jobId: String(job.id) };
});
app.get('/align/whisper/jobs/:jobId', async (request, reply) => {
const params = z.object({ jobId: z.string().trim().min(1) }).safeParse(request.params);
if (!params.success) {
reply.code(400);
return { error: 'Invalid job id' };
}
const job = await alignQueue.getJob(params.data.jobId);
if (!job) {
reply.code(404);
return { error: 'Job not found' };
}
return mapJobState<WhisperAlignJobResult>(job);
});
app.post('/layout/pdf/jobs', async (request, reply) => {
const parsed = layoutSchema.safeParse(request.body);
if (!parsed.success) {
reply.code(400);
return {
error: 'Invalid request body',
issues: parsed.error.issues,
};
}
if (await rejectIfSaturated(layoutQueue, reply)) return;
const job = await layoutQueue.add('layout', parsed.data);
reply.code(202);
return { jobId: String(job.id) };
});
app.get('/layout/pdf/jobs/:jobId', async (request, reply) => {
const params = z.object({ jobId: z.string().trim().min(1) }).safeParse(request.params);
if (!params.success) {
reply.code(400);
return { error: 'Invalid job id' };
}
const job = await layoutQueue.getJob(params.data.jobId);
if (!job) {
reply.code(404);
return { error: 'Job not found' };
}
return mapJobState<PdfLayoutJobResult>(job);
});
const close = async (): Promise<void> => {
await app.close();
await Promise.allSettled([
alignWorker.close(),
layoutWorker.close(),
alignQueue.close(),
layoutQueue.close(),
redis.quit(),
]);
};
process.once('SIGINT', () => {
void close().finally(() => process.exit(0));
});
process.once('SIGTERM', () => {
void close().finally(() => process.exit(0));
});
await app.listen({ host, port });
app.log.info({ host, port }, 'compute worker listening');
}
void main().catch((error) => {
console.error('[compute-worker] fatal startup error', error);
process.exit(1);
});

View file

@ -0,0 +1,7 @@
{
"extends": "../../tsconfig.json",
"compilerOptions": {
"noEmit": true
},
"include": ["src/**/*.ts"]
}

View file

@ -0,0 +1,64 @@
---
title: Compute Worker (Redis + BullMQ)
---
Use this guide for `COMPUTE_MODE=worker` deployments where heavy compute runs outside the Next.js app server.
## Overview
The compute worker handles:
- Whisper word alignment (`/align/whisper/jobs`)
- PDF layout parsing (`/layout/pdf/jobs`)
The app server enqueues jobs and polls status. Queue durability and retries are backed by Redis + BullMQ.
## Published image
- App server image: `ghcr.io/richardr1126/openreader`
- Compute worker image: `ghcr.io/richardr1126/openreader-compute-worker`
## Worker environment variables
Required:
- `COMPUTE_WORKER_TOKEN`: bearer token expected by worker routes
- `REDIS_URL`: BullMQ Redis connection string
- `S3_BUCKET`
- `S3_REGION`
- `S3_ACCESS_KEY_ID`
- `S3_SECRET_ACCESS_KEY`
Common optional:
- `S3_ENDPOINT` (for non-AWS S3-compatible storage)
- `S3_FORCE_PATH_STYLE=true` (for many S3-compatible providers)
- `S3_PREFIX=openreader`
- `COMPUTE_WORKER_HOST=0.0.0.0`
- `COMPUTE_WORKER_PORT=8081`
- `COMPUTE_LOG_FORMAT=pretty` (default) or `json`
- `COMPUTE_QUEUE_MAX_DEPTH=64`
- `COMPUTE_PREWARM_MODELS=true`
## App server environment variables (worker mode)
Set on the Next.js app server:
```env
COMPUTE_MODE=worker
COMPUTE_WORKER_URL=http://<worker-host>:8081
COMPUTE_WORKER_TOKEN=<same-token-as-worker>
```
`COMPUTE_MODE=worker` has no local fallback. If worker is unavailable, affected requests fail.
## Production notes
- Worker mode assumes shared object storage is reachable by both app server and worker.
- Non-exposed embedded `weed mini` is not supported with external worker mode.
- Protect `COMPUTE_WORKER_TOKEN` and avoid exposing worker routes publicly without auth.
## Health endpoints
- `GET /health/live`
- `GET /health/ready`

View file

@ -124,6 +124,40 @@ If you need mirrors or pinned artifact locations, set `WHISPER_MODEL_BASE_URL` i
</details>
<details>
<summary><strong>External compute worker dev stack (optional)</strong></summary>
Use this when you want durable compute with Redis/BullMQ while keeping Next.js on native host `pnpm dev`.
Full worker deployment details are in [Compute Worker (Redis + BullMQ)](./compute-worker).
Start only Redis + compute-worker via compose watch:
```bash
docker compose --env-file compute/worker/.env -f compute/worker/docker-compose.yml up --watch
# or: pnpm compute:dev:watch
```
`compute/worker/.env.example` contains a starter config. Copy it to `compute/worker/.env` and adjust values for your environment.
Run the main app separately on the host:
```bash
pnpm dev
```
For app server worker mode, set:
```env
COMPUTE_MODE=worker
COMPUTE_WORKER_URL=http://localhost:8081
COMPUTE_WORKER_TOKEN=<same-token-used-by-worker>
```
Worker mode requires worker-reachable shared object storage (S3-compatible endpoint).
Non-exposed embedded `weed mini` is not a supported topology for external worker mode.
</details>
## Steps
### Required flow
@ -206,6 +240,25 @@ S3_SECRET_ACCESS_KEY=your-secret-key
# Optional for non-AWS providers:
# S3_ENDPOINT=https://your-s3-compatible-endpoint
# S3_FORCE_PATH_STYLE=true
```
</TabItem>
<TabItem value="worker-mode" label="External Worker + Redis">
```env
API_BASE=http://host.docker.internal:8880/v1
API_KEY=none
COMPUTE_MODE=worker
COMPUTE_WORKER_URL=http://localhost:8081
COMPUTE_WORKER_TOKEN=<same-token-used-by-worker>
USE_EMBEDDED_WEED_MINI=false
S3_BUCKET=your-bucket
S3_REGION=us-east-1
S3_ACCESS_KEY_ID=your-access-key
S3_SECRET_ACCESS_KEY=your-secret-key
# Optional for non-AWS providers:
# S3_ENDPOINT=https://your-s3-compatible-endpoint
# S3_FORCE_PATH_STYLE=true
```
</TabItem>

View file

@ -8,6 +8,8 @@ This guide covers deploying OpenReader to Vercel with external Postgres and S3-c
- Documents (PDF/EPUB/TXT/MD) work with `POSTGRES_URL` + external S3 storage.
- Audiobook routes work on Node.js serverless functions using `ffmpeg-static`.
- Heavy compute features (Whisper alignment + PDF layout parsing) work through `COMPUTE_MODE=worker` with an external compute worker service.
- For worker setup details and worker-specific env vars, see [Compute Worker (Redis + BullMQ)](./compute-worker).
:::warning DOCX Conversion Limitation
`docx` conversion requires `soffice` (LibreOffice), which is not available in a standard Vercel runtime.
@ -36,12 +38,15 @@ AUTH_SECRET=...
ADMIN_EMAILS=you@example.com # comma-separated; admins manage TTS + features in-app
# Heavy compute (recommended on Vercel in v1)
# local = requires native binaries/models in-process
# local = requires native binaries/models in-process (not recommended on Vercel)
# worker = external durable compute worker (recommended)
# none = disable ONNX whisper alignment + PDF layout parsing
COMPUTE_MODE=none
COMPUTE_MODE=worker
COMPUTE_WORKER_URL=https://your-compute-worker.example.com
COMPUTE_WORKER_TOKEN=...
# First-boot seed for the TTS shared provider (optional; manage in-app afterwards)
API_KEY=your_replicate_key
# API_KEY=your_replicate_key
# API_BASE only needed for OpenAI-compatible self-hosted providers
```
@ -129,4 +134,4 @@ Adjust memory per route if your files are larger or your plan differs.
1. Upload and read a PDF/EPUB document.
2. Confirm sync/blob fetch works across refreshes/devices.
3. Generate at least one audiobook chapter and play/download it.
4. If you run with local compute (`COMPUTE_MODE=local`) outside Vercel, verify word highlighting timestamps on a TTS run.
4. Verify worker-backed word highlighting and PDF parsing in `COMPUTE_MODE=worker`.

View file

@ -22,6 +22,12 @@ OpenReader currently pins embedded SeaweedFS to `4.18` in CI and Docker builds.
`4.19` introduced intermittent `InternalError` responses on S3 `PutObject` in our upload flow.
:::
## Published images
- App server: `ghcr.io/richardr1126/openreader:latest`
- Compute worker (Optional): `ghcr.io/richardr1126/openreader-compute-worker:latest`
- Legacy app alias: `ghcr.io/richardr1126/openreader-webui:latest`
## 1. Start the Docker container
<Tabs groupId="docker-start-mode">
@ -135,6 +141,7 @@ Visit [http://localhost:3003](http://localhost:3003) after startup.
## 3. Update Docker image
Legacy image compatibility: `ghcr.io/richardr1126/openreader-webui:latest` remains available as an alias.
For external compute mode image details, see [Compute Worker (Redis + BullMQ)](./deploy/compute-worker).
```bash
docker stop openreader || true && \

View file

@ -54,8 +54,8 @@ For auth-enabled deployments, use **Settings → Admin** as the primary source o
| `IMPORT_LIBRARY_DIR` | Library import | `docstore/library` fallback | Set a single server library root |
| `IMPORT_LIBRARY_DIRS` | Library import | unset | Set multiple roots (comma/colon/semicolon separated) |
| `COMPUTE_MODE` | Heavy compute backend | `local` | Set to `none` to disable ONNX word alignment + PDF layout parsing |
| `COMPUTE_WORKER_URL` | Heavy compute backend | unset | Reserved for future worker backend mode (`worker`) |
| `COMPUTE_WORKER_TOKEN` | Heavy compute backend | unset | Reserved for future worker backend mode (`worker`) |
| `COMPUTE_WORKER_URL` | Heavy compute backend | unset | Required when `COMPUTE_MODE=worker`; base URL for external compute worker |
| `COMPUTE_WORKER_TOKEN` | Heavy compute backend | unset | Required bearer token for external compute worker auth |
| `PDF_LAYOUT_MODEL_BASE_URL` | PDF layout model | PP-DocLayoutV3 ONNX base URL | Optional base URL override for `ensureModel()` |
| `WHISPER_MODEL_BASE_URL` | Whisper ONNX model | onnx-community defaults | Optional base URL override for ONNX whisper-base_timestamped int8 downloads |
| `FFMPEG_BIN` | Audio runtime | auto-detected (`ffmpeg-static`) | Override ffmpeg binary path |
@ -357,22 +357,26 @@ Selects the backend for heavy compute features (ONNX word alignment + PDF layout
- Default: `local`
- Supported in v1:
- `local`: run compute in-process on the app server
- `worker`: enqueue async jobs in an external durable compute worker (Redis + BullMQ)
- `none`: disable these features cleanly
- `worker` is reserved for a future external worker backend and currently fails fast at startup if selected
- `worker` requires `COMPUTE_WORKER_URL` and `COMPUTE_WORKER_TOKEN`
- `worker` assumes the external worker can directly reach shared object storage (S3-compatible endpoint)
- `worker` is not compatible with non-exposed embedded `weed mini` storage topologies
- Worker service env vars are documented in [Compute Worker (Redis + BullMQ)](../deploy/compute-worker)
### COMPUTE_WORKER_URL
Reserved for future external compute worker mode.
Base URL for external compute worker mode.
- Used only when `COMPUTE_MODE=worker` (not implemented in v1)
- Leave unset in v1
- Used only when `COMPUTE_MODE=worker`
- Example: `http://localhost:8081`
### COMPUTE_WORKER_TOKEN
Reserved bearer token for future external compute worker mode.
Bearer token for external compute worker auth.
- Used only when `COMPUTE_MODE=worker` (not implemented in v1)
- Leave unset in v1
- Used only when `COMPUTE_MODE=worker`
- Must match worker service `COMPUTE_WORKER_TOKEN`
### PDF_LAYOUT_MODEL_BASE_URL

View file

@ -4,9 +4,10 @@ title: Stack
## Framework
- [Next.js](https://nextjs.org/) 15 (App Router)
- [Next.js](https://nextjs.org/) 15 (App Router, Turbopack in dev)
- [React](https://react.dev/) 19
- [TypeScript](https://www.typescriptlang.org/)
- [pnpm](https://pnpm.io/) workspaces monorepo
## Containerization and runtime
@ -17,24 +18,48 @@ title: Stack
- UI: [Tailwind CSS](https://tailwindcss.com), [Headless UI](https://headlessui.com), [@tailwindcss/typography](https://tailwindcss.com/docs/typography-plugin)
- Interactions: `react-dnd`, `react-dropzone`
- Server state: [TanStack Query](https://tanstack.com/query) (React Query v5)
- Authentication: [Better Auth](https://www.better-auth.com/) client SDK
- Local storage/cache: [Dexie.js](https://dexie.org/) (IndexedDB)
- Audio playback: [Howler.js](https://howlerjs.com/)
- Notifications: `react-hot-toast`
- Document rendering:
- PDF: [react-pdf](https://github.com/wojtekmaj/react-pdf), [pdf.js](https://mozilla.github.io/pdf.js/)
- EPUB: [react-reader](https://github.com/gerhardsletten/react-reader), [epubjs](https://github.com/futurepress/epub.js/)
- Markdown/Text: [react-markdown](https://github.com/remarkjs/react-markdown), [remark-gfm](https://github.com/remarkjs/remark-gfm)
- Text preprocessing/matching: [compromise](https://github.com/spencermountain/compromise), [cmpstr](https://github.com/remsky/cmpstr)
- Analytics: [Vercel Analytics](https://vercel.com/analytics)
## Next.js server
- APIs: Route Handlers for sync, blob/content access, migrations, audiobook export, TTS/Whisper proxying
- State sync: request-based today (not realtime push updates)
- Authentication: [Better Auth](https://www.better-auth.com/) server handlers/adapters
- Authentication: [Better Auth](https://www.better-auth.com/) server handlers/adapters with anonymous session support
- Metadata DB: [Drizzle ORM](https://orm.drizzle.team/) with SQLite (`better-sqlite3`) by default and optional Postgres (`pg`)
- App tables are manually maintained in Drizzle schema files
- Auth tables are auto-generated by the [Better Auth](https://www.better-auth.com/) CLI and migrated alongside app tables via Drizzle
- Blob storage: embedded [SeaweedFS](https://github.com/seaweedfs/seaweedfs) (`weed mini`) by default, or external S3-compatible storage via AWS SDK v3
- Audio/processing pipeline: OpenAI-compatible TTS providers, [ffmpeg](https://ffmpeg.org/) for audiobook assembly, built-in ONNX Whisper (`onnx-community/whisper-base_timestamped` int8) for word timestamps
- TTS providers: OpenAI-compatible API (`openai` SDK), [Replicate](https://replicate.com/) (`replicate` client), DeepInfra, and custom OpenAI-compatible endpoints — credentials are encrypted at rest
- Audio pipeline: [ffmpeg](https://ffmpeg.org/) (`ffmpeg-static`) for audiobook assembly, `archiver` for export packaging
- Utilities: `lru-cache` for in-process caching, `fast-xml-parser` for EPUB/XML parsing, `uuid` for identifier generation, `zod` for schema validation
## External compute worker (optional)
Monorepo packages under `compute/`:
- **`@openreader/compute-core`** — ONNX runtime lifecycle, model management, and inference logic shared between local and worker modes
- ONNX runtime: `onnxruntime-node` with `@huggingface/tokenizers`
- Whisper alignment: `onnx-community/whisper-base_timestamped` (int8) for word-level timestamps
- PDF layout: `Bei0001/PP-DocLayoutV3-ONNX` for document block detection and layout parsing
- PDF rendering: `pdfjs-dist`, `@napi-rs/canvas` for server-side page rasterization
- Utilities: `jszip`, `ffmpeg-static`
- **`@openreader/compute-worker`** — standalone Node.js worker service
- HTTP server: [Fastify](https://fastify.dev/) v5
- Job queue: [BullMQ](https://bullmq.io/) + [ioredis](https://github.com/redis/ioredis) (queues: `whisper-align`, `pdf-layout`)
- Storage: AWS SDK v3 S3 client for reading/writing blobs
- Logging: [Pino](https://getpino.io/)
- Validation: [Zod](https://zod.dev/)
- Compute mode is controlled by `COMPUTE_MODE` env var: `local` (in-process), `worker` (remote queue via HTTP + Redis), or disabled
## Tooling and testing
@ -42,3 +67,4 @@ title: Stack
- TypeScript
- [Playwright](https://playwright.dev/) end-to-end tests
- Drizzle migration/generation scripts
- [Docusaurus](https://docusaurus.io/) documentation site (`docs-site/`)

View file

@ -53,7 +53,7 @@ const sidebars: SidebarsConfig = {
{
type: 'category',
label: '🚀 Deploy',
items: ['deploy/local-development', 'deploy/vercel-deployment'],
items: ['deploy/local-development', 'deploy/compute-worker', 'deploy/vercel-deployment'],
},
{
type: 'category',

View file

@ -15,13 +15,16 @@ const securityHeaders = [
},
];
const computeMode = (process.env.COMPUTE_MODE || 'local').trim().toLowerCase();
const computeDisabled = computeMode === 'none';
const computeModeRaw = (process.env.COMPUTE_MODE || 'local').trim().toLowerCase();
const computeMode = computeModeRaw === 'none' || computeModeRaw === 'worker' || computeModeRaw === 'local'
? computeModeRaw
: 'local';
const computeLocal = computeMode === 'local';
const serverExternalPackages = [
'@napi-rs/canvas',
'ffmpeg-static',
'better-sqlite3',
...(computeDisabled ? [] : ['onnxruntime-node', '@huggingface/tokenizers']),
...(computeLocal ? ['onnxruntime-node', '@huggingface/tokenizers'] : []),
];
const nextConfig: NextConfig = {
@ -39,6 +42,7 @@ const nextConfig: NextConfig = {
canvas: '@napi-rs/canvas',
},
},
transpilePackages: ['@openreader/compute-core'],
serverExternalPackages,
outputFileTracingIncludes: {
'/api/audiobook': [
@ -60,7 +64,7 @@ const nextConfig: NextConfig = {
'./node_modules/pdfjs-dist/legacy/build/pdf.worker.mjs',
],
},
...(computeDisabled
...(!computeLocal
? {
outputFileTracingExcludes: {
'/*': [

View file

@ -20,7 +20,11 @@
"docs:build": "pnpm --dir docs-site build",
"docs:serve": "pnpm --dir docs-site serve",
"docs:version": "pnpm --dir docs-site docusaurus docs:version",
"docs:clear": "pnpm --dir docs-site clear"
"docs:clear": "pnpm --dir docs-site clear",
"compute:worker:dev": "pnpm --filter @openreader/compute-worker dev",
"compute:worker:start": "pnpm --filter @openreader/compute-worker start",
"compute:dev:compose": "docker compose --env-file compute/worker/.env -f compute/worker/docker-compose.yml up --build",
"compute:dev:watch": "docker compose --env-file compute/worker/.env -f compute/worker/docker-compose.yml up --watch"
},
"dependencies": {
"@aws-sdk/client-s3": "^3.1045.0",

File diff suppressed because it is too large Load diff

View file

@ -1,5 +1,6 @@
packages:
- .
- compute/*
allowBuilds:
'@napi-rs/canvas': true

View file

@ -372,10 +372,8 @@ export async function POST(request: NextRequest) {
// previously unavailable, retry alignment using the current segment text.
if (!alignment) {
try {
const audioBuffer = await getTtsSegmentAudioObject(existing.audioKey);
const whisperBytes = Uint8Array.from(audioBuffer);
const aligned = (await getCompute().alignWords({
audioBuffer: whisperBytes.buffer,
audioObjectKey: existing.audioKey,
text: segment.text,
})).alignments;
alignment = aligned[0] ? { ...aligned[0], sentenceIndex: segment.original.segmentIndex } : null;
@ -527,6 +525,7 @@ export async function POST(request: NextRequest) {
const whisperBytes = Uint8Array.from(persistedBuffer);
const aligned = (await getCompute().alignWords({
audioBuffer: whisperBytes.buffer,
audioObjectKey: audioKey,
text: segment.text,
})).alignments;
alignment = aligned[0] ? { ...aligned[0], sentenceIndex: segment.original.segmentIndex } : null;

View file

@ -1,17 +1,14 @@
import type { ComputeBackend, ComputeMode } from '@/lib/server/compute/types';
import { NoneComputeBackend } from '@/lib/server/compute/none';
import { isComputeModeAvailable, readComputeMode } from '@/lib/server/compute/mode';
import { WorkerComputeBackend } from '@/lib/server/compute/worker';
let backend: ComputeBackend | null = null;
function createBackend(): ComputeBackend {
const mode: ComputeMode = readComputeMode();
if (mode === 'none') return new NoneComputeBackend();
if (mode === 'worker') {
throw new Error(
'COMPUTE_MODE=worker is not implemented yet in v1. Switch to local/none or implement WorkerComputeBackend (v2).',
);
}
if (mode === 'worker') return new WorkerComputeBackend();
// Intentionally lazy-load local compute so COMPUTE_MODE=none builds
// can avoid tracing heavy ONNX dependencies.
// eslint-disable-next-line @typescript-eslint/no-require-imports

View file

@ -1,21 +1,40 @@
import type { ComputeBackend, PdfLayoutInput, WhisperAlignInput, WhisperAlignResult } from '@/lib/server/compute/types';
import { getDocumentBlob } from '@/lib/server/documents/blobstore';
import { getTtsSegmentAudioObject } from '@/lib/server/tts/segments-blobstore';
import {
runPdfLayoutFromPdfBuffer,
runWhisperAlignmentFromAudioBuffer,
} from '@openreader/compute-core';
export class LocalComputeBackend implements ComputeBackend {
readonly mode = 'local' as const;
async alignWords(input: WhisperAlignInput): Promise<WhisperAlignResult> {
const { alignAudioWithText } = await import('@/lib/server/whisper/alignment');
const alignments = await alignAudioWithText(
input.audioBuffer,
input.text,
input.cacheKey,
{ lang: input.lang },
);
return { alignments };
let audioBuffer = input.audioBuffer ?? null;
if (!audioBuffer && input.audioObjectKey) {
const bytes = new Uint8Array(await getTtsSegmentAudioObject(input.audioObjectKey));
audioBuffer = bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength);
}
if (!audioBuffer) {
throw new Error('Local compute alignment requires audioBuffer or audioObjectKey');
}
return runWhisperAlignmentFromAudioBuffer({
audioBuffer,
text: input.text,
cacheKey: input.cacheKey,
lang: input.lang,
});
}
async parsePdfLayout(input: PdfLayoutInput) {
const { parsePdf } = await import('@/lib/server/pdf-layout/parsePdf');
return parsePdf({ documentId: input.documentId, pdfBytes: input.pdfBytes });
let pdfBytes = input.pdfBytes ?? null;
if (!pdfBytes && input.documentId && typeof input.namespace !== 'undefined') {
const bytes = new Uint8Array(await getDocumentBlob(input.documentId, input.namespace));
pdfBytes = bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength);
}
if (!pdfBytes) {
throw new Error('Local compute PDF layout requires pdfBytes or (documentId + namespace)');
}
return (await runPdfLayoutFromPdfBuffer({ documentId: input.documentId, pdfBytes })).parsed;
}
}

View file

@ -7,10 +7,5 @@ export function readComputeMode(): ComputeMode {
}
export function isComputeModeAvailable(mode: ComputeMode): boolean {
if (mode === 'worker') {
throw new Error(
'COMPUTE_MODE=worker is not implemented yet in v1. Switch to local/none or implement WorkerComputeBackend (v2).',
);
}
return mode !== 'none';
}

View file

@ -1,10 +1,10 @@
import type { TTSAudioBuffer, TTSSentenceAlignment } from '@/types/tts';
import type { ParsedPdfDocument } from '@/types/parsed-pdf';
import type { TTSAudioBuffer, TTSSentenceAlignment, ParsedPdfDocument } from '@openreader/compute-core';
export type ComputeMode = 'local' | 'worker' | 'none';
export interface WhisperAlignInput {
audioBuffer: TTSAudioBuffer;
audioBuffer?: TTSAudioBuffer;
audioObjectKey?: string;
text: string;
cacheKey?: string;
lang?: string;
@ -16,7 +16,9 @@ export interface WhisperAlignResult {
export interface PdfLayoutInput {
documentId: string;
pdfBytes: ArrayBuffer;
namespace?: string | null;
documentObjectKey?: string;
pdfBytes?: ArrayBuffer;
}
export interface ComputeBackend {

View file

@ -0,0 +1,11 @@
export {
ALIGN_QUEUE_NAME,
PDF_LAYOUT_QUEUE_NAME,
type PdfLayoutJobRequest,
type PdfLayoutJobResult,
type WhisperAlignJobRequest,
type WhisperAlignJobResult,
type WorkerJobErrorShape,
type WorkerJobState,
type WorkerJobStatusResponse,
} from '@openreader/compute-core';

View file

@ -0,0 +1,165 @@
import type { ComputeBackend, PdfLayoutInput, WhisperAlignInput, WhisperAlignResult } from '@/lib/server/compute/types';
import type {
PdfLayoutJobRequest,
PdfLayoutJobResult,
WhisperAlignJobRequest,
WhisperAlignJobResult,
WorkerJobStatusResponse,
} from '@/lib/server/compute/worker-contract';
class WorkerHttpError extends Error {
status: number;
retryAfterMs: number | null;
constructor(message: string, status: number, retryAfterMs: number | null = null) {
super(message);
this.name = 'WorkerHttpError';
this.status = status;
this.retryAfterMs = retryAfterMs;
}
}
const DEFAULT_WAIT_TIMEOUT_MS = 45_000;
const DEFAULT_RETRIES = 2;
const POLL_INTERVAL_MS = 400;
const POLL_MAX_INTERVAL_MS = 1_500;
function readRequiredEnv(name: string): string {
const value = process.env[name]?.trim();
if (!value) throw new Error(`${name} is required when COMPUTE_MODE=worker`);
return value;
}
function parseRetryAfterMs(value: string | null): number | null {
if (!value) return null;
const asNum = Number(value);
if (Number.isFinite(asNum)) {
return Math.max(0, Math.floor(asNum * 1000));
}
const when = Date.parse(value);
if (Number.isNaN(when)) return null;
return Math.max(0, when - Date.now());
}
function sleep(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
function shouldRetry(error: unknown): boolean {
if (error instanceof WorkerHttpError) {
return error.status === 429 || error.status === 502 || error.status === 503 || error.status === 504;
}
if (error instanceof Error) {
const msg = error.message.toLowerCase();
return msg.includes('network') || msg.includes('timeout') || msg.includes('fetch failed');
}
return false;
}
async function withRetries<T>(attempts: number, operation: () => Promise<T>): Promise<T> {
let lastError: unknown = null;
for (let attempt = 0; attempt < attempts; attempt += 1) {
try {
return await operation();
} catch (error) {
lastError = error;
if (attempt === attempts - 1 || !shouldRetry(error)) break;
if (error instanceof WorkerHttpError && typeof error.retryAfterMs === 'number') {
await sleep(error.retryAfterMs);
} else {
await sleep((attempt + 1) * 250);
}
}
}
throw lastError instanceof Error ? lastError : new Error('Unknown worker compute failure');
}
export class WorkerComputeBackend implements ComputeBackend {
readonly mode = 'worker' as const;
private readonly baseUrl: string;
private readonly token: string;
private readonly waitTimeoutMs: number;
private readonly retries: number;
constructor() {
this.baseUrl = readRequiredEnv('COMPUTE_WORKER_URL').replace(/\/+$/, '');
this.token = readRequiredEnv('COMPUTE_WORKER_TOKEN');
this.waitTimeoutMs = DEFAULT_WAIT_TIMEOUT_MS;
this.retries = DEFAULT_RETRIES;
}
async alignWords(input: WhisperAlignInput): Promise<WhisperAlignResult> {
if (!input.audioObjectKey) {
throw new Error('Worker compute alignment requires audioObjectKey');
}
const payload: WhisperAlignJobRequest = {
text: input.text,
lang: input.lang,
cacheKey: input.cacheKey,
audioObjectKey: input.audioObjectKey,
};
return withRetries(this.retries, async () => {
const { jobId } = await this.requestJson<{ jobId: string }>('POST', '/align/whisper/jobs', payload);
const status = await this.waitForJob<WhisperAlignJobResult>(`/align/whisper/jobs/${encodeURIComponent(jobId)}`);
if (status.status !== 'succeeded' || !status.result) {
throw new Error(status.error?.message || 'Whisper worker job did not complete');
}
return { alignments: status.result.alignments };
});
}
async parsePdfLayout(input: PdfLayoutInput) {
if (!input.documentObjectKey) {
throw new Error('Worker compute PDF layout requires documentObjectKey');
}
const payload: PdfLayoutJobRequest = {
documentId: input.documentId,
namespace: input.namespace ?? null,
documentObjectKey: input.documentObjectKey,
};
return withRetries(this.retries, async () => {
const { jobId } = await this.requestJson<{ jobId: string }>('POST', '/layout/pdf/jobs', payload);
const status = await this.waitForJob<PdfLayoutJobResult>(`/layout/pdf/jobs/${encodeURIComponent(jobId)}`);
if (status.status !== 'succeeded' || !status.result) {
throw new Error(status.error?.message || 'PDF layout worker job did not complete');
}
return status.result.parsed;
});
}
private async requestJson<T>(method: 'GET' | 'POST', path: string, body?: unknown): Promise<T> {
const res = await fetch(`${this.baseUrl}${path}`, {
method,
headers: {
Authorization: `Bearer ${this.token}`,
...(method === 'POST' ? { 'Content-Type': 'application/json' } : {}),
},
...(method === 'POST' ? { body: JSON.stringify(body ?? {}) } : {}),
});
if (!res.ok) {
const retryAfterMs = parseRetryAfterMs(res.headers.get('retry-after'));
const detail = await res.text().catch(() => '');
throw new WorkerHttpError(
`Worker request failed (${method} ${path}): ${res.status}${detail ? ` ${detail}` : ''}`,
res.status,
retryAfterMs,
);
}
return res.json() as Promise<T>;
}
private async waitForJob<Result>(path: string): Promise<WorkerJobStatusResponse<Result>> {
const started = Date.now();
let interval = POLL_INTERVAL_MS;
while ((Date.now() - started) < this.waitTimeoutMs) {
const status = await this.requestJson<WorkerJobStatusResponse<Result>>('GET', path);
if (status.status === 'succeeded' || status.status === 'failed') return status;
await sleep(interval);
interval = Math.min(POLL_MAX_INTERVAL_MS, Math.floor(interval * 1.5));
}
throw new Error(`Timed out waiting for worker job after ${this.waitTimeoutMs}ms`);
}
}

View file

@ -2,7 +2,7 @@ import { and, eq } from 'drizzle-orm';
import { db } from '@/db';
import { documents } from '@/db/schema';
import { UnsupportedComputeError } from '@/lib/server/compute/types';
import { getDocumentBlob, putParsedDocumentBlob } from '@/lib/server/documents/blobstore';
import { documentKey, putParsedDocumentBlob } from '@/lib/server/documents/blobstore';
import { getCompute } from '@/lib/server/compute';
import { clearTtsSegmentCache } from '@/lib/server/tts/segments-cache';
@ -29,10 +29,10 @@ export async function parsePdfJob(input: ParsePdfJobInput): Promise<void> {
.set({ parseStatus: 'running' })
.where(and(eq(documents.id, input.documentId), eq(documents.userId, input.userId)));
const pdfBytes = await getDocumentBlob(input.documentId, input.namespace);
const parsed = await getCompute().parsePdfLayout({
documentId: input.documentId,
pdfBytes: new Uint8Array(pdfBytes).buffer,
namespace: input.namespace,
documentObjectKey: documentKey(input.documentId, input.namespace),
});
const parsedJson = Buffer.from(JSON.stringify(parsed));

View file

@ -19,7 +19,8 @@
}
],
"paths": {
"@/*": ["./src/*"]
"@/*": ["./src/*"],
"@openreader/compute-core": ["./compute/core/src/index.ts"]
}
},
"include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts"],