Expose OnnxRuntime's getMetadata() in DJL API (#3596)

* Expose OnnxRuntime getMetadata() in DJL API * Revert "Expose OnnxRuntime getMetadata() in DJL API" This reverts commit 738a329. * Expose OnnxRuntime getMetadata() in DJL API --------- Co-authored-by: Vanja Radulovic <[email protected]>
deepjavalibrary · Jan 28, 2025 · 2ad90ea · 2ad90ea
1 parent d381c4d
commit 2ad90ea
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/api/src/main/java/ai/djl/nn/Block.java b/api/src/main/java/ai/djl/nn/Block.java
@@ -26,6 +26,8 @@
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
 import java.util.function.Predicate;
 
 /**
@@ -353,4 +355,13 @@ static void validateLayout(LayoutType[] expectedLayout, LayoutType[] actualLayou
             }
         }
     }
+
+    /**
+     * Returns a map of all the custom metadata of the block.
+     *
+     * @return the map of {@link PairList}
+     */
+    default Map<String, String> getCustomMetadata() {
+        return Collections.emptyMap();
+    }
 }
diff --git a/...nnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtSymbolBlock.java b/...nnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtSymbolBlock.java
@@ -26,6 +26,7 @@
 import ai.djl.util.PairList;
 import ai.onnxruntime.OnnxJavaType;
 import ai.onnxruntime.OnnxMap;
+import ai.onnxruntime.OnnxModelMetadata;
 import ai.onnxruntime.OnnxSequence;
 import ai.onnxruntime.OnnxTensor;
 import ai.onnxruntime.OnnxValue;
@@ -128,6 +129,17 @@ public PairList<String, Shape> describeInput() {
         return result;
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public Map<String, String> getCustomMetadata() {
+        try {
+            OnnxModelMetadata modelMetadata = session.getMetadata();
+            return modelMetadata.getCustomMetadata();
+        } catch (OrtException e) {
+            throw new EngineException(e);
+        }
+    }
+
     private NDList evaluateOutput(OrtSession.Result results) {
         NDList output = new NDList();
         for (Map.Entry<String, OnnxValue> r : results) {