Repository: deepjavalibrary/djl
Pull request #3132 (2024-04-26 04:48:19)

[rust] Adds Rust model zoo
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/api/src/main/java/ai/djl/ndarray/types/Shape.java b/api/src/main/java/ai/djl/ndarray/types/Shape.java
index e7751f7d5..51b4d2ad6 100644
--- a/api/src/main/java/ai/djl/ndarray/types/Shape.java
+++ b/api/src/main/java/ai/djl/ndarray/types/Shape.java
@@ -23,6 +23,8 @@ import java.util.Arrays;
 import java.util.List;
 import java.util.function.Function;
 import java.util.function.Predicate;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.LongStream;
 import java.util.stream.Stream;
@@ -535,4 +537,64 @@ public class Shape {
         }
         return max == ans;
     }
+
+    /**
+     * Parses a string representation of shapes for NDList.
+     *
+     * @param value a string representation of shapes for NDList
+     * @return a list of Shape and datatype pairs
+     */
+    public static PairList<DataType, Shape> parseShapes(String value) {
+        PairList<DataType, Shape> inputShapes = new PairList<>();
+        if (value != null) {
+            if (value.contains("(")) {
+                Pattern pattern = Pattern.compile("\\((\\s*(\\d+)([,\\s]+\\d+)*\\s*)\\)(\\w?)");
+                Matcher matcher = pattern.matcher(value);
+                while (matcher.find()) {
+                    String[] tokens = matcher.group(1).split(",");
+                    long[] array = Arrays.stream(tokens).mapToLong(Long::parseLong).toArray();
+                    DataType dataType;
+                    String dataTypeStr = matcher.group(4);
+                    if (dataTypeStr == null || dataTypeStr.isEmpty()) {
+                        dataType = DataType.FLOAT32;
+                    } else {
+                        switch (dataTypeStr) {
+                            case "s":
+                                dataType = DataType.FLOAT16;
+                                break;
+                            case "d":
+                                dataType = DataType.FLOAT64;
+                                break;
+                            case "u":
+                                dataType = DataType.UINT8;
+                                break;
+                            case "b":
+                                dataType = DataType.INT8;
+                                break;
+                            case "i":
+                                dataType = DataType.INT32;
+                                break;
+                            case "l":
+                                dataType = DataType.INT64;
+                                break;
+                            case "B":
+                                dataType = DataType.BOOLEAN;
+                                break;
+                            case "f":
+                                dataType = DataType.FLOAT32;
+                                break;
+                            default:
+                                throw new IllegalArgumentException("Invalid input-shape: " + value);
+                        }
+                    }
+                    inputShapes.add(dataType, new Shape(array));
+                }
+            } else {
+                String[] tokens = value.split(",");
+                long[] shapes = Arrays.stream(tokens).mapToLong(Long::parseLong).toArray();
+                inputShapes.add(DataType.FLOAT32, new Shape(shapes));
+            }
+        }
+        return inputShapes;
+    }
 }
diff --git a/api/src/main/java/ai/djl/nn/Blocks.java b/api/src/main/java/ai/djl/nn/Blocks.java
index 47cd58435..136c9bd0f 100644
--- a/api/src/main/java/ai/djl/nn/Blocks.java
+++ b/api/src/main/java/ai/djl/nn/Blocks.java
@@ -14,8 +14,11 @@ package ai.djl.nn;
 
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.util.Pair;
+import ai.djl.util.PairList;
 
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -87,6 +90,31 @@ public final class Blocks {
         return new LambdaBlock(x -> x, "identity");
     }
 
+    /**
+     * Creates a {@link LambdaBlock} that return all-ones NDList.
+     *
+     * @return an all-ones {@link Block}
+     */
+    public static Block onesBlock(PairList<DataType, Shape> shapes, String[] names) {
+        return new LambdaBlock(
+                a -> {
+                    NDManager manager = a.getManager();
+                    NDList list = new NDList(shapes.size());
+                    int index = 0;
+                    for (Pair<DataType, Shape> pair : shapes) {
+                        DataType dataType = pair.getKey();
+                        Shape shape = pair.getValue();
+                        NDArray arr = manager.ones(shape, dataType);
+                        if (names.length == list.size()) {
+                            arr.setName(names[index++]);
+                        }
+                        list.add(arr);
+                    }
+                    return list;
+                },
+                "ones");
+    }
+
     /**
      * Returns a string representation of the passed {@link Block} describing the input axes, output
      * axes, and the block's children.
diff --git a/api/src/main/java/ai/djl/nn/OnesBlockFactory.java b/api/src/main/java/ai/djl/nn/OnesBlockFactory.java
new file mode 100644
index 000000000..55f17269f
--- /dev/null
+++ b/api/src/main/java/ai/djl/nn/OnesBlockFactory.java
@@ -0,0 +1,45 @@
+/*
+ * 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.nn;
+
+import ai.djl.Model;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.util.PairList;
+import ai.djl.util.Utils;
+
+import java.nio.file.Path;
+import java.util.Map;
+
+/** A {@link BlockFactory} class that creates LambdaBlock. */
+public class OnesBlockFactory implements BlockFactory {
+
+    private static final long serialVersionUID = 1L;
+
+    /** {@inheritDoc} */
+    @Override
+    public Block newBlock(Model model, Path modelPath, Map<String, ?> arguments) {
+        String shapes = ArgumentsUtil.stringValue(arguments, "block_shapes");
+        String blockNames = ArgumentsUtil.stringValue(arguments, "block_names");
+        PairList<DataType, Shape> pairs = Shape.parseShapes(shapes);
+        String[] names;
+        if (blockNames != null) {
+            names = blockNames.split(",");
+        } else {
+            names = Utils.EMPTY_ARRAY;
+        }
+
+        return Blocks.onesBlock(pairs, names);
+    }
+}
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/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/api/src/test/java/ai/djl/nn/BlockFactoryTest.java b/api/src/test/java/ai/djl/nn/BlockFactoryTest.java
index 230631486..54502d171 100644
--- a/api/src/test/java/ai/djl/nn/BlockFactoryTest.java
+++ b/api/src/test/java/ai/djl/nn/BlockFactoryTest.java
@@ -13,12 +13,19 @@
 package ai.djl.nn;
 
 import ai.djl.Model;
+import ai.djl.ModelException;
+import ai.djl.ndarray.NDList;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import java.io.IOException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 @SuppressWarnings("PMD.TestClassWithoutTestCases")
 public class BlockFactoryTest {
@@ -32,4 +39,35 @@ public class BlockFactoryTest {
             Assert.assertEquals(((LambdaBlock) block).getName(), "identity");
         }
     }
+
+    @Test
+    public void testOnesBlockFactory() throws ModelException, IOException {
+        OnesBlockFactory factory = new OnesBlockFactory();
+        Path path = Paths.get("build");
+        Criteria<NDList, NDList> criteria =
+                Criteria.builder()
+                        .setTypes(NDList.class, NDList.class)
+                        .optModelPath(path)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1)s,(1)d,(1)u,(1)b,(1)i,(1)l,(1)B,(1)f,(1)")
+                        .optArgument("block_names", "1,2,3,4,5,6,7,8,9")
+                        .optOption("hasParameter", "false")
+                        .build();
+
+        try (ZooModel<NDList, NDList> model = criteria.loadModel()) {
+            Block block = model.getBlock();
+            Assert.assertTrue(block instanceof LambdaBlock);
+
+            Map<String, String> args = new ConcurrentHashMap<>();
+            args.put("block_shapes", "1,2");
+            block = factory.newBlock(model, path, args);
+            Assert.assertTrue(block instanceof LambdaBlock);
+
+            args.put("block_shapes", "(1)a");
+            Assert.assertThrows(
+                    () -> {
+                        factory.newBlock(model, path, args);
+                    });
+        }
+    }
 }
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;
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
index 476eed24c..2712dc424 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
@@ -21,14 +21,18 @@ import ai.djl.modality.Output;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.Block;
+import ai.djl.nn.Blocks;
 import ai.djl.nn.LambdaBlock;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.testing.Assertions;
 import ai.djl.translate.TranslateException;
 import ai.djl.util.JsonUtils;
+import ai.djl.util.PairList;
+import ai.djl.util.Utils;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -50,15 +54,6 @@ public class TextEmbeddingTranslatorTest {
             throws ModelException, IOException, TranslateException {
         String text = "This is an example sentence";
 
-        Block block =
-                new LambdaBlock(
-                        a -> {
-                            NDManager manager = a.getManager();
-                            NDArray arr = manager.ones(new Shape(1, 7, 384));
-                            arr.setName("last_hidden_state");
-                            return new NDList(arr);
-                        },
-                        "model");
         Path modelDir = Paths.get("build/model");
         Files.createDirectories(modelDir);
         try (NDManager manager = NDManager.newBaseManager("Rust")) {
@@ -84,8 +79,10 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(String.class, float[].class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
                         .optEngine("PyTorch")
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optOption("hasParameter", "false")
                         .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
@@ -103,7 +100,9 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(String.class, float[].class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("pooling", "max")
@@ -123,7 +122,9 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(String.class, float[].class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("pooling", "mean_sqrt_len")
@@ -143,7 +144,9 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(String.class, float[].class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("pooling", "weightedmean")
@@ -163,7 +166,9 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(String.class, float[].class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("dense", "linear.safetensors")
@@ -183,7 +188,9 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(Input.class, Output.class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("pooling", "cls")
@@ -212,7 +219,9 @@ public class TextEmbeddingTranslatorTest {
         }
 
         try (Model model = Model.newInstance("test")) {
-            model.setBlock(block);
+            PairList<DataType, Shape> pairs = new PairList<>();
+            pairs.add(DataType.FLOAT32, new Shape(1, 7, 384));
+            model.setBlock(Blocks.onesBlock(pairs, Utils.EMPTY_ARRAY));
             Map<String, String> options = new HashMap<>();
             options.put("hasParameter", "false");
             model.load(modelDir, "test", options);