Added support for int64 -> string CategoryMapper updated (#2181)

* add int64->string support
Signed-off-by: pagran <[email protected]>

---------

Signed-off-by: pagran <[email protected]>
Co-authored-by: Jay Zhang <[email protected]>

Author: pagran

tests/test_backend.py

@@ -5440,6 +5440,24 @@ def func(query_holder):
         self._run_test_case(func, [_OUTPUT], {_INPUT: query}, as_session=True)
         os.remove(filnm)
 
+    @check_opset_min_version(8, "CategoryMapper")
+    @skip_tfjs("TFJS does not initialize table")
+    @skip_onnx_checker("ONNX can't do type inference on CategoryMapper")
+    def test_hashtable_lookup_invert(self):
+        filnm = "vocab.tmp"
+        words = ["apple", "pear", "banana", "cherry", "grape"]
+        query = np.array([3], dtype=np.int64)
+        with open(filnm, "w") as f:
+            for word in words:
+                f.write(word + "\n")
+        def func(query_holder):
+            hash_table = lookup_ops.index_to_string_table_from_file(filnm)
+            lookup_results = hash_table.lookup(query_holder)
+            ret = tf.identity(lookup_results, name=_TFOUTPUT)
+            return ret
+        self._run_test_case(func, [_OUTPUT], {_INPUT: query}, as_session=True)
+        os.remove(filnm)
+
     @check_opset_min_version(8, "CategoryMapper")
     @skip_tfjs("TFJS does not initialize table")
     def test_hashtable_lookup_const(self):
@@ -5458,6 +5476,24 @@ def func():
         self._run_test_case(func, [_OUTPUT], {}, as_session=True)
         os.remove(filnm)
 
+    @check_opset_min_version(8, "CategoryMapper")
+    @skip_tfjs("TFJS does not initialize table")
+    def test_hashtable_lookup_invert_const(self):
+        filnm = "vocab.tmp"
+        words = ["apple", "pear", "banana", "cherry", "grape"]
+        query_val = np.array([3, 2], dtype=np.int64).reshape((1, 2, 1))
+        with open(filnm, "w", encoding='UTF-8') as f:
+            for word in words:
+                f.write(word + "\n")
+        def func():
+            hash_table = lookup_ops.index_to_string_table_from_file(filnm)
+            query = tf.constant(query_val)
+            lookup_results = hash_table.lookup(query)
+            ret = tf.identity(lookup_results, name=_TFOUTPUT)
+            return ret
+        self._run_test_case(func, [_OUTPUT], {}, as_session=True)
+        os.remove(filnm)
+
     @skip_tfjs("TFJS does not initialize table")
     def test_hashtable_size(self):
         filnm = "vocab.tmp"

tf2onnx/custom_opsets/onnx_ml.py

@@ -42,10 +42,18 @@ def version_8(cls, ctx, node, initialized_tables, **kwargs):
 
         dtype = ctx.get_dtype(node.output[0])
         in_dtype = ctx.get_dtype(node.input[1])
-        utils.make_sure(dtype == TensorProto.INT64 and in_dtype == TensorProto.STRING,
-                        "Only lookup tables of type string->int64 are currently supported.")
+        utils.make_sure((dtype == TensorProto.INT64 and in_dtype == TensorProto.STRING) or
+                        (dtype == TensorProto.STRING and in_dtype == TensorProto.INT64),
+                        f"Only lookup tables of type string<->int64 are currently supported.")
+
+        if in_dtype == TensorProto.STRING:
+            cats_strings, cats_int64s = initialized_tables[shared_name]
+            default_key = 'default_int64'
+        else:
+            cats_int64s, cats_strings = initialized_tables[shared_name]
+            default_key = 'default_string'
+        attr = {'cats_int64s': cats_int64s, 'cats_strings': cats_strings, default_key: default_val}
 
-        cats_strings, cats_int64s = initialized_tables[shared_name]
         shape = ctx.get_shape(node.input[1])
 
         node_name = node.name
@@ -56,18 +64,19 @@ def version_8(cls, ctx, node, initialized_tables, **kwargs):
             # Handle explicitly since const folding doesn't work for tables
             key_np = node.inputs[1].get_tensor_value(as_list=False)
             ctx.remove_node(node.name)
-            key_to_val = dict(zip(cats_strings, cats_int64s))
-            def lookup_value(key):
-                return key_to_val.get(key.encode("UTF-8"), default_val_np)
-            lookup_result = np.vectorize(lookup_value)(key_np)
+            if in_dtype == TensorProto.STRING:
+                key_to_val = dict(zip(cats_strings, cats_int64s))
+                lookup_result = np.vectorize(lambda key: key_to_val.get(key.encode("UTF-8"), default_val_np))(key_np)
+            else:
+                key_to_val = dict(zip(cats_int64s, cats_strings))
+                lookup_result = np.vectorize(lambda key: key_to_val.get(key, default_val_np))(key_np).astype(object)
             onnx_tensor = numpy_helper.from_array(lookup_result, node_name)
             ctx.make_node("Const", name=node_name, inputs=[], outputs=node_outputs,
                           attr={"value": onnx_tensor}, shapes=[lookup_result.shape], dtypes=[dtype])
         else:
             ctx.remove_node(node.name)
             ctx.make_node("CategoryMapper", domain=constants.AI_ONNX_ML_DOMAIN,
-                          name=node_name, inputs=[node_inputs[1]], outputs=node_outputs,
-                          attr={'cats_int64s': cats_int64s, 'cats_strings': cats_strings, 'default_int64': default_val},
+                          name=node_name, inputs=[node_inputs[1]], outputs=node_outputs, attr=attr,
                           shapes=[shape], dtypes=[dtype])
 
         customer_nodes = ctx.find_output_consumers(table_node.output[0])