Repository: deepjavalibrary/djl
Pull request #3034 (2024-03-19 05:17:10)

[api] Support encode/decode String tensor
Fixes: #3033

## 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


codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3034?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: Patch coverage is `90.00000%` with `2 lines` in your changes are missing coverage. Please review.
> Project coverage is 72.24%. Comparing base [(`b6b8729`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/b6b8729920ee5b3170939be71f057e7293d29f4d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) to head [(`7f786dd`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3034?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> Report is 1 commits behind head on master.

| [Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3034?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Patch % | Lines |
|---|---|---|
| [...src/main/java/ai/djl/pytorch/engine/PtNDArray.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3034?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy9weXRvcmNoL3B5dG9yY2gtZW5naW5lL3NyYy9tYWluL2phdmEvYWkvZGpsL3B5dG9yY2gvZW5naW5lL1B0TkRBcnJheS5qYXZh) | 60.00% | [1 Missing and 1 partial :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3034?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #3034      +/-   ##
============================================
- Coverage     72.29%   72.24%   -0.06%     
- Complexity     7293     7295       +2     
============================================
  Files           722      722              
  Lines         32512    32533      +21     
  Branches       3396     3401       +5     
============================================
- Hits          23506    23504       -2     
- Misses         7386     7409      +23     
  Partials       1620     1620              
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3034?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/ndarray/NDSerializer.java b/api/src/main/java/ai/djl/ndarray/NDSerializer.java
index f42be7f08..26dde5d92 100644
--- a/api/src/main/java/ai/djl/ndarray/NDSerializer.java
+++ b/api/src/main/java/ai/djl/ndarray/NDSerializer.java
@@ -83,6 +83,16 @@ final class NDSerializer {
         Shape shape = array.getShape();
         dos.write(shape.getEncoded());
 
+        if (array.getDataType() == DataType.STRING) {
+            String[] data = array.toStringArray();
+            dos.writeInt(data.length);
+            for (String str : data) {
+                dos.writeUTF(str);
+            }
+            dos.flush();
+            return;
+        }
+
         ByteBuffer bb = array.toByteBuffer();
         dos.write(bb.order() == ByteOrder.BIG_ENDIAN ? '>' : '<');
         int length = bb.remaining();
@@ -167,6 +177,17 @@ final class NDSerializer {
         // Shape
         Shape shape = Shape.decode(bb);
 
+        if (dataType == DataType.STRING) {
+            int size = bb.getInt();
+            String[] data = new String[size];
+            for (int i = 0; i < size; ++i) {
+                data[i] = readUTF(bb);
+            }
+            NDArray array = manager.create(data, StandardCharsets.UTF_8, shape);
+            array.setName(name);
+            return array;
+        }
+
         // Data
         ByteOrder order;
         if (version > 2) {
diff --git a/bom/build.gradle b/bom/build.gradle
index 313173161..0c509740f 100644
--- a/bom/build.gradle
+++ b/bom/build.gradle
@@ -19,7 +19,6 @@ dependencies {
         api "ai.djl:basicdataset:${version}"
         api "ai.djl:model-zoo:${version}"
         api "ai.djl:djl-zero:${version}"
-        api "ai.djl:serving:${version}"
         api "ai.djl.android:core:${version}"
         api "ai.djl.android:onnxruntime:${version}"
         api "ai.djl.android:pytorch-native:${version}"
@@ -44,6 +43,9 @@ dependencies {
         api "ai.djl.pytorch:pytorch-model-zoo:${version}"
         api "ai.djl.sentencepiece:sentencepiece:${version}"
         api "ai.djl.spark:spark_2.12:${version}"
+        api "ai.djl.serving:prometheus:${version}"
+        api "ai.djl.serving:serving:${version}"
+        api "ai.djl.serving:wlm:${version}"
         api "ai.djl.tablesaw:tablesaw:${version}"
         api "ai.djl.tensorflow:tensorflow-api:${version}"
         api "ai.djl.tensorflow:tensorflow-engine:${version}"
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
index 499f51eba..551a16d03 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
@@ -100,6 +100,7 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
         super(-1L);
         this.manager = manager;
         this.strs = strs;
+        this.sparseFormat = SparseFormat.DENSE;
         this.shape = shape;
         this.dataType = DataType.STRING;
         NDScope.register(this);
@@ -225,6 +226,10 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
     /** {@inheritDoc} */
     @Override
     public ByteBuffer toByteBuffer() {
+        if (getDataType() == DataType.STRING) {
+            throw new UnsupportedOperationException(
+                    "toByteBuffer is not supported for String tensor.");
+        }
         return JniUtils.getByteBuffer(this);
     }
 
@@ -429,6 +434,9 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
         if (getDataType() != other.getDataType()) {
             return false;
         }
+        if (getDataType() == DataType.STRING) {
+            return Arrays.equals(toStringArray(), other.toStringArray());
+        }
         return JniUtils.contentEqual(this, manager.from(other));
     }
 


Test patch:
diff --git a/api/src/test/java/ai/djl/ndarray/NDSerializerTest.java b/api/src/test/java/ai/djl/ndarray/NDSerializerTest.java
index e89f22442..98ba89688 100644
--- a/api/src/test/java/ai/djl/ndarray/NDSerializerTest.java
+++ b/api/src/test/java/ai/djl/ndarray/NDSerializerTest.java
@@ -97,6 +97,23 @@ public class NDSerializerTest {
         }
     }
 
+    @Test
+    public void testStringTensor() {
+        try (NDManager manager = NDManager.newBaseManager("PyTorch")) {
+            NDArray array = manager.create("hello");
+            byte[] buf = array.encode();
+            NDArray decoded = NDArray.decode(manager, buf);
+            Assert.assertTrue(decoded.getShape().isScalar());
+
+            array = manager.create(new String[] {"hello", "world"});
+            buf = array.encode();
+            decoded = NDArray.decode(manager, buf);
+            Assert.assertEquals(decoded.getShape(), array.getShape());
+            Assert.assertEquals(decoded.toStringArray()[1], "world");
+            Assert.assertEquals(decoded, array);
+        }
+    }
+
     private static byte[] encode(NDArray array) throws IOException {
         try (ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
             NDSerializer.encodeAsNumpy(array, bos);