[rust] Adds Rust model zoo

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) } }
+            }
         }
     }
 }

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"));
+        }
+    }
+}

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();
+    }
+}

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;

extensions/tokenizers/src/main/resources/META-INF/services/ai.djl.repository.zoo.ZooProvider

@@ -1 +1,2 @@
 ai.djl.huggingface.zoo.HfZooProvider
+ai.djl.engine.rust.zoo.RsZooProvider

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");
+        }
+    }
+}

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;