feat: creating inference_api

- Exposing an user friendly interface to consume the `onnx` backend

Author: kallebysantos

examples/ort-raw-session/index.ts

@@ -0,0 +1,41 @@
+const { Tensor, RawSession } = Supabase.ai;
+
+const session = await RawSession.fromHuggingFace('kallebysantos/vehicle-emission', {
+  path: {
+    modelFile: 'model.onnx',
+  },
+});
+
+Deno.serve(async (_req: Request) => {
+  // sample data could be a JSON request
+  const carsBatchInput = [{
+    'Model_Year': 2021,
+    'Engine_Size': 2.9,
+    'Cylinders': 6,
+    'Fuel_Consumption_in_City': 13.9,
+    'Fuel_Consumption_in_City_Hwy': 10.3,
+    'Fuel_Consumption_comb': 12.3,
+    'Smog_Level': 3,
+  }, {
+    'Model_Year': 2023,
+    'Engine_Size': 2.4,
+    'Cylinders': 4,
+    'Fuel_Consumption_in_City': 9.9,
+    'Fuel_Consumption_in_City_Hwy': 7.0,
+    'Fuel_Consumption_comb': 8.6,
+    'Smog_Level': 3,
+  }];
+
+  // Parsing objects to tensor input
+  const inputTensors = {};
+  session.inputs.forEach((inputKey) => {
+    const values = carsBatchInput.map((item) => item[inputKey]);
+
+    inputTensors[inputKey] = new Tensor('float32', values, [values.length, 1]);
+  });
+
+  const { emissions } = await session.run(inputTensors);
+  // [ 289.01, 199.53]
+
+  return Response.json({ result: emissions });
+});

ext/ai/js/ai.js

@@ -1,5 +1,6 @@
-import "ext:ai/onnxruntime/onnx.js";
-import EventSourceStream from "ext:ai/util/event_source_stream.mjs";
+import 'ext:ai/onnxruntime/onnx.js';
+import InferenceAPI from 'ext:ai/onnxruntime/inference_api.js';
+import EventSourceStream from 'ext:ai/util/event_source_stream.mjs';
 
 const core = globalThis.Deno.core;
 
@@ -257,7 +258,8 @@ const MAIN_WORKER_API = {
 };
 
 const USER_WORKER_API = {
-  Session,
+    Session,
+    ...InferenceAPI
 };
 
 export { MAIN_WORKER_API, USER_WORKER_API };

ext/ai/js/onnxruntime/inference_api.js

@@ -0,0 +1,75 @@
+import { InferenceSession, Tensor } from 'ext:ai/onnxruntime/onnx.js';
+
+const DEFAULT_HUGGING_FACE_OPTIONS = {
+  hostname: 'https://huggingface.co',
+  path: {
+    template: '{REPO_ID}/resolve/{REVISION}/onnx/{MODEL_FILE}?donwload=true',
+    revision: 'main',
+    modelFile: 'model_quantized.onnx',
+  },
+};
+
+/**
+ * An user friendly API for onnx backend
+ */
+class UserInferenceSession {
+  inner;
+
+  id;
+  inputs;
+  outputs;
+
+  constructor(session) {
+    this.inner = session;
+
+    this.id = session.sessionId;
+    this.inputs = session.inputNames;
+    this.outputs = session.outputNames;
+  }
+
+  static async fromUrl(modelUrl) {
+    if (modelUrl instanceof URL) {
+      modelUrl = modelUrl.toString();
+    }
+
+    const encoder = new TextEncoder();
+    const modelUrlBuffer = encoder.encode(modelUrl);
+    const session = await InferenceSession.fromBuffer(modelUrlBuffer);
+
+    return new UserInferenceSession(session);
+  }
+
+  static async fromHuggingFace(repoId, opts = {}) {
+    const hostname = opts?.hostname ?? DEFAULT_HUGGING_FACE_OPTIONS.hostname;
+    const pathOpts = {
+      ...DEFAULT_HUGGING_FACE_OPTIONS.path,
+      ...opts?.path,
+    };
+
+    const modelPath = pathOpts.template
+      .replaceAll('{REPO_ID}', repoId)
+      .replaceAll('{REVISION}', pathOpts.revision)
+      .replaceAll('{MODEL_FILE}', pathOpts.modelFile);
+
+    if (!URL.canParse(modelPath, hostname)) {
+      throw Error(`[Invalid URL] Couldn't parse the model path: "${modelPath}"`);
+    }
+
+    return await UserInferenceSession.fromUrl(new URL(modelPath, hostname));
+  }
+
+  async run(inputs) {
+    return await this.inner.run(inputs);
+  }
+}
+
+class UserTensor extends Tensor {
+  constructor(type, data, dim) {
+    super(type, data, dim);
+  }
+}
+
+export default {
+  RawSession: UserInferenceSession,
+  Tensor: UserTensor,
+};

ext/ai/js/onnxruntime/onnx.js

@@ -31,7 +31,7 @@ class TensorProxy {
   }
 }
 
-class Tensor {
+export class Tensor {
   /** @type {DataType} Type of the tensor. */
   type;
 
@@ -67,7 +67,7 @@ class Tensor {
   }
 }
 
-class InferenceSession {
+export class InferenceSession {
   sessionId;
   inputNames;
   outputNames;

ext/ai/lib.rs

@@ -55,7 +55,8 @@ deno_core::extension!(
     "util/event_stream_parser.mjs",
     "util/event_source_stream.mjs",
     "onnxruntime/onnx.js",
-    "onnxruntime/cache_adapter.js"
+    "onnxruntime/cache_adapter.js",
+    "onnxruntime/inference_api.js"
   ]
 );