Repository: deepjavalibrary/djl (java)

Issue:
Facing issues when Tensorflow model involves a String as input/output to serving function
Hi,
I am testing out a tensorflow model which takes a string as input and returns the spell corrected version of this string and it's confidence score. The signature_def retuned for the model using `saved_model_cli` is as following

```
signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['query'] tensor_info:
        dtype: DT_STRING
        shape: (-1)
        name: serving_default_query:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['corrected_query'] tensor_info:
        dtype: DT_STRING
        shape: ()
        name: StatefulPartitionedCall_5:0
    outputs['score'] tensor_info:
        dtype: DT_DOUBLE
        shape: unknown_rank
        name: StatefulPartitionedCall_5:1
  Method name is: tensorflow/serving/predict
```

But when loading the model on DJL, the loaded model is only identifying `score` as model output. `corrected_query` is not getting returned when checking with `model.describeOutput().keys()` and only `score` is getting returned upon model prediction call.

Any support or document around solving this issue would be really helpful.

Thanks

Pull request description:
[tensorflow] Remove String tensor limitation for model output
Fixes #2053

## 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/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java
index 331e1ff52..3eec14fbc 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java
@@ -14,17 +14,31 @@
 package ai.djl.tensorflow.engine;
 
 import org.tensorflow.internal.c_api.TF_Graph;
+import org.tensorflow.internal.c_api.TF_Operation;
 import org.tensorflow.internal.c_api.TF_Session;
+import org.tensorflow.internal.c_api.global.tensorflow;
+import org.tensorflow.proto.framework.CollectionDef;
 import org.tensorflow.proto.framework.MetaGraphDef;
+import org.tensorflow.proto.framework.SignatureDef;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /** The wrapper class for native resources required for SavedModelBundle. */
 public class SavedModelBundle implements AutoCloseable {
 
+    private static final String INIT_OP_SIGNATURE_KEY = "__saved_model_init_op";
+    private static final String MAIN_OP_COLLECTION_KEY = "saved_model_main_op";
+    private static final String LEGACY_INIT_OP_COLLECTION_KEY = "legacy_init_op";
+    private static final String TABLE_INITIALIZERS_COLLECTION_KEY = "table_initializer";
+
     private TF_Graph graphHandle;
     private TF_Session sessionHandle;
     private MetaGraphDef metaGraphDef;
+    private TF_Operation[] targetOpHandles;
     private AtomicBoolean closed;
 
     public SavedModelBundle(
@@ -33,6 +47,61 @@ public class SavedModelBundle implements AutoCloseable {
         this.sessionHandle = sessionHandle;
         this.metaGraphDef = metaGraphDef;
         closed = new AtomicBoolean(false);
+
+        Map<String, SignatureDef> functions = new ConcurrentHashMap<>();
+        metaGraphDef
+                .getSignatureDefMap()
+                .forEach(
+                        (signatureName, signatureDef) -> {
+                            if (!functions.containsKey(signatureName)) {
+                                functions.put(signatureName, signatureDef);
+                            }
+                        });
+
+        List<TF_Operation> initOps = new ArrayList<>();
+        TF_Operation initOp = findInitOp(functions, metaGraphDef.getCollectionDefMap());
+        if (initOp != null) {
+            initOps.add(initOp);
+        }
+
+        if (metaGraphDef.containsCollectionDef(TABLE_INITIALIZERS_COLLECTION_KEY)) {
+            metaGraphDef
+                    .getCollectionDefMap()
+                    .get(TABLE_INITIALIZERS_COLLECTION_KEY)
+                    .getNodeList()
+                    .getValueList()
+                    .forEach(
+                            node -> {
+                                initOps.add(tensorflow.TF_GraphOperationByName(graphHandle, node));
+                            });
+        }
+        targetOpHandles = initOps.toArray(new TF_Operation[0]);
+    }
+
+    private TF_Operation findInitOp(
+            Map<String, SignatureDef> signatures, Map<String, CollectionDef> collections) {
+        SignatureDef initSig = signatures.get(INIT_OP_SIGNATURE_KEY);
+        if (initSig != null) {
+            String opName = initSig.getOutputsMap().get(INIT_OP_SIGNATURE_KEY).getName();
+            return tensorflow.TF_GraphOperationByName(graphHandle, opName);
+        }
+
+        CollectionDef initCollection;
+        if (collections.containsKey(MAIN_OP_COLLECTION_KEY)) {
+            initCollection = collections.get(MAIN_OP_COLLECTION_KEY);
+        } else {
+            initCollection = collections.get(LEGACY_INIT_OP_COLLECTION_KEY);
+        }
+
+        if (initCollection != null) {
+            CollectionDef.NodeList nodes = initCollection.getNodeList();
+            if (nodes.getValueCount() != 1) {
+                throw new IllegalArgumentException("Expected exactly one main op in saved model.");
+            }
+            String opName = nodes.getValue(0);
+            return tensorflow.TF_GraphOperationByName(graphHandle, opName);
+        }
+        return null;
     }
 
     /**
@@ -62,6 +131,10 @@ public class SavedModelBundle implements AutoCloseable {
         return metaGraphDef;
     }
 
+    TF_Operation[] getTargetOpHandles() {
+        return targetOpHandles;
+    }
+
     /** {@inheritDoc} */
     @Override
     public void close() {
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java
index b6a853bdc..2453e58c2 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java
@@ -67,6 +67,7 @@ public class TfSymbolBlock extends AbstractSymbolBlock implements AutoCloseable
         this.bundle = bundle;
         graphHandle = bundle.getGraph();
         sessionHandle = bundle.getSession();
+        targetOpHandles = bundle.getTargetOpHandles();
         MetaGraphDef metaGraphDef = bundle.getMetaGraphDef();
         Map<String, SignatureDef> signatureDefMap = metaGraphDef.getSignatureDefMap();
         if (signatureDefMap.containsKey(signatureDefKey)) {
@@ -89,8 +90,6 @@ public class TfSymbolBlock extends AbstractSymbolBlock implements AutoCloseable
         }
         describeInput();
         describeOutput();
-        // we don't use target for now
-        targetOpHandles = new TF_Operation[0];
     }
 
     /** {@inheritDoc} */
@@ -215,10 +214,6 @@ public class TfSymbolBlock extends AbstractSymbolBlock implements AutoCloseable
             for (String key : keys) {
                 TensorInfo tensorInfo = outputsMap.get(key);
                 TensorShapeProto shapeProto = tensorInfo.getTensorShape();
-                // does not support string tensors
-                if (tensorInfo.getDtype() == org.tensorflow.proto.framework.DataType.DT_STRING) {
-                    continue;
-                }
                 outputDescriptions.add(
                         key,
                         new Shape(


Test patch:
diff --git a/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java b/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java
index 649ca193b..3e8937de7 100644
--- a/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java
+++ b/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java
@@ -19,11 +19,16 @@ import ai.djl.modality.Classifications;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
 import ai.djl.modality.cv.output.DetectedObjects;
+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.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.testing.TestRequirements;
 import ai.djl.training.util.ProgressBar;
+import ai.djl.translate.NoopTranslator;
 import ai.djl.translate.TranslateException;
 import ai.djl.util.Pair;
 
@@ -59,14 +64,21 @@ public class TfSsdTest {
                 Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
             Assert.assertEquals(model.describeInput().get(0).getValue(), new Shape(-1, -1, -1, 3));
             for (Pair<String, Shape> pair : model.describeOutput()) {
-                if (pair.getKey().contains("label")) {
-                    Assert.assertEquals(pair.getValue(), new Shape(-1, 1));
-                } else if (pair.getKey().contains("box")) {
-                    Assert.assertEquals(pair.getValue(), new Shape(-1, 4));
-                } else if (pair.getKey().contains("score")) {
-                    Assert.assertEquals(pair.getValue(), new Shape(-1, 1));
-                } else {
-                    throw new IllegalStateException("Unexpected output name:" + pair.getKey());
+                switch (pair.getKey()) {
+                    case "box":
+                    case "detection_boxes":
+                        Assert.assertEquals(pair.getValue(), new Shape(-1, 4));
+                        break;
+                    case "label":
+                    case "score":
+                    case "detection_class_entities":
+                    case "detection_class_labels":
+                    case "detection_class_names":
+                    case "detection_scores":
+                        Assert.assertEquals(pair.getValue(), new Shape(-1, 1));
+                        break;
+                    default:
+                        throw new IllegalStateException("Unexpected output name: " + pair.getKey());
                 }
             }
 
@@ -82,6 +94,32 @@ public class TfSsdTest {
         }
     }
 
+    @Test
+    public void testStringInputOutput() throws IOException, ModelException, TranslateException {
+        TestRequirements.notArm();
+
+        Criteria<NDList, NDList> criteria =
+                Criteria.builder()
+                        .optApplication(Application.CV.OBJECT_DETECTION)
+                        .setTypes(NDList.class, NDList.class)
+                        .optArtifactId("ssd")
+                        .optFilter("backbone", "mobilenet_v2")
+                        .optEngine("TensorFlow")
+                        .optProgress(new ProgressBar())
+                        .optTranslator(new NoopTranslator())
+                        .build();
+
+        try (ZooModel<NDList, NDList> model = criteria.loadModel();
+                Predictor<NDList, NDList> predictor = model.newPredictor();
+                NDManager manager = NDManager.newBaseManager()) {
+            NDArray array = manager.zeros(new Shape(1, 224, 224, 3));
+            NDList output = predictor.predict(new NDList(array));
+            Assert.assertEquals(output.size(), 5);
+            NDArray entities = output.get("detection_class_entities");
+            Assert.assertEquals(entities.getDataType(), DataType.STRING);
+        }
+    }
+
     private static void saveBoundingBoxImage(Image img, DetectedObjects detection)
             throws IOException {
         Path outputDir = Paths.get("build/output");