[onnxruntime] Load onnx extenstion if available

engines/onnxruntime/onnxruntime-engine/build.gradle.kts

@@ -1,5 +1,3 @@
-import java.net.URL
-
 plugins {
     ai.djl.javaProject
     ai.djl.publish
@@ -14,13 +12,14 @@ dependencies {
     testImplementation(project(":testing"))
     testImplementation(project(":engines:pytorch:pytorch-engine"))
     testImplementation(project(":extensions:tokenizers"))
+    testImplementation("com.microsoft.onnxruntime:onnxruntime-extensions:${libs.versions.onnxruntimeExtensions.get()}")
 
     testRuntimeOnly(libs.slf4j.simple)
 }
 
 tasks {
     processResources {
-        var basePath = "${project.projectDir}/build/resources/main/nlp"
+        val basePath = "${project.projectDir}/build/resources/main/nlp"
         outputs.dir(basePath)
         doLast {
             val url = "https://mlrepo.djl.ai/model/nlp"
@@ -34,7 +33,7 @@ tasks {
             for (task in tasks) {
                 val file = File("$basePath/$task/ai.djl.huggingface.onnxruntime.json")
                 if (file.exists())
-                    project.logger.lifecycle("model zoo metadata alrady exists: $task")
+                    project.logger.lifecycle("model zoo metadata already exists: $task")
                 else {
                     project.logger.lifecycle("Downloading model zoo metadata: $task")
                     file.parentFile.mkdirs()

engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java

@@ -18,6 +18,7 @@
 import ai.djl.Model;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.DataType;
+import ai.djl.util.ClassLoaderUtils;
 import ai.djl.util.Utils;
 import ai.onnxruntime.OrtEnvironment;
 import ai.onnxruntime.OrtException;
@@ -26,9 +27,13 @@
 import ai.onnxruntime.OrtSession.SessionOptions.ExecutionMode;
 import ai.onnxruntime.OrtSession.SessionOptions.OptLevel;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
+import java.lang.reflect.Method;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Map;
@@ -41,6 +46,8 @@
  */
 public class OrtModel extends BaseModel {
 
+    private static final Logger logger = LoggerFactory.getLogger(OrtModel.class);
+
     private OrtEnvironment env;
     private SessionOptions sessionOptions;
 
@@ -190,6 +197,9 @@ private SessionOptions getSessionOptions(Map<String, ?> options) throws OrtExcep
         }
 
         String customOpLibrary = (String) options.get("customOpLibrary");
+        if (customOpLibrary == null) {
+            customOpLibrary = getOrtxLibraryPath();
+        }
         if (customOpLibrary != null) {
             ortSession.registerCustomOpLibrary(customOpLibrary);
         }
@@ -223,4 +233,16 @@ private SessionOptions getSessionOptions(Map<String, ?> options) throws OrtExcep
         }
         return ortSession;
     }
+
+    private String getOrtxLibraryPath() {
+        ClassLoader cl = ClassLoaderUtils.getContextClassLoader();
+        try {
+            Class<?> clazz = Class.forName("ai.onnxruntime.extensions.OrtxPackage", true, cl);
+            Method method = clazz.getDeclaredMethod("getLibraryPath");
+            return (String) method.invoke(null);
+        } catch (Throwable e) {
+            logger.debug("Failed to load onnx extension", e);
+        }
+        return null;
+    }
 }

gradle/libs.versions.toml

@@ -8,6 +8,7 @@ tensorflow = "2.16.1"
 tensorflowCore = "1.0.0-rc.1"
 mxnet = "1.9.1"
 onnxruntime = "1.18.0"
+onnxruntimeExtensions = "0.11.0"
 xgboost = "2.0.3"
 lightgbm = "3.2.110"
 tensorrt = "8.4.1"