onnxruntime/js/react_native/lib/binding.ts

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

// eslint-disable-next-line @typescript-eslint/no-unused-vars
import type {InferenceSession} from 'onnxruntime-common';
import {NativeModules} from 'react-native';

/**
 * model loading information
 */
interface ModelLoadInfo {
  /**
   * Key for an instance of InferenceSession, which is passed to run() function as parameter.
   */
  readonly key: string;

  /**
   * Get input names of the loaded model.
   */
  readonly inputNames: string[];

  /**
   * Get output names of the loaded model.
   */
  readonly outputNames: string[];
}

/**
 * JSIBlob is a blob object that exchange ArrayBuffer by OnnxruntimeJSIHelper.
 */
export type JSIBlob = {
  blobId: string; offset: number; size: number;
};

/**
 * Tensor type for react native, which doesn't allow ArrayBuffer in native bridge, so data will be stored as JSIBlob.
 */
interface EncodedTensor {
  /**
   * the dimensions of the tensor.
   */
  readonly dims: readonly number[];
  /**
   * the data type of the tensor.
   */
  readonly type: string;
  /**
   * the JSIBlob object of the buffer data of the tensor.
   * if data is string array, it won't be stored as JSIBlob.
   */
  readonly data: JSIBlob|string[];
}

/**
 * Binding exports a simple synchronized inference session object wrap.
 */
export declare namespace Binding {
  type ModelLoadInfoType = ModelLoadInfo;
  type EncodedTensorType = EncodedTensor;

  type SessionOptions = InferenceSession.SessionOptions;
  type RunOptions = InferenceSession.RunOptions;

  type FeedsType = {[name: string]: EncodedTensor};

  // SessionHanlder FetchesType is different from native module's one.
  // It's because Java API doesn't support preallocated output values.
  type FetchesType = string[];

  type ReturnType = {[name: string]: EncodedTensor};

  interface InferenceSession {
    loadModel(modelPath: string, options: SessionOptions): Promise<ModelLoadInfoType>;
    loadModelFromBlob?(blob: JSIBlob, options: SessionOptions): Promise<ModelLoadInfoType>;
    dispose(key: string): Promise<void>;
    run(key: string, feeds: FeedsType, fetches: FetchesType, options: RunOptions): Promise<ReturnType>;
  }
}

// export native binding
const {Onnxruntime, OnnxruntimeJSIHelper} = NativeModules;
export const binding = Onnxruntime as Binding.InferenceSession;

// install JSI helper global functions
OnnxruntimeJSIHelper.install();

declare global {
  // eslint-disable-next-line no-var
  var jsiOnnxruntimeStoreArrayBuffer: ((buffer: ArrayBuffer) => JSIBlob)|undefined;
  // eslint-disable-next-line no-var
  var jsiOnnxruntimeResolveArrayBuffer: ((blob: JSIBlob) => ArrayBuffer)|undefined;
}

export const jsiHelper = {
  storeArrayBuffer: globalThis.jsiOnnxruntimeStoreArrayBuffer || (() => {
                      throw new Error(
                          'jsiOnnxruntimeStoreArrayBuffer is not found, ' +
                          'please make sure OnnxruntimeJSIHelper installation is successful.');
                    }),
  resolveArrayBuffer: globalThis.jsiOnnxruntimeResolveArrayBuffer || (() => {
                        throw new Error(
                            'jsiOnnxruntimeResolveArrayBuffer is not found, ' +
                            'please make sure OnnxruntimeJSIHelper installation is successful.');
                      }),
};

// Remove global functions after installation
{
  delete globalThis.jsiOnnxruntimeStoreArrayBuffer;
  delete globalThis.jsiOnnxruntimeResolveArrayBuffer;
}