From b3f0486c5a848a3a862654c7938a8143c4db4c8c Mon Sep 17 00:00:00 2001 From: Miltiadis Vasilakis Date: Sun, 14 Jun 2026 14:11:38 +0300 Subject: [PATCH] Add host ONNX runtime option --- src/index.js | 20 +++- src/pdf/index.js | 14 ++- .../structure/model/onnx/host-onnx-runtime.js | 91 +++++++++++++++++++ src/pdf/structure/model/onnx/runtime.js | 4 + 4 files changed, 122 insertions(+), 7 deletions(-) create mode 100644 src/pdf/structure/model/onnx/host-onnx-runtime.js diff --git a/src/index.js b/src/index.js index 2ad4077..fe404e2 100644 --- a/src/index.js +++ b/src/index.js @@ -40,6 +40,7 @@ async function createStructuredDocumentText(buf, options = {}) { contentType, password, dataProvider, + nativeONNXRun, sourceHash, } = options; assertSourceHash(sourceHash); @@ -50,6 +51,7 @@ async function createStructuredDocumentText(buf, options = {}) { : await getSnapshotStructure(buf, contentType, { sourceHash }); } return await pdfGetStructure(buf, password, dataProvider, { + nativeONNXRun, sourceHash, }); } @@ -59,6 +61,7 @@ async function getStructuredDocumentText(buf, options = {}) { contentType: options.contentType, password: options.password, dataProvider: options.dataProvider, + nativeONNXRun: options.nativeONNXRun, sourceHash: options.sourceHash, }); let buffer = packStructuredDocumentText(structure, { @@ -110,9 +113,9 @@ if (typeof self !== 'undefined') { let waitingPromises = {}; self.query = async function (action, data, transfer) { - return new Promise(function (resolve) { + return new Promise(function (resolve, reject) { promiseID++; - waitingPromises[promiseID] = resolve; + waitingPromises[promiseID] = { resolve, reject }; self.postMessage({ id: promiseID, action, data }, transfer); }); }; @@ -121,9 +124,15 @@ if (typeof self !== 'undefined') { let message = e.data; if (message.responseID) { - let resolve = waitingPromises[message.responseID]; - if (resolve) { - resolve(message.data); + let waiting = waitingPromises[message.responseID]; + if (waiting) { + delete waitingPromises[message.responseID]; + if (message.error) { + waiting.reject(message.error); + } + else { + waiting.resolve(message.data); + } } return; } @@ -274,6 +283,7 @@ if (typeof self !== 'undefined') { contentType: message.data.contentType, password: message.data.password, dataProvider: fetchData, + nativeONNXRun: message.data.nativeONNX ? data => query('NativeONNXRun', data) : null, sourceHash: message.data.sourceHash, }); self.postMessage({ diff --git a/src/pdf/index.js b/src/pdf/index.js index 385ea90..50f7b09 100644 --- a/src/pdf/index.js +++ b/src/pdf/index.js @@ -659,8 +659,18 @@ async function getStructure(buf, password, dataProvider, options = {}) { let pdfManager = await getPdfManager(buf); setHandler(pdfManager.pdfDocument, dataProvider); - let onnxRuntimeProvider = () => dataProvider('onnx/ort-wasm-simd.wasm'); - let modelProvider = (name) => dataProvider(name.includes('/') ? name : name + '/model.onnx'); + let useNativeONNX = typeof options.nativeONNXRun === 'function'; + let onnxRuntimeProvider = useNativeONNX + ? { type: 'native', run: options.nativeONNXRun } + : () => dataProvider('onnx/ort-wasm-simd.wasm'); + let modelProvider = async (name) => { + let path = name.includes('/') ? name : name + '/model.onnx'; + if (useNativeONNX && path.endsWith('.onnx')) { + return { nativeONNXModel: path }; + } + let data = await dataProvider(path); + return data; + }; return await getFullStructure(pdfManager.pdfDocument, onnxRuntimeProvider, modelProvider, options); } diff --git a/src/pdf/structure/model/onnx/host-onnx-runtime.js b/src/pdf/structure/model/onnx/host-onnx-runtime.js new file mode 100644 index 0000000..cb0b728 --- /dev/null +++ b/src/pdf/structure/model/onnx/host-onnx-runtime.js @@ -0,0 +1,91 @@ +class NativeTensor { + constructor(type, data, dims) { + this.type = type; + this.data = data; + this.dims = dims; + } +} + +class NativeInferenceSession { + constructor(model, nativeONNXRun) { + this.model = model; + this.nativeONNXRun = nativeONNXRun; + } + + static async create(model, _options) { + const modelName = model?.nativeONNXModel; + const nativeONNXRun = model?.nativeONNXRun; + if (!modelName || typeof nativeONNXRun !== 'function') { + throw new Error('Native ONNX sessions require a native model provider'); + } + return new NativeInferenceSession(modelName, nativeONNXRun); + } + + async run(feeds, outputNames) { + const inputs = Object.entries(feeds).map(([name, tensor]) => ({ + name, + type: tensor.type, + dims: Array.from(tensor.dims), + values: serializeTensorValues(tensor), + })); + const result = await this.nativeONNXRun({ + model: this.model, + inputs, + outputNames: outputNames ?? null, + }); + return deserializeOutputs(result?.outputs); + } + + async release() { + // Native sessions are cached and owned by the host runtime. + } +} + +export function createNativeRuntime(nativeONNXRun) { + return { + Tensor: NativeTensor, + InferenceSession: { + create: (model, options) => NativeInferenceSession.create({ + ...model, + nativeONNXRun, + }, options), + }, + }; +} + +function serializeTensorValues(tensor) { + if (tensor.type === 'int64' && typeof BigInt64Array !== 'undefined' && tensor.data instanceof BigInt64Array) { + return Array.from(tensor.data, Number); + } + return Array.from(tensor.data); +} + +function deserializeOutputs(outputs) { + if (!outputs || typeof outputs !== 'object') { + throw new Error('Native ONNX did not return outputs'); + } + const result = {}; + for (const [name, output] of Object.entries(outputs)) { + result[name] = new NativeTensor(output.type, typedArray(output.type, output.values), output.dims); + } + return result; +} + +function typedArray(type, values) { + switch (type) { + case 'float32': + return Float32Array.from(values); + + case 'int64': + if (typeof BigInt64Array === 'undefined') { + throw new Error('BigInt64Array is not available'); + } + return BigInt64Array.from(values.map(BigInt)); + + case 'bool': + return Uint8Array.from(values); + + default: + throw new Error(`Unsupported native ONNX output type: ${type}`); + } +} diff --git a/src/pdf/structure/model/onnx/runtime.js b/src/pdf/structure/model/onnx/runtime.js index 85491ca..8510fd6 100644 --- a/src/pdf/structure/model/onnx/runtime.js +++ b/src/pdf/structure/model/onnx/runtime.js @@ -1,9 +1,13 @@ import * as ort from './ort.wasm.min.js'; import { Mutex } from '../../../mutex.js'; +import { createNativeRuntime } from './host-onnx-runtime.js'; export let onnxMutex = new Mutex() export async function getRuntime(onnxRuntimeProvider) { + if (onnxRuntimeProvider?.type === 'native') { + return createNativeRuntime(onnxRuntimeProvider.run); + } return await onnxMutex.runExclusive(async () => { ort.env.wasm.simd = true; ort.env.wasm.numThreads = 1;