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:
parent
3a21f2a5f5
commit
f1aa1c3e3b
53 changed files with 4880 additions and 96 deletions
|
|
@ -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
|
||||
|
|
|
|||
126
.github/workflows/docker-publish.yml
vendored
126
.github/workflows/docker-publish.yml
vendored
|
|
@ -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 }}"
|
||||
|
|
|
|||
|
|
@ -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
17
compute/core/package.json
Normal 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
87
compute/core/src/index.ts
Normal 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 };
|
||||
}
|
||||
132
compute/core/src/pdf-layout/ensureModel.ts
Normal file
132
compute/core/src/pdf-layout/ensureModel.ts
Normal 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;
|
||||
}
|
||||
153
compute/core/src/pdf-layout/mergeTextWithRegions.ts
Normal file
153
compute/core/src/pdf-layout/mergeTextWithRegions.ts
Normal 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;
|
||||
}
|
||||
3
compute/core/src/pdf-layout/model/LICENSE.txt
Normal file
3
compute/core/src/pdf-layout/model/LICENSE.txt
Normal 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
|
||||
31
compute/core/src/pdf-layout/model/manifest.json
Normal file
31
compute/core/src/pdf-layout/model/manifest.json
Normal 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
|
||||
}
|
||||
]
|
||||
}
|
||||
165
compute/core/src/pdf-layout/parsePdf.ts
Normal file
165
compute/core/src/pdf-layout/parsePdf.ts
Normal 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);
|
||||
}
|
||||
}
|
||||
162
compute/core/src/pdf-layout/renderPage.ts
Normal file
162
compute/core/src/pdf-layout/renderPage.ts
Normal 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);
|
||||
}
|
||||
}
|
||||
331
compute/core/src/pdf-layout/runLayoutModel.ts
Normal file
331
compute/core/src/pdf-layout/runLayoutModel.ts
Normal 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)}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
144
compute/core/src/pdf-layout/stitchCrossPageBlocks.ts
Normal file
144
compute/core/src/pdf-layout/stitchCrossPageBlocks.ts
Normal 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,
|
||||
};
|
||||
}
|
||||
15
compute/core/src/pdf-layout/types.ts
Normal file
15
compute/core/src/pdf-layout/types.ts
Normal 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;
|
||||
}
|
||||
7
compute/core/src/runtime/docstore.ts
Normal file
7
compute/core/src/runtime/docstore.ts
Normal 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;
|
||||
}
|
||||
36
compute/core/src/runtime/ffmpeg.ts
Normal file
36
compute/core/src/runtime/ffmpeg.ts
Normal 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',
|
||||
);
|
||||
}
|
||||
53
compute/core/src/types/parsed-pdf.ts
Normal file
53
compute/core/src/types/parsed-pdf.ts
Normal 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[];
|
||||
}
|
||||
16
compute/core/src/types/tts.ts
Normal file
16
compute/core/src/types/tts.ts
Normal 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[];
|
||||
}
|
||||
54
compute/core/src/whisper/alignment-mapping.ts
Normal file
54
compute/core/src/whisper/alignment-mapping.ts
Normal 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,
|
||||
};
|
||||
}
|
||||
1012
compute/core/src/whisper/alignment.ts
Normal file
1012
compute/core/src/whisper/alignment.ts
Normal file
File diff suppressed because it is too large
Load diff
242
compute/core/src/whisper/ensureModel.ts
Normal file
242
compute/core/src/whisper/ensureModel.ts
Normal 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();
|
||||
}
|
||||
21
compute/core/src/whisper/model/LICENSE.txt
Normal file
21
compute/core/src/whisper/model/LICENSE.txt
Normal 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.
|
||||
76
compute/core/src/whisper/model/manifest.json
Normal file
76
compute/core/src/whisper/model/manifest.json
Normal 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
|
||||
}
|
||||
]
|
||||
}
|
||||
BIN
compute/core/src/whisper/model/mel_filters.npz
Normal file
BIN
compute/core/src/whisper/model/mel_filters.npz
Normal file
Binary file not shown.
21
compute/core/src/whisper/spectral.ts
Normal file
21
compute/core/src/whisper/spectral.ts
Normal 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;
|
||||
}
|
||||
449
compute/core/src/whisper/token-timestamps.ts
Normal file
449
compute/core/src/whisper/token-timestamps.ts
Normal 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);
|
||||
}
|
||||
7
compute/core/tsconfig.json
Normal file
7
compute/core/tsconfig.json
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
{
|
||||
"extends": "../../tsconfig.json",
|
||||
"compilerOptions": {
|
||||
"noEmit": true
|
||||
},
|
||||
"include": ["src/**/*.ts"]
|
||||
}
|
||||
25
compute/worker/.env.example
Normal file
25
compute/worker/.env.example
Normal 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
20
compute/worker/Dockerfile
Normal 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"]
|
||||
41
compute/worker/docker-compose.yml
Normal file
41
compute/worker/docker-compose.yml
Normal 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
|
||||
23
compute/worker/package.json
Normal file
23
compute/worker/package.json
Normal 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"
|
||||
}
|
||||
}
|
||||
432
compute/worker/src/server.ts
Normal file
432
compute/worker/src/server.ts
Normal 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);
|
||||
});
|
||||
7
compute/worker/tsconfig.json
Normal file
7
compute/worker/tsconfig.json
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
{
|
||||
"extends": "../../tsconfig.json",
|
||||
"compilerOptions": {
|
||||
"noEmit": true
|
||||
},
|
||||
"include": ["src/**/*.ts"]
|
||||
}
|
||||
64
docs-site/docs/deploy/compute-worker.md
Normal file
64
docs-site/docs/deploy/compute-worker.md
Normal 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`
|
||||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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`.
|
||||
|
|
|
|||
|
|
@ -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 && \
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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/`)
|
||||
|
|
|
|||
|
|
@ -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',
|
||||
|
|
|
|||
|
|
@ -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: {
|
||||
'/*': [
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
593
pnpm-lock.yaml
593
pnpm-lock.yaml
File diff suppressed because it is too large
Load diff
|
|
@ -1,5 +1,6 @@
|
|||
packages:
|
||||
- .
|
||||
- compute/*
|
||||
|
||||
allowBuilds:
|
||||
'@napi-rs/canvas': true
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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';
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
11
src/lib/server/compute/worker-contract.ts
Normal file
11
src/lib/server/compute/worker-contract.ts
Normal 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';
|
||||
165
src/lib/server/compute/worker.ts
Normal file
165
src/lib/server/compute/worker.ts
Normal 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`);
|
||||
}
|
||||
}
|
||||
|
|
@ -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));
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
Loading…
Reference in a new issue