From 6e69b481034aa7eb58b88ee7a59697e2cdac78b4 Mon Sep 17 00:00:00 2001 From: Richard R Date: Fri, 12 Jun 2026 14:16:13 -0600 Subject: [PATCH] chore(worker): restructure inference modules and centralize config - Remove legacy inference and platform files from `src/inference/` and migrate responsibilities to new modular locations under `src/infrastructure/`, `src/inference/pdf/`, and `src/inference/whisper/` - Consolidate environment/config logic into `src/infrastructure/config.ts` - Move docstore and ffmpeg platform utilities to `src/infrastructure/platform.ts` - Refactor PDF and Whisper inference code to use new document layout, layout model, and timestamp utilities - Update API contracts to define and export types and constants previously scattered in inference/types - Adjust all imports in jobs, storage, and tests to reference new module structure and type locations - Inline parser version and encoding logic into API contracts - Remove obsolete files and update test fixtures for new type locations This update improves codebase modularity, maintainability, and separation of concerns. No changes to inference or orchestration functionality. BREAKING CHANGE: inference module structure, type imports, and config utilities have changed; downstream code must update imports and integration points. --- compute-worker/src/api/app.ts | 4 +- compute-worker/src/api/contracts.ts | 14 +++--- compute-worker/src/api/operation-keys.ts | 2 +- .../types/parsed-pdf.ts => api/types.ts} | 17 +++++++ .../src/inference/config/cpu-budget.ts | 28 ------------ compute-worker/src/inference/index.ts | 26 ----------- .../pdf/{merge.ts => document-layout.ts} | 16 ++++++- .../{runLayoutModel.ts => layout-model.ts} | 4 +- compute-worker/src/inference/pdf/model.ts | 2 +- .../src/inference/pdf/normalize-text.ts | 2 +- compute-worker/src/inference/pdf/parse.ts | 10 ++--- .../src/inference/pdf/parser-version-key.ts | 9 ---- .../src/inference/pdf/parser-version.ts | 1 - .../pdf/{pdfjs-runtime.ts => pdfjs.ts} | 0 compute-worker/src/inference/pdf/render.ts | 2 +- compute-worker/src/inference/pdf/stitch.ts | 2 +- compute-worker/src/inference/pdf/types.ts | 15 ------- .../src/inference/platform/docstore.ts | 21 --------- .../{local-runtime.ts => runtime.ts} | 0 compute-worker/src/inference/types/index.ts | 17 ------- compute-worker/src/inference/types/tts.ts | 16 ------- compute-worker/src/inference/whisper/align.ts | 33 +++++++++++--- .../src/inference/whisper/alignment-map.ts | 44 ------------------- compute-worker/src/inference/whisper/model.ts | 2 +- .../src/inference/whisper/spectral.ts | 21 --------- .../{token-timestamps.ts => timestamps.ts} | 24 ++++++++++ .../timeout.ts => infrastructure/config.ts} | 20 +++++++++ .../ffmpeg.ts => infrastructure/platform.ts} | 29 +++++++++--- compute-worker/src/jobs/handlers.ts | 4 +- .../src/storage/artifact-addressing.ts | 4 +- .../pdf-merge-text-with-regions.test.ts | 2 +- .../pdf-parse-normalize-text-items.test.ts | 2 +- .../pdf-stitch-cross-page-blocks.test.ts | 4 +- .../algorithms/support/document-fixtures.ts | 2 +- .../whisper-alignment-mapping.test.ts | 2 +- .../algorithms/whisper-spectral.test.ts | 2 +- .../whisper-token-timestamps.test.ts | 2 +- .../control-plane/run-layout-model.test.ts | 6 +-- 38 files changed, 163 insertions(+), 248 deletions(-) rename compute-worker/src/{inference/types/parsed-pdf.ts => api/types.ts} (78%) delete mode 100644 compute-worker/src/inference/config/cpu-budget.ts delete mode 100644 compute-worker/src/inference/index.ts rename compute-worker/src/inference/pdf/{merge.ts => document-layout.ts} (94%) rename compute-worker/src/inference/pdf/{runLayoutModel.ts => layout-model.ts} (98%) delete mode 100644 compute-worker/src/inference/pdf/parser-version-key.ts delete mode 100644 compute-worker/src/inference/pdf/parser-version.ts rename compute-worker/src/inference/pdf/{pdfjs-runtime.ts => pdfjs.ts} (100%) delete mode 100644 compute-worker/src/inference/pdf/types.ts delete mode 100644 compute-worker/src/inference/platform/docstore.ts rename compute-worker/src/inference/{local-runtime.ts => runtime.ts} (100%) delete mode 100644 compute-worker/src/inference/types/index.ts delete mode 100644 compute-worker/src/inference/types/tts.ts delete mode 100644 compute-worker/src/inference/whisper/alignment-map.ts delete mode 100644 compute-worker/src/inference/whisper/spectral.ts rename compute-worker/src/inference/whisper/{token-timestamps.ts => timestamps.ts} (93%) rename compute-worker/src/{inference/config/timeout.ts => infrastructure/config.ts} (84%) rename compute-worker/src/{inference/platform/ffmpeg.ts => infrastructure/platform.ts} (60%) diff --git a/compute-worker/src/api/app.ts b/compute-worker/src/api/app.ts index 1775f4e..77b5a55 100644 --- a/compute-worker/src/api/app.ts +++ b/compute-worker/src/api/app.ts @@ -15,13 +15,13 @@ import { import { Kvm } from '@nats-io/kv'; import { ensureComputeModels, -} from '../inference/local-runtime'; +} from '../inference/runtime'; import { getComputeTimeoutConfig, getComputeOpStaleMs, getAvailableCpuCores, getOnnxThreadsPerJob, -} from '../inference'; +} from '../infrastructure/config'; import { encodeSseFrame, OperationOrchestrator } from '../operations'; import type { PdfLayoutJobRequest, diff --git a/compute-worker/src/api/contracts.ts b/compute-worker/src/api/contracts.ts index d6e87f3..4f78782 100644 --- a/compute-worker/src/api/contracts.ts +++ b/compute-worker/src/api/contracts.ts @@ -1,24 +1,26 @@ -import type { TTSSentenceAlignment } from '../inference/types/tts'; -import type { ParsedPdfDocument } from '../inference/types/parsed-pdf'; +import type { TTSSentenceAlignment, ParsedPdfDocument } from './types'; export type { TTSAudioBuffer, TTSAudioBytes, TTSSentenceAlignment, TTSSentenceWord, -} from '../inference/types/tts'; +} from './types'; export type { ParsedPdfBlockKind, ParsedPdfBlockFragment, ParsedPdfBlock, ParsedPdfPage, ParsedPdfDocument, -} from '../inference/types/parsed-pdf'; +} from './types'; export const ALIGN_QUEUE_NAME = 'whisper-align'; export const PDF_LAYOUT_QUEUE_NAME = 'pdf-layout'; -export { PDF_PARSER_VERSION } from '../inference/pdf/parser-version'; -export { encodeParserVersion } from '../inference/pdf/parser-version-key'; +export const PDF_PARSER_VERSION = 'pp-doclayoutv3-onnx@800+pdfjs@4.8.69'; + +export function encodeParserVersion(parserVersion: string, defaultVersion = PDF_PARSER_VERSION): string { + return encodeURIComponent(parserVersion.trim() || defaultVersion); +} export interface WhisperAlignJobBase { text: string; diff --git a/compute-worker/src/api/operation-keys.ts b/compute-worker/src/api/operation-keys.ts index a6763fe..d814fb2 100644 --- a/compute-worker/src/api/operation-keys.ts +++ b/compute-worker/src/api/operation-keys.ts @@ -1,5 +1,5 @@ import { createHash } from 'node:crypto'; -import { PDF_PARSER_VERSION } from '../inference/pdf/parser-version'; +import { PDF_PARSER_VERSION } from '../api/contracts'; function sha256Hex(input: string): string { return createHash('sha256').update(input).digest('hex'); diff --git a/compute-worker/src/inference/types/parsed-pdf.ts b/compute-worker/src/api/types.ts similarity index 78% rename from compute-worker/src/inference/types/parsed-pdf.ts rename to compute-worker/src/api/types.ts index 01d7067..16ad911 100644 --- a/compute-worker/src/inference/types/parsed-pdf.ts +++ b/compute-worker/src/api/types.ts @@ -61,3 +61,20 @@ export interface PdfParseProgress { currentPage?: number; phase: PdfParsePhase; } + +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[]; +} diff --git a/compute-worker/src/inference/config/cpu-budget.ts b/compute-worker/src/inference/config/cpu-budget.ts deleted file mode 100644 index 50d161b..0000000 --- a/compute-worker/src/inference/config/cpu-budget.ts +++ /dev/null @@ -1,28 +0,0 @@ -import os from 'node:os'; - -function readPositiveInt(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); -} - -export function getComputeJobConcurrency(): number { - return readPositiveInt('COMPUTE_JOB_CONCURRENCY', 1); -} - -export function getAvailableCpuCores(): number { - if (typeof os.availableParallelism === 'function') { - const value = os.availableParallelism(); - if (Number.isFinite(value) && value >= 1) return Math.floor(value); - } - const fallback = os.cpus().length; - return Number.isFinite(fallback) && fallback >= 1 ? Math.floor(fallback) : 1; -} - -export function getOnnxThreadsPerJob(): number { - const concurrency = getComputeJobConcurrency(); - const usableCores = Math.max(1, getAvailableCpuCores() - 1); - return Math.max(1, Math.floor(usableCores / concurrency)); -} diff --git a/compute-worker/src/inference/index.ts b/compute-worker/src/inference/index.ts deleted file mode 100644 index 255fdce..0000000 --- a/compute-worker/src/inference/index.ts +++ /dev/null @@ -1,26 +0,0 @@ -export * from '../api/contracts'; -export { - getComputeJobConcurrency, - getAvailableCpuCores, - getOnnxThreadsPerJob, -} from './config/cpu-budget'; -export { - getComputeTimeoutConfig, - getComputeOpStaleMs, - getWorkerClientWaitTimeoutMs, - withTimeout, - withIdleTimeoutAndHardCap, - type ComputeTimeoutConfig, - type ComputeOperationKind, - type IdleTimeoutAndHardCapInput, -} from './config/timeout'; -export { renderPage } from './pdf/render'; -export { mergeTextWithRegions } from './pdf/merge'; -export { PDF_PARSER_VERSION } from './pdf/parser-version'; -export { encodeParserVersion } from './pdf/parser-version-key'; -export { stitchCrossPageBlocks } from './pdf/stitch'; -export { normalizeTextItemsForLayout } from './pdf/normalize-text'; -export { mapWordsToSentenceOffsets, type WhisperWord } from './whisper/alignment-map'; -export { buildGoertzelCoefficients, goertzelPower } from './whisper/spectral'; -export { buildWordsFromTimestampedTokens, extractTokenStartTimestamps } from './whisper/token-timestamps'; -export * from '../operations'; diff --git a/compute-worker/src/inference/pdf/merge.ts b/compute-worker/src/inference/pdf/document-layout.ts similarity index 94% rename from compute-worker/src/inference/pdf/merge.ts rename to compute-worker/src/inference/pdf/document-layout.ts index e1f71b2..8d48fd9 100644 --- a/compute-worker/src/inference/pdf/merge.ts +++ b/compute-worker/src/inference/pdf/document-layout.ts @@ -1,4 +1,18 @@ -import type { LayoutRegion, PdfTextItem } from './types'; +import type { ParsedPdfBlockKind } from '../../api/types'; + +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; +} const NON_TEXT_REGION_LABELS = new Set(['chart', 'image', 'table', 'seal']); const TEXT_ASSIGNABLE_LABELS = new Set([ diff --git a/compute-worker/src/inference/pdf/runLayoutModel.ts b/compute-worker/src/inference/pdf/layout-model.ts similarity index 98% rename from compute-worker/src/inference/pdf/runLayoutModel.ts rename to compute-worker/src/inference/pdf/layout-model.ts index 5c726e0..4e1e6e2 100644 --- a/compute-worker/src/inference/pdf/runLayoutModel.ts +++ b/compute-worker/src/inference/pdf/layout-model.ts @@ -1,8 +1,8 @@ import * as ort from 'onnxruntime-node'; import { readFile } from 'fs/promises'; -import type { LayoutRegion, PdfTextItem } from './types'; +import type { LayoutRegion, PdfTextItem } from './document-layout'; import { ensureModel, MODEL_CONFIG_PATH, MODEL_PREPROCESSOR_PATH } from './model'; -import { getOnnxThreadsPerJob } from '../config/cpu-budget'; +import { getOnnxThreadsPerJob } from '../../infrastructure/config'; interface RunLayoutInput { pageWidth: number; diff --git a/compute-worker/src/inference/pdf/model.ts b/compute-worker/src/inference/pdf/model.ts index ca09540..2d13cda 100644 --- a/compute-worker/src/inference/pdf/model.ts +++ b/compute-worker/src/inference/pdf/model.ts @@ -3,7 +3,7 @@ import { fileURLToPath } from 'url'; import { createHash } from 'crypto'; import { readFileSync } from 'fs'; import { access, mkdir, rename, writeFile, readFile, unlink, copyFile } from 'fs/promises'; -import { DOCSTORE_DIR } from '../platform/docstore'; +import { DOCSTORE_DIR } from '../../infrastructure/platform'; 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'; diff --git a/compute-worker/src/inference/pdf/normalize-text.ts b/compute-worker/src/inference/pdf/normalize-text.ts index 1e91101..f8e5dc0 100644 --- a/compute-worker/src/inference/pdf/normalize-text.ts +++ b/compute-worker/src/inference/pdf/normalize-text.ts @@ -1,5 +1,5 @@ import type { TextItem } from 'pdfjs-dist/types/src/display/api'; -import type { PdfTextItem } from './types'; +import type { PdfTextItem } from './document-layout'; interface ViewportLike { height: number; diff --git a/compute-worker/src/inference/pdf/parse.ts b/compute-worker/src/inference/pdf/parse.ts index 69d10d2..b4491b2 100644 --- a/compute-worker/src/inference/pdf/parse.ts +++ b/compute-worker/src/inference/pdf/parse.ts @@ -1,10 +1,10 @@ import type { TextItem } from 'pdfjs-dist/types/src/display/api'; -import type { ParsedPdfDocument, ParsedPdfPage } from '../types/parsed-pdf'; +import type { ParsedPdfDocument, ParsedPdfPage } from '../../api/types'; import { ensureModel } from './model'; -import { runLayoutModel } from './runLayoutModel'; -import { mergeTextWithRegions } from './merge'; -import { configurePdfjsNodeRuntime, resolvePdfjsStandardFontDataUrl } from './pdfjs-runtime'; -import { PDF_PARSER_VERSION } from './parser-version'; +import { runLayoutModel } from './layout-model'; +import { mergeTextWithRegions } from './document-layout'; +import { configurePdfjsNodeRuntime, resolvePdfjsStandardFontDataUrl } from './pdfjs'; +import { PDF_PARSER_VERSION } from '../../api/contracts'; import { stitchCrossPageBlocks } from './stitch'; import { renderPage } from './render'; import { normalizeTextItemsForLayout } from './normalize-text'; diff --git a/compute-worker/src/inference/pdf/parser-version-key.ts b/compute-worker/src/inference/pdf/parser-version-key.ts deleted file mode 100644 index f9f82b2..0000000 --- a/compute-worker/src/inference/pdf/parser-version-key.ts +++ /dev/null @@ -1,9 +0,0 @@ -import { PDF_PARSER_VERSION } from './parser-version'; - -export function encodeParserVersion( - parserVersion: string, - defaultVersion = PDF_PARSER_VERSION, -): string { - const normalized = parserVersion.trim() || defaultVersion; - return encodeURIComponent(normalized); -} diff --git a/compute-worker/src/inference/pdf/parser-version.ts b/compute-worker/src/inference/pdf/parser-version.ts deleted file mode 100644 index b5fa7d7..0000000 --- a/compute-worker/src/inference/pdf/parser-version.ts +++ /dev/null @@ -1 +0,0 @@ -export const PDF_PARSER_VERSION = 'pp-doclayoutv3-onnx@800+pdfjs@4.8.69'; diff --git a/compute-worker/src/inference/pdf/pdfjs-runtime.ts b/compute-worker/src/inference/pdf/pdfjs.ts similarity index 100% rename from compute-worker/src/inference/pdf/pdfjs-runtime.ts rename to compute-worker/src/inference/pdf/pdfjs.ts diff --git a/compute-worker/src/inference/pdf/render.ts b/compute-worker/src/inference/pdf/render.ts index 33ec1a1..eb23993 100644 --- a/compute-worker/src/inference/pdf/render.ts +++ b/compute-worker/src/inference/pdf/render.ts @@ -1,5 +1,5 @@ import type { Canvas } from '@napi-rs/canvas'; -import { configurePdfjsNodeRuntime, resolvePdfjsStandardFontDataUrl } from './pdfjs-runtime'; +import { configurePdfjsNodeRuntime, resolvePdfjsStandardFontDataUrl } from './pdfjs'; type CanvasRuntime = { DOMMatrixCtor: unknown; diff --git a/compute-worker/src/inference/pdf/stitch.ts b/compute-worker/src/inference/pdf/stitch.ts index 98a3d77..d3da6f6 100644 --- a/compute-worker/src/inference/pdf/stitch.ts +++ b/compute-worker/src/inference/pdf/stitch.ts @@ -1,4 +1,4 @@ -import type { ParsedPdfDocument, ParsedPdfBlock } from '../types/parsed-pdf'; +import type { ParsedPdfDocument, ParsedPdfBlock } from '../../api/types'; const STITCHABLE_KINDS = new Set([ 'text', diff --git a/compute-worker/src/inference/pdf/types.ts b/compute-worker/src/inference/pdf/types.ts deleted file mode 100644 index 38089a6..0000000 --- a/compute-worker/src/inference/pdf/types.ts +++ /dev/null @@ -1,15 +0,0 @@ -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; -} diff --git a/compute-worker/src/inference/platform/docstore.ts b/compute-worker/src/inference/platform/docstore.ts deleted file mode 100644 index 0b62401..0000000 --- a/compute-worker/src/inference/platform/docstore.ts +++ /dev/null @@ -1,21 +0,0 @@ -import fs from 'fs'; -import path from 'path'; - -function findMonorepoRoot(startDir: string): string | null { - let current = path.resolve(startDir); - for (;;) { - const marker = path.join(current, 'pnpm-workspace.yaml'); - if (fs.existsSync(marker)) return current; - const parent = path.dirname(current); - if (parent === current) return null; - current = parent; - } -} - -function resolveDocstoreDir(): string { - const repoRoot = findMonorepoRoot(process.cwd()); - if (repoRoot) return path.join(repoRoot, 'docstore'); - return path.join(process.cwd(), 'docstore'); -} - -export const DOCSTORE_DIR = resolveDocstoreDir(); diff --git a/compute-worker/src/inference/local-runtime.ts b/compute-worker/src/inference/runtime.ts similarity index 100% rename from compute-worker/src/inference/local-runtime.ts rename to compute-worker/src/inference/runtime.ts diff --git a/compute-worker/src/inference/types/index.ts b/compute-worker/src/inference/types/index.ts deleted file mode 100644 index 7408914..0000000 --- a/compute-worker/src/inference/types/index.ts +++ /dev/null @@ -1,17 +0,0 @@ -export type { - TTSAudioBuffer, - TTSAudioBytes, - TTSSentenceAlignment, - TTSSentenceWord, -} from './tts'; - -export type { - ParsedPdfBlock, - ParsedPdfBlockFragment, - ParsedPdfBlockKind, - ParsedPdfDocument, - ParsedPdfPage, - PdfParsePhase, - PdfParseProgress, - PdfParseStatus, -} from './parsed-pdf'; diff --git a/compute-worker/src/inference/types/tts.ts b/compute-worker/src/inference/types/tts.ts deleted file mode 100644 index 73db759..0000000 --- a/compute-worker/src/inference/types/tts.ts +++ /dev/null @@ -1,16 +0,0 @@ -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[]; -} diff --git a/compute-worker/src/inference/whisper/align.ts b/compute-worker/src/inference/whisper/align.ts index d18fa07..073afc4 100644 --- a/compute-worker/src/inference/whisper/align.ts +++ b/compute-worker/src/inference/whisper/align.ts @@ -7,19 +7,18 @@ import { fileURLToPath } from 'url'; import * as ort from 'onnxruntime-node'; import { Tokenizer } from '@huggingface/tokenizers'; import JSZip from 'jszip'; -import type { TTSAudioBuffer, TTSAudioBytes, TTSSentenceAlignment } from '../types/tts'; -import { getFFmpegPath } from '../platform/ffmpeg'; -import { getOnnxThreadsPerJob } from '../config/cpu-budget'; -import { getComputeTimeoutConfig } from '../config/timeout'; +import type { TTSAudioBuffer, TTSAudioBytes, TTSSentenceAlignment } from '../../api/types'; +import { getFFmpegPath } from '../../infrastructure/platform'; +import { getOnnxThreadsPerJob } from '../../infrastructure/config'; +import { getComputeTimeoutConfig } from '../../infrastructure/config'; import { mapWordsToSentenceOffsets, type WhisperWord, -} from './alignment-map'; -import { buildGoertzelCoefficients, goertzelPower } from './spectral'; +} from './timestamps'; import { buildWordsFromTimestampedTokens, extractTokenStartTimestamps, -} from './token-timestamps'; +} from './timestamps'; import { ensureWhisperModel, WHISPER_CONFIG_PATH, @@ -31,6 +30,26 @@ import { WHISPER_DECODER_WITH_PAST_MODEL_PATH, } from './model'; +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); + return !Number.isFinite(power) || power < 0 ? 0 : power; +} + interface WhisperAlignmentOptions { lang?: string; textHint?: string; diff --git a/compute-worker/src/inference/whisper/alignment-map.ts b/compute-worker/src/inference/whisper/alignment-map.ts deleted file mode 100644 index 098374c..0000000 --- a/compute-worker/src/inference/whisper/alignment-map.ts +++ /dev/null @@ -1,44 +0,0 @@ -import type { TTSSentenceAlignment, TTSSentenceWord } from '../types/tts'; - -export interface WhisperWord { - start: number; - end: number; - word: string; -} - -export function mapWordsToSentenceOffsets(sentence: string, words: WhisperWord[]): TTSSentenceAlignment { - const lowerSentence = sentence.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(sentence.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, - }; -} diff --git a/compute-worker/src/inference/whisper/model.ts b/compute-worker/src/inference/whisper/model.ts index 4c5a021..9cf0c4f 100644 --- a/compute-worker/src/inference/whisper/model.ts +++ b/compute-worker/src/inference/whisper/model.ts @@ -3,7 +3,7 @@ import { fileURLToPath } from 'url'; import { createHash } from 'crypto'; import { readFileSync } from 'fs'; import { access, copyFile, mkdir, readFile, rename, unlink, writeFile } from 'fs/promises'; -import { DOCSTORE_DIR } from '../platform/docstore'; +import { DOCSTORE_DIR } from '../../infrastructure/platform'; const MODULE_DIR = path.dirname(fileURLToPath(import.meta.url)); const MODEL_DIR = path.join(DOCSTORE_DIR, 'model', 'whisper-base_timestamped'); diff --git a/compute-worker/src/inference/whisper/spectral.ts b/compute-worker/src/inference/whisper/spectral.ts deleted file mode 100644 index b7223cc..0000000 --- a/compute-worker/src/inference/whisper/spectral.ts +++ /dev/null @@ -1,21 +0,0 @@ -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; -} diff --git a/compute-worker/src/inference/whisper/token-timestamps.ts b/compute-worker/src/inference/whisper/timestamps.ts similarity index 93% rename from compute-worker/src/inference/whisper/token-timestamps.ts rename to compute-worker/src/inference/whisper/timestamps.ts index 47edc1d..e6e2b96 100644 --- a/compute-worker/src/inference/whisper/token-timestamps.ts +++ b/compute-worker/src/inference/whisper/timestamps.ts @@ -1,5 +1,6 @@ import type { Tokenizer } from '@huggingface/tokenizers'; import type * as ort from 'onnxruntime-node'; +import type { TTSSentenceAlignment, TTSSentenceWord } from '../../api/types'; const PUNCTUATION_REGEX = '\\p{P}\\u0021-\\u002F\\u003A-\\u0040\\u005B-\\u0060\\u007B-\\u007E'; const PUNCTUATION_ONLY_REGEX = new RegExp(`^[${PUNCTUATION_REGEX}]+$`, 'gu'); @@ -12,6 +13,29 @@ export interface WhisperWordTiming { endSec: number; } +export interface WhisperWord { + start: number; + end: number; + word: string; +} + +export function mapWordsToSentenceOffsets(sentence: string, words: WhisperWord[]): TTSSentenceAlignment { + const lowerSentence = sentence.toLowerCase(); + let cursor = 0; + const alignedWords: TTSSentenceWord[] = words.map((word) => { + const token = word.word.trim(); + if (!token) { + return { text: '', startSec: word.start, endSec: word.end, charStart: cursor, charEnd: cursor }; + } + const index = lowerSentence.indexOf(token.toLowerCase(), cursor); + const start = index >= 0 ? index : cursor; + const end = Math.min(sentence.length, start + token.length); + cursor = Math.max(cursor, end); + return { text: token, startSec: word.start, endSec: word.end, charStart: start, charEnd: end }; + }).filter((word) => word.text.length > 0); + return { sentence, sentenceIndex: 0, words: alignedWords }; +} + function medianFilter(data: Float32Array, windowSize: number): Float32Array { if (windowSize % 2 === 0 || windowSize <= 0) { throw new Error('Window size must be a positive odd number'); diff --git a/compute-worker/src/inference/config/timeout.ts b/compute-worker/src/infrastructure/config.ts similarity index 84% rename from compute-worker/src/inference/config/timeout.ts rename to compute-worker/src/infrastructure/config.ts index 40fd406..2547f3b 100644 --- a/compute-worker/src/inference/config/timeout.ts +++ b/compute-worker/src/infrastructure/config.ts @@ -117,3 +117,23 @@ export async function withIdleTimeoutAndHardCap(input: IdleTimeoutAndHardCapI throw error; } } + +export function getComputeJobConcurrency(): number { + return readPositiveIntEnv('COMPUTE_JOB_CONCURRENCY', 1); +} + +export function getAvailableCpuCores(): number { + if (typeof os.availableParallelism === 'function') { + const value = os.availableParallelism(); + if (Number.isFinite(value) && value >= 1) return Math.floor(value); + } + const fallback = os.cpus().length; + return Number.isFinite(fallback) && fallback >= 1 ? Math.floor(fallback) : 1; +} + +export function getOnnxThreadsPerJob(): number { + const concurrency = getComputeJobConcurrency(); + const usableCores = Math.max(1, getAvailableCpuCores() - 1); + return Math.max(1, Math.floor(usableCores / concurrency)); +} +import os from 'node:os'; diff --git a/compute-worker/src/inference/platform/ffmpeg.ts b/compute-worker/src/infrastructure/platform.ts similarity index 60% rename from compute-worker/src/inference/platform/ffmpeg.ts rename to compute-worker/src/infrastructure/platform.ts index b4ed669..9f029e4 100644 --- a/compute-worker/src/inference/platform/ffmpeg.ts +++ b/compute-worker/src/infrastructure/platform.ts @@ -1,6 +1,26 @@ -import { existsSync } from 'fs'; +import fs from 'fs'; +import path from 'path'; import ffmpegStatic from 'ffmpeg-static'; +function findMonorepoRoot(startDir: string): string | null { + let current = path.resolve(startDir); + for (;;) { + const marker = path.join(current, 'pnpm-workspace.yaml'); + if (fs.existsSync(marker)) return current; + const parent = path.dirname(current); + if (parent === current) return null; + current = parent; + } +} + +function resolveDocstoreDir(): string { + const repoRoot = findMonorepoRoot(process.cwd()); + if (repoRoot) return path.join(repoRoot, 'docstore'); + return path.join(process.cwd(), 'docstore'); +} + +export const DOCSTORE_DIR = resolveDocstoreDir(); + function normalizePath(value: unknown): string | null { if (typeof value !== 'string') return null; const trimmed = value.trim(); @@ -9,20 +29,17 @@ function normalizePath(value: unknown): string | null { function resolveBinary(envValue: string | null, bundledValue: string | null, envVarName: string, packageName: string): string { if (envValue) { - if ((envValue.includes('/') || envValue.includes('\\')) && !existsSync(envValue)) { + if ((envValue.includes('/') || envValue.includes('\\')) && !fs.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)) { + if ((bundledValue.includes('/') || bundledValue.includes('\\')) && !fs.existsSync(bundledValue)) { throw new Error(`${packageName} resolved to a missing binary path: ${bundledValue}`); } - return bundledValue; } diff --git a/compute-worker/src/jobs/handlers.ts b/compute-worker/src/jobs/handlers.ts index 90a4068..a013602 100644 --- a/compute-worker/src/jobs/handlers.ts +++ b/compute-worker/src/jobs/handlers.ts @@ -2,8 +2,8 @@ import { z } from 'zod'; import { runPdfLayoutFromPdfBuffer, runWhisperAlignmentFromAudioBuffer, -} from '../inference/local-runtime'; -import { withIdleTimeoutAndHardCap, withTimeout } from '../inference'; +} from '../inference/runtime'; +import { withIdleTimeoutAndHardCap, withTimeout } from '../infrastructure/config'; import type { PdfLayoutJobRequest, PdfLayoutJobResult, diff --git a/compute-worker/src/storage/artifact-addressing.ts b/compute-worker/src/storage/artifact-addressing.ts index e341f23..bf4e312 100644 --- a/compute-worker/src/storage/artifact-addressing.ts +++ b/compute-worker/src/storage/artifact-addressing.ts @@ -1,5 +1,5 @@ -import { PDF_PARSER_VERSION } from '../inference/pdf/parser-version'; -import { encodeParserVersion } from '../inference/pdf/parser-version-key'; +import { PDF_PARSER_VERSION } from '../api/contracts'; +import { encodeParserVersion } from '../api/contracts'; const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i; const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/; diff --git a/compute-worker/tests/compute/algorithms/pdf-merge-text-with-regions.test.ts b/compute-worker/tests/compute/algorithms/pdf-merge-text-with-regions.test.ts index d035613..cb18eee 100644 --- a/compute-worker/tests/compute/algorithms/pdf-merge-text-with-regions.test.ts +++ b/compute-worker/tests/compute/algorithms/pdf-merge-text-with-regions.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from 'vitest'; -import { mergeTextWithRegions } from '../../../src/inference'; +import { mergeTextWithRegions } from '../../../src/inference/pdf/document-layout'; describe('mergeTextWithRegions', () => { test('assigns text items to containing regions by centroid and joins in reading order', () => { diff --git a/compute-worker/tests/compute/algorithms/pdf-parse-normalize-text-items.test.ts b/compute-worker/tests/compute/algorithms/pdf-parse-normalize-text-items.test.ts index 167dc46..f93cf25 100644 --- a/compute-worker/tests/compute/algorithms/pdf-parse-normalize-text-items.test.ts +++ b/compute-worker/tests/compute/algorithms/pdf-parse-normalize-text-items.test.ts @@ -1,6 +1,6 @@ import { describe, expect, test } from 'vitest'; -import { normalizeTextItemsForLayout } from '../../../src/inference'; +import { normalizeTextItemsForLayout } from '../../../src/inference/pdf/normalize-text'; import type { TextItem } from 'pdfjs-dist/types/src/display/api'; function makeTextItem( diff --git a/compute-worker/tests/compute/algorithms/pdf-stitch-cross-page-blocks.test.ts b/compute-worker/tests/compute/algorithms/pdf-stitch-cross-page-blocks.test.ts index 97d4b19..3fd35a6 100644 --- a/compute-worker/tests/compute/algorithms/pdf-stitch-cross-page-blocks.test.ts +++ b/compute-worker/tests/compute/algorithms/pdf-stitch-cross-page-blocks.test.ts @@ -1,6 +1,6 @@ import { describe, expect, test } from 'vitest'; -import { stitchCrossPageBlocks } from '../../../src/inference'; -import type { ParsedPdfBlock, ParsedPdfDocument, ParsedPdfBlockKind } from '../../../src/inference/types'; +import { stitchCrossPageBlocks } from '../../../src/inference/pdf/stitch'; +import type { ParsedPdfBlock, ParsedPdfDocument, ParsedPdfBlockKind } from '../../../src/api/types'; import { makeParsedPdfDocument, makeParsedPdfPage } from './support/document-fixtures'; function makeBlock( diff --git a/compute-worker/tests/compute/algorithms/support/document-fixtures.ts b/compute-worker/tests/compute/algorithms/support/document-fixtures.ts index f8db418..5a3d697 100644 --- a/compute-worker/tests/compute/algorithms/support/document-fixtures.ts +++ b/compute-worker/tests/compute/algorithms/support/document-fixtures.ts @@ -1,5 +1,5 @@ import type { BaseDocument } from '../../../../../src/types/documents'; -import type { ParsedPdfBlock, ParsedPdfBlockKind, ParsedPdfDocument, ParsedPdfPage } from '../../../../src/inference/types'; +import type { ParsedPdfBlock, ParsedPdfBlockKind, ParsedPdfDocument, ParsedPdfPage } from '../../../../src/api/types'; export function makeBaseDocument(overrides: Partial = {}): BaseDocument { return { diff --git a/compute-worker/tests/compute/algorithms/whisper-alignment-mapping.test.ts b/compute-worker/tests/compute/algorithms/whisper-alignment-mapping.test.ts index fedea8b..201c9c6 100644 --- a/compute-worker/tests/compute/algorithms/whisper-alignment-mapping.test.ts +++ b/compute-worker/tests/compute/algorithms/whisper-alignment-mapping.test.ts @@ -1,7 +1,7 @@ import { describe, expect, test } from 'vitest'; import { mapWordsToSentenceOffsets, -} from '../../../src/inference'; +} from '../../../src/inference/whisper/timestamps'; describe('whisper alignment mapping', () => { test('maps words to sentence offsets with punctuation and repeated spaces', () => { diff --git a/compute-worker/tests/compute/algorithms/whisper-spectral.test.ts b/compute-worker/tests/compute/algorithms/whisper-spectral.test.ts index 6b51d62..1794290 100644 --- a/compute-worker/tests/compute/algorithms/whisper-spectral.test.ts +++ b/compute-worker/tests/compute/algorithms/whisper-spectral.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from 'vitest'; -import { buildGoertzelCoefficients, goertzelPower } from '../../../src/inference'; +import { buildGoertzelCoefficients, goertzelPower } from '../../../src/inference/whisper/align'; function dftPower(samples: Float32Array, k: number): number { const n = samples.length; diff --git a/compute-worker/tests/compute/algorithms/whisper-token-timestamps.test.ts b/compute-worker/tests/compute/algorithms/whisper-token-timestamps.test.ts index 6f8bced..e1cdc18 100644 --- a/compute-worker/tests/compute/algorithms/whisper-token-timestamps.test.ts +++ b/compute-worker/tests/compute/algorithms/whisper-token-timestamps.test.ts @@ -3,7 +3,7 @@ import * as ort from 'onnxruntime-node'; import { buildWordsFromTimestampedTokens, extractTokenStartTimestamps, -} from '../../../src/inference'; +} from '../../../src/inference/whisper/timestamps'; describe('whisper token timestamp alignment', () => { test('extracts monotonic token timestamps from cross-attention maps', () => { diff --git a/compute-worker/tests/compute/control-plane/run-layout-model.test.ts b/compute-worker/tests/compute/control-plane/run-layout-model.test.ts index c8278c9..cc6a34d 100644 --- a/compute-worker/tests/compute/control-plane/run-layout-model.test.ts +++ b/compute-worker/tests/compute/control-plane/run-layout-model.test.ts @@ -76,7 +76,7 @@ vi.mock('../../../src/inference/pdf/model', () => ({ MODEL_PREPROCESSOR_PATH: '/tmp/model-preprocessor.json', })); -vi.mock('../../../src/inference/config/cpu-budget', () => ({ +vi.mock('../../../src/infrastructure/config', () => ({ getOnnxThreadsPerJob: vi.fn(() => 1), })); @@ -107,7 +107,7 @@ describe('runLayoutModel', () => { }, }; - const { runLayoutModel } = await import('../../../src/inference/pdf/runLayoutModel'); + const { runLayoutModel } = await import('../../../src/inference/pdf/layout-model'); const regions = await runLayoutModel({ pageWidth: 100, pageHeight: 100, @@ -154,7 +154,7 @@ describe('runLayoutModel', () => { }, }; - const { runLayoutModel } = await import('../../../src/inference/pdf/runLayoutModel'); + const { runLayoutModel } = await import('../../../src/inference/pdf/layout-model'); const regions = await runLayoutModel({ pageWidth: 100, pageHeight: 100,