Repository: deepjavalibrary/djl
Pull request #1968 (2022-08-26 19:38:41)

Add string tensor support for PyTorch
## Description ##

Brief description of what this PR is about

Add string tensor support for PyTorch

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1968?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1968](https://codecov.io/gh/deepjavalibrary/djl/pull/1968?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (d4bb90d) into [master](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (bb5073f) will **decrease** coverage by `2.03%`.
> The diff coverage is `67.91%`.

```diff
@@             Coverage Diff              @@
##             master    #1968      +/-   ##
============================================
- Coverage     72.08%   70.05%   -2.04%     
- Complexity     5126     5827     +701     
============================================
  Files           473      574     +101     
  Lines         21970    25854    +3884     
  Branches       2351     2786     +435     
============================================
+ Hits          15838    18111    +2273     
- Misses         4925     6375    +1450     
- Partials       1207     1368     +161     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1968?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1968/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1968/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1968/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1968/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...main/java/ai/djl/modality/cv/output/Rectangle.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1968/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9SZWN0YW5nbGUuamF2YQ==) | `72.41% <0.00%> (ø)` | |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1968/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [...odality/cv/translator/BigGANTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1968/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `33.33% <0.00%> (+8.33%)` | :arrow_up: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1968/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `14.28% <0.00%> (-3.90%)` | :arrow_down: |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1968/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU2VtYW50aWNTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../cv/translator/StyleTransferTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1968/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU3R5bGVUcmFuc2ZlclRyYW5zbGF0b3JGYWN0b3J5LmphdmE=) | `40.00% <ø> (ø)` | |
| ... and [476 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1968/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)


Patch:
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 562d1d38c..d61ecad3f 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
@@ -46,6 +46,7 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
     private Boolean hasGradient;
     private PtNDManager manager;
     private PtNDArrayEx ptNDArrayEx;
+    private String[] strs;
 
     // keep a reference to direct buffer to avoid GC release the memory
     @SuppressWarnings("PMD.UnusedPrivateField")
@@ -81,6 +82,22 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
         dataRef = data;
     }
 
+    /**
+     * Constructs a PyTorch {@code NDArray} to hold string array with a dummy native handle
+     * (internal. Use {@link NDManager} instead) with the data that is hold on Java side.
+     *
+     * @param manager the manager to attach the new array to
+     * @param strs the string array
+     * @param shape the {@link Shape} of the {@link NDArray}
+     */
+    public PtNDArray(PtNDManager manager, String[] strs, Shape shape) {
+        super(-1L);
+        this.manager = manager;
+        this.strs = strs;
+        this.shape = shape;
+        this.dataType = DataType.STRING;
+    }
+
     /** {@inheritDoc} */
     @Override
     public PtNDManager getManager() {
@@ -203,7 +220,7 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
     /** {@inheritDoc} */
     @Override
     public String[] toStringArray(Charset charset) {
-        throw new UnsupportedOperationException("String NDArray is not supported!");
+        return strs;
     }
 
     /** {@inheritDoc} */
@@ -1477,7 +1494,7 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
     @Override
     public void close() {
         Long pointer = handle.getAndSet(null);
-        if (pointer != null) {
+        if (pointer != null && pointer != -1) {
             JniUtils.deleteNDArray(pointer);
         }
         manager.detachInternal(getUid());
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java
index 7c4844542..307edb414 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDManager.java
@@ -26,6 +26,7 @@ import ai.djl.pytorch.jni.JniUtils;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.nio.charset.Charset;
 
 /** {@code PtNDManager} is the PyTorch implementation of {@link NDManager}. */
 public class PtNDManager extends BaseNDManager {
@@ -76,6 +77,12 @@ public class PtNDManager extends BaseNDManager {
                 this, buf, shape, dataType, SparseFormat.DENSE, device);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray create(String[] data, Charset charset, Shape shape) {
+        return new PtNDArray(this, data, shape);
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray createCoo(Buffer data, long[][] indices, Shape shape) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValue.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValue.java
index 79cb3f577..b1c34688c 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValue.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValue.java
@@ -13,6 +13,8 @@
 package ai.djl.pytorch.jni;
 
 import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
 import ai.djl.pytorch.engine.PtNDArray;
 import ai.djl.pytorch.engine.PtNDManager;
 import ai.djl.util.NativeResource;
@@ -156,11 +158,26 @@ public class IValue extends NativeResource<Long> {
     /**
      * Creates a new {@code IValue} of type {@code PtNDArray}.
      *
-     * @param value the NDArray value
+     * @param array the NDArray
      * @return a new {@code IValue} of type {@code PtNDArray}
      */
-    public static IValue from(PtNDArray value) {
-        return new IValue(PyTorchLibrary.LIB.iValueFromTensor(value.getHandle()));
+    public static IValue from(PtNDArray array) {
+        if (array.getDataType() == DataType.STRING) {
+            Shape shape = array.getShape();
+            String[] strs = array.toStringArray();
+            if (shape.isScalar()) {
+                return from(strs[0]);
+            }
+            IValue[] list = new IValue[strs.length];
+            PtNDManager manager = array.getManager();
+            for (int i = 0; i < strs.length; i++) {
+                IValue ivalue = from(strs[i]);
+                manager.attachUncappedInternal(ivalue.getUid(), ivalue);
+                list[i] = ivalue;
+            }
+            return listFrom(list);
+        }
+        return new IValue(PyTorchLibrary.LIB.iValueFromTensor(array.getHandle()));
     }
 
     /**
@@ -477,6 +494,8 @@ public class IValue extends NativeResource<Long> {
                 ivalue.close();
             }
             return list;
+        } else if (isString()) {
+            return new NDList(manager.create(toStringValue()));
         }
         throw new UnsupportedOperationException("Unsupported IValue type.");
     }


Test patch:
diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java
index 4caca064f..286ac0200 100644
--- a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java
@@ -40,6 +40,8 @@ public class IValueTest {
         try (PtNDManager manager = (PtNDManager) NDManager.newBaseManager()) {
             PtNDArray array1 = (PtNDArray) manager.zeros(new Shape(1));
             PtNDArray array2 = (PtNDArray) manager.ones(new Shape(1));
+            PtNDArray array3 = (PtNDArray) manager.create("test");
+            PtNDArray array4 = (PtNDArray) manager.create(new String[] {"test1", "test2"});
 
             try (IValue ivalue = IValue.from(array1)) {
                 Assert.assertTrue(ivalue.isTensor());
@@ -69,6 +71,17 @@ public class IValueTest {
                 Assert.assertEquals(ivalue.toDouble(), 1d);
             }
 
+            try (IValue ivalue = IValue.from(array3)) {
+                Assert.assertTrue(ivalue.isString());
+                Assert.assertEquals(ivalue.getType(), "str");
+                Assert.assertEquals(ivalue.toStringValue(), "test");
+            }
+
+            try (IValue ivalue = IValue.from(array4)) {
+                Assert.assertTrue(ivalue.isList());
+                Assert.assertEquals(ivalue.getType(), "str[]");
+            }
+
             try (IValue ivalue = IValue.from("test")) {
                 Assert.assertTrue(ivalue.isString());
                 Assert.assertEquals(ivalue.getType(), "str");