Repository: deepjavalibrary/djl
Pull request #3138 (2024-04-29 04:20:36)

Ensure the alternative ND manager can use GPUs
This makes a massive performance difference when using OnnxRuntime, which delegates to PyTorch for some NDArray operations.  This almost makes OnnxRuntime run as fast compared to PyTorch.  Without this change, OnnxRuntime was over 3 times slower.



Patch:
diff --git a/api/src/main/java/ai/djl/ndarray/BaseNDManager.java b/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
index c3df1ef33..726ecf757 100644
--- a/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
@@ -62,7 +62,8 @@ public abstract class BaseNDManager implements NDManager {
         uid = UUID.randomUUID().toString();
         Engine engine = getEngine().getAlternativeEngine();
         if (engine != null) {
-            alternativeManager = engine.newBaseManager(Device.cpu());
+            // Use the default device
+            alternativeManager = engine.newBaseManager();
         }
     }
 
diff --git a/examples/src/main/java/ai/djl/examples/inference/BertTextEmbedding.java b/examples/src/main/java/ai/djl/examples/inference/BertTextEmbedding.java
deleted file mode 100644
index 387b1eed5..000000000
--- a/examples/src/main/java/ai/djl/examples/inference/BertTextEmbedding.java
+++ /dev/null
@@ -1,61 +0,0 @@
-/*
- * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
- * with the License. A copy of the License is located at
- *
- * http://aws.amazon.com/apache2.0/
- *
- * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
- * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
- * and limitations under the License.
- */
-
-package ai.djl.examples.inference;
-
-import ai.djl.ModelException;
-import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
-import ai.djl.inference.Predictor;
-import ai.djl.repository.zoo.Criteria;
-import ai.djl.repository.zoo.ZooModel;
-import ai.djl.training.util.ProgressBar;
-import ai.djl.translate.TranslateException;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.util.Arrays;
-import java.util.List;
-
-/** An example of inference using Bert Text Embedding. */
-public final class BertTextEmbedding {
-
-    private static final Logger logger = LoggerFactory.getLogger(BertTextEmbedding.class);
-
-    private BertTextEmbedding() {}
-
-    public static void main(String[] args) throws IOException, TranslateException, ModelException {
-        float[][] output = BertTextEmbedding.predict();
-        logger.info("Output: {}", Arrays.deepToString(output));
-    }
-
-    public static float[][] predict() throws IOException, TranslateException, ModelException {
-        List<String> input =
-                Arrays.asList("What is Deep Learning?", "The movie got Oscar this year");
-        Criteria<String, float[]> criteria =
-                Criteria.builder()
-                        .optModelUrls(
-                                "https://alpha-djl-demos.s3.amazonaws.com/model/examples/bge-base-en-v1.5.zip")
-                        .setTypes(String.class, float[].class)
-                        .optEngine("Rust")
-                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
-                        .optProgress(new ProgressBar())
-                        .build();
-
-        try (ZooModel<String, float[]> model = criteria.loadModel();
-                Predictor<String, float[]> predictor = model.newPredictor()) {
-            return predictor.batchPredict(input).toArray(new float[0][]);
-        }
-    }
-}
diff --git a/extensions/tokenizers/build.gradle b/extensions/tokenizers/build.gradle
index 39163a409..0fae92696 100644
--- a/extensions/tokenizers/build.gradle
+++ b/extensions/tokenizers/build.gradle
@@ -59,13 +59,27 @@ processResources {
         for (String task : tasks) {
             def file = new File("${prefix}/${task}/ai.djl.huggingface.pytorch.json")
             if (file.exists()) {
-                project.logger.lifecycle("model zoo metadata alrady exists: ${task}")
+                project.logger.lifecycle("PyTorch model zoo metadata alrady exists: ${task}")
             } else {
-                project.logger.lifecycle("Downloading model zoo metadata: ${task}")
+                project.logger.lifecycle("Downloading PyTorch model zoo metadata: ${task}")
                 file.getParentFile().mkdirs()
                 def downloadPath = new URL("${url}/${task}/ai/djl/huggingface/pytorch/models.json.gz")
                 downloadPath.withInputStream { i -> file.withOutputStream { it << new GZIPInputStream(i) } }
             }
+
+            if ("text_embedding" != task) {
+                continue
+            }
+
+            file = new File("${prefix}/${task}/ai.djl.huggingface.rust.json")
+            if (file.exists()) {
+                project.logger.lifecycle("Rust model zoo metadata alrady exists: ${task}")
+            } else {
+                project.logger.lifecycle("Downloading Rust model zoo metadata: ${task}")
+                file.getParentFile().mkdirs()
+                def downloadPath = new URL("${url}/${task}/ai/djl/huggingface/rust/models.json.gz")
+                downloadPath.withInputStream { i -> file.withOutputStream { it << new GZIPInputStream(i) } }
+            }
         }
     }
 }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/zoo/RsModelZoo.java b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/zoo/RsModelZoo.java
new file mode 100644
index 000000000..f706bc865
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/zoo/RsModelZoo.java
@@ -0,0 +1,96 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.engine.rust.zoo;
+
+import ai.djl.Application;
+import ai.djl.engine.Engine;
+import ai.djl.repository.Repository;
+import ai.djl.repository.Version;
+import ai.djl.repository.VersionRange;
+import ai.djl.repository.zoo.ModelLoader;
+import ai.djl.repository.zoo.ModelZoo;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+
+/** RsModelZoo is a repository that contains HuggingFace models. */
+public class RsModelZoo extends ModelZoo {
+
+    private static final String REPO = "https://mlrepo.djl.ai/";
+    private static final Repository REPOSITORY = Repository.newInstance("Rust", REPO);
+    private static final String GROUP_ID = "ai.djl.huggingface.rust";
+
+    private volatile boolean initialized; // NOPMD
+
+    RsModelZoo() {}
+
+    /** {@inheritDoc} */
+    @Override
+    public String getGroupId() {
+        return GROUP_ID;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Set<String> getSupportedEngines() {
+        return Collections.singleton("Rust");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Collection<ModelLoader> getModelLoaders() {
+        init();
+        return super.getModelLoaders();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public ModelLoader getModelLoader(String name) {
+        init();
+        return super.getModelLoader(name);
+    }
+
+    private void init() {
+        if (!initialized) {
+            synchronized (RsModelZoo.class) {
+                if (!initialized) {
+                    Version version = new Version(Engine.getDjlVersion());
+                    addModels(Application.NLP.TEXT_EMBEDDING, version);
+                    initialized = true;
+                }
+            }
+        }
+    }
+
+    private void addModels(Application app, Version version) {
+        Map<String, Map<String, Object>> map = listModels(REPOSITORY, app);
+        for (Map.Entry<String, Map<String, Object>> entry : map.entrySet()) {
+            Map<String, Object> model = entry.getValue();
+            if ("failed".equals(model.get("result"))) {
+                continue;
+            }
+            String requires = (String) model.get("requires");
+            if (requires != null) {
+                // the model requires specific DJL version
+                VersionRange range = VersionRange.parse(requires);
+                if (!range.contains(version)) {
+                    continue;
+                }
+            }
+            String artifactId = entry.getKey();
+            addModel(REPOSITORY.model(app, GROUP_ID, artifactId, "0.0.1"));
+        }
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/zoo/RsZooProvider.java b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/zoo/RsZooProvider.java
new file mode 100644
index 000000000..68604566c
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/zoo/RsZooProvider.java
@@ -0,0 +1,28 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.engine.rust.zoo;
+
+import ai.djl.repository.zoo.ModelZoo;
+import ai.djl.repository.zoo.ZooProvider;
+
+/**
+ * An Rust model zoo provider implements the {@link ai.djl.repository.zoo.ZooProvider} interface.
+ */
+public class RsZooProvider implements ZooProvider {
+
+    /** {@inheritDoc} */
+    @Override
+    public ModelZoo getModelZoo() {
+        return new RsModelZoo();
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/zoo/package-info.java b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/zoo/package-info.java
new file mode 100644
index 000000000..8349163ee
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/zoo/package-info.java
@@ -0,0 +1,14 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+/** Contains the built-in {@link ai.djl.huggingface.zoo.RsModelZoo}. */
+package ai.djl.engine.rust.zoo;
diff --git a/extensions/tokenizers/src/main/python/huggingface_converter.py b/extensions/tokenizers/src/main/python/huggingface_converter.py
index c856d1880..167b4ad81 100644
--- a/extensions/tokenizers/src/main/python/huggingface_converter.py
+++ b/extensions/tokenizers/src/main/python/huggingface_converter.py
@@ -100,6 +100,8 @@ class HuggingfaceConverter:
                 include_types = False
             else:
                 return False, f"Unsupported model_type: {config.model_type}", -1
+        else:
+            return False, f"Unknown model_type: {model_id}", -1
 
         logging.info(f"Saving rust model: {model_id} ...")
 
@@ -116,6 +118,10 @@ class HuggingfaceConverter:
             logging.warning(e, exc_info=True)
             return False, "Failed to save tokenizer", -1
 
+        # Save config.json
+        config_file = hf_hub_download(repo_id=model_id, filename="config.json")
+        shutil.copyfile(config_file, os.path.join(temp_dir, "config.json"))
+
         target = os.path.join(temp_dir, "model.safetensors")
         model = self.api.model_info(model_id, files_metadata=True)
         has_sf_file = False
@@ -188,14 +194,17 @@ class HuggingfaceConverter:
             if not result:
                 return False, reason, -1
 
-        size = self.save_to_model_zoo(model_info, args.output_dir, temp_dir,
-                                      hf_pipeline, include_types)
+        size = self.save_to_model_zoo(model_info, args.output_dir, "PyTorch",
+                                      temp_dir, hf_pipeline, include_types)
 
         return True, None, size
 
     @staticmethod
     def save_tokenizer(hf_pipeline, temp_dir: str):
         hf_pipeline.tokenizer.save_pretrained(temp_dir)
+        if not os.path.exists(os.path.join(temp_dir, "tokenizer.json")):
+            raise ValueError("no fast tokenizer found.")
+
         # only keep tokenizer.json file
         for path in os.listdir(temp_dir):
             if path != "tokenizer.json" and path != "tokenizer_config.json":
diff --git a/extensions/tokenizers/src/main/resources/META-INF/services/ai.djl.repository.zoo.ZooProvider b/extensions/tokenizers/src/main/resources/META-INF/services/ai.djl.repository.zoo.ZooProvider
index 956e00dcf..b9708f05e 100644
--- a/extensions/tokenizers/src/main/resources/META-INF/services/ai.djl.repository.zoo.ZooProvider
+++ b/extensions/tokenizers/src/main/resources/META-INF/services/ai.djl.repository.zoo.ZooProvider
@@ -1,1 +1,2 @@
 ai.djl.huggingface.zoo.HfZooProvider
+ai.djl.engine.rust.zoo.RsZooProvider


Test patch:
diff --git a/extensions/tokenizers/src/test/java/ai/djl/engine/rust/zoo/RsModelZooTest.java b/extensions/tokenizers/src/test/java/ai/djl/engine/rust/zoo/RsModelZooTest.java
new file mode 100644
index 000000000..3c0624173
--- /dev/null
+++ b/extensions/tokenizers/src/test/java/ai/djl/engine/rust/zoo/RsModelZooTest.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.engine.rust.zoo;
+
+import ai.djl.ModelException;
+import ai.djl.inference.Predictor;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ModelZoo;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.testing.TestRequirements;
+import ai.djl.translate.TranslateException;
+import ai.djl.util.Utils;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.nio.file.Paths;
+import java.util.Set;
+
+public class RsModelZooTest {
+
+    @Test
+    public void testRsModelZoo() throws ModelException, IOException, TranslateException {
+        TestRequirements.nightly();
+
+        String text = "What is deep learning?";
+        String url = "djl://ai.djl.huggingface.rust/TaylorAI/bge-micro-v2";
+        Criteria<String, float[]> criteria =
+                Criteria.builder().setTypes(String.class, float[].class).optModelUrls(url).build();
+
+        try (ZooModel<String, float[]> model = criteria.loadModel();
+                Predictor<String, float[]> predictor = model.newPredictor()) {
+            float[] res = predictor.predict(text);
+            Assert.assertEquals(res.length, 384);
+        }
+    }
+
+    @Test
+    public void testOffLine() {
+        System.setProperty("DJL_CACHE_DIR", "build/cache");
+        System.setProperty("ai.djl.offline", "true");
+        try {
+            Utils.deleteQuietly(Paths.get("build/cache"));
+            // static variables cannot not be initialized properly if directly use new HfModelZoo()
+            ModelZoo.getModelZoo("ai.djl.huggingface.rust");
+
+            ModelZoo zoo = new RsModelZoo();
+            Assert.assertFalse(zoo.getModelLoaders().isEmpty());
+
+            Set<String> engines = zoo.getSupportedEngines();
+            Assert.assertEquals(engines.size(), 1);
+            Assert.assertEquals(engines.iterator().next(), "Rust");
+        } finally {
+            System.clearProperty("DJL_CACHE_DIR");
+            System.clearProperty("ai.djl.offline");
+        }
+    }
+}
diff --git a/extensions/tokenizers/src/test/java/ai/djl/engine/rust/zoo/package-info.java b/extensions/tokenizers/src/test/java/ai/djl/engine/rust/zoo/package-info.java
new file mode 100644
index 000000000..ce2cfde67
--- /dev/null
+++ b/extensions/tokenizers/src/test/java/ai/djl/engine/rust/zoo/package-info.java
@@ -0,0 +1,13 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.engine.rust.zoo;