[api] Adds load model from Inpustream to public API

Change-Id: I2106712d175ea1c6ec31d9561f2809473df2fbc9
deepjavalibrary · Dec 3, 2021 · cbaef03 · cbaef03
1 parent be5289d
commit cbaef03
Show file tree

Hide file tree

Showing 5 changed files with 50 additions and 10 deletions.
diff --git a/api/src/main/java/ai/djl/BaseModel.java b/api/src/main/java/ai/djl/BaseModel.java
@@ -116,6 +116,13 @@ public DataType getDataType() {
         return dataType;
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void load(InputStream is, Map<String, ?> options)
+            throws IOException, MalformedModelException {
+        throw new UnsupportedOperationException("Not supported!");
+    }
+
     /** {@inheritDoc} */
     @Override
     public void close() {

diff --git a/api/src/main/java/ai/djl/Model.java b/api/src/main/java/ai/djl/Model.java
@@ -133,6 +133,27 @@ default void load(Path modelPath, String prefix) throws IOException, MalformedMo
     void load(Path modelPath, String prefix, Map<String, ?> options)
             throws IOException, MalformedModelException;
 
+    /**
+     * Loads the model from the {@code InputStream}.
+     *
+     * @param is the {@code InputStream} to load the model from
+     * @throws IOException when IO operation fails in loading a resource
+     * @throws MalformedModelException if model file is corrupted
+     */
+    default void load(InputStream is) throws IOException, MalformedModelException {
+        load(is, null);
+    }
+
+    /**
+     * Loads the model from the {@code InputStream} with the options provided.
+     *
+     * @param is the {@code InputStream} to load the model from
+     * @param options engine specific load model options, see documentation for each engine
+     * @throws IOException when IO operation fails in loading a resource
+     * @throws MalformedModelException if model file is corrupted
+     */
+    void load(InputStream is, Map<String, ?> options) throws IOException, MalformedModelException;
+
     /**
      * Saves the model to the specified {@code modelPath} with the name provided.
      *

diff --git a/api/src/main/java/ai/djl/repository/zoo/ZooModel.java b/api/src/main/java/ai/djl/repository/zoo/ZooModel.java
@@ -59,6 +59,12 @@ public void load(Path modelPath, String prefix, Map<String, ?> options) {
         throw new IllegalArgumentException("ZooModel should not be re-loaded.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void load(InputStream modelStream, Map<String, ?> options) throws IOException {
+        throw new IllegalArgumentException("ZooModel should not be re-loaded.");
+    }
+
     /**
      * Returns the wrapped model.
      *

diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtModel.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtModel.java
@@ -108,16 +108,14 @@ public void load(Path modelPath, String prefix, Map<String, ?> options)
         }
     }
 
-    /**
-     * Load PyTorch model from {@link InputStream}.
-     *
-     * <p>Currently, only TorchScript file are supported
-     *
-     * @param modelStream the stream of the model file
-     * @throws IOException model loading error
-     */
-    public void load(InputStream modelStream) throws IOException {
-        load(modelStream, true);
+    /** {@inheritDoc} */
+    @Override
+    public void load(InputStream modelStream, Map<String, ?> options) throws IOException {
+        boolean mapLocation = true;
+        if (options != null && options.containsKey("mapLocation")) {
+            mapLocation = Boolean.parseBoolean(options.get("mapLocation").toString());
+        }
+        load(modelStream, mapLocation);
     }
 
     /**
@@ -128,6 +126,8 @@ public void load(InputStream modelStream) throws IOException {
      * @throws IOException model loading error
      */
     public void load(InputStream modelStream, boolean mapLocation) throws IOException {
+        modelDir = Files.createTempDirectory("pt-model");
+        modelDir.toFile().deleteOnExit();
         block = JniUtils.loadModule((PtNDManager) manager, modelStream, mapLocation, false);
     }
 

diff --git a/extensions/fasttext/src/main/java/ai/djl/fasttext/FtModel.java b/extensions/fasttext/src/main/java/ai/djl/fasttext/FtModel.java
@@ -93,6 +93,12 @@ public void load(Path modelPath, String prefix, Map<String, ?> options)
         properties.put("model-type", fta.getModelType());
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void load(InputStream modelStream, Map<String, ?> options) {
+        throw new UnsupportedOperationException("Not supported.");
+    }
+
     private Path findModelFile(String prefix) {
         if (Files.isRegularFile(modelDir)) {
             Path file = modelDir;