[examples] Fixes broken example (#3119)

deepjavalibrary · Apr 24, 2024 · a2f075f · a2f075f
1 parent 2409d8b
commit a2f075f
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/examples/src/main/java/ai/djl/examples/inference/BertTextEmbedding.java b/examples/src/main/java/ai/djl/examples/inference/BertTextEmbedding.java
@@ -26,6 +26,7 @@
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.List;
 
 /** An example of inference using Bert Text Embedding. */
 public final class BertTextEmbedding {
@@ -40,21 +41,21 @@ public static void main(String[] args) throws IOException, TranslateException, M
     }
 
     public static float[][] predict() throws IOException, TranslateException, ModelException {
-        String[] input = {"What is Deep Learning?", "The movie got Oscar this year"};
-        Criteria<String[], float[][]> criteria =
+        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)
+                        .setTypes(String.class, float[].class)
                         .optEngine("Rust")
                         .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                         .optProgress(new ProgressBar())
                         .build();
 
-        try (ZooModel<String[], float[][]> model = criteria.loadModel()) {
-            try (Predictor<String[], float[][]> predictor = model.newPredictor()) {
-                return predictor.predict(input);
-            }
+        try (ZooModel<String, float[]> model = criteria.loadModel();
+                Predictor<String, float[]> predictor = model.newPredictor()) {
+            return predictor.batchPredict(input).toArray(new float[0][]);
         }
     }
 }