Check the associated files in the process units.

PiperOrigin-RevId: 318554536

Author: lu-wang-g

tensorflow_lite_support/metadata/BUILD

@@ -79,36 +79,33 @@ py_library(
     ],
 )
 
-# TODO(b/160127851): Temporarily disable the test because it doesn't build with TF internal libs.
-# py_test(
-#     name = "metadata_test",
-#     srcs = ["metadata_test.py"],
-#     data = ["testdata/golden_json.json"],
-#     python_version = "PY3",
-#     srcs_version = "PY2AND3",
-#     tags = [
-#         "no_mac",  # TODO(b/148247402): flatbuffers import broken on Mac OS.
-#     ],
-#     deps = [
-#         ":metadata",
-#         ":metadata_schema_py",
-#         ":schema_py",
-#         "@flatbuffers//:runtime_py",
-#         "@six_archive//:six",
-#         "@org_tensorflow//tensorflow/python:client_testlib",
-#         "@org_tensorflow//tensorflow/python:platform",
-#         "@org_tensorflow//tensorflow/python:platform_test",
-#     ],
-# )
+py_test(
+    name = "metadata_test",
+    srcs = ["metadata_test.py"],
+    data = ["testdata/golden_json.json"],
+    python_version = "PY3",
+    srcs_version = "PY2AND3",
+    tags = [
+        "no_mac",  # TODO(b/148247402): flatbuffers import broken on Mac OS.
+    ],
+    deps = [
+        ":metadata",
+        ":metadata_schema_py",
+        ":schema_py",
+        "//third_party/py/tensorflow",
+        "@flatbuffers//:runtime_py",
+        "@org_tensorflow//tensorflow/python:platform",
+        "@six_archive//:six",
+    ],
+)
 
-# TODO(b/160127851): Temporarily disable the test because it doesn't build with TF internal libs.
-# py_test(
-#     name = "metadata_parser_test",
-#     srcs = ["metadata_parser_test.py"],
-#     python_version = "PY3",
-#     srcs_version = "PY2AND3",
-#     deps = [
-#         ":metadata",
-#         "@org_tensorflow//tensorflow/python:client_testlib",
-#     ],
-# )
+py_test(
+    name = "metadata_parser_test",
+    srcs = ["metadata_parser_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY2AND3",
+    deps = [
+        ":metadata",
+        "//third_party/py/tensorflow",
+    ],
+)

tensorflow_lite_support/metadata/metadata.py

@@ -179,23 +179,37 @@ def get_recorded_associated_file_list(self):
     metadata = _metadata_fb.ModelMetadata.GetRootAsModelMetadata(
         self._metadata_buf, 0)
 
-    # Add associated files attached to ModelMetadata
-    self._get_associated_files_from_metadata_struct(metadata, recorded_files)
+    # Add associated files attached to ModelMetadata.
+    recorded_files += self._get_associated_files_from_table(
+        metadata, "AssociatedFiles")
 
-    # Add associated files attached to each SubgraphMetadata
+    # Add associated files attached to each SubgraphMetadata.
     for j in range(metadata.SubgraphMetadataLength()):
       subgraph = metadata.SubgraphMetadata(j)
-      self._get_associated_files_from_metadata_struct(subgraph, recorded_files)
+      recorded_files += self._get_associated_files_from_table(
+          subgraph, "AssociatedFiles")
 
-      # Add associated files attached to each input tensor
+      # Add associated files attached to each input tensor.
       for k in range(subgraph.InputTensorMetadataLength()):
-        tensor = subgraph.InputTensorMetadata(k)
-        self._get_associated_files_from_metadata_struct(tensor, recorded_files)
+        recorded_files += self._get_associated_files_from_table(
+            subgraph.InputTensorMetadata(k), "AssociatedFiles")
+        recorded_files += self._get_associated_files_from_process_units(
+            subgraph.InputTensorMetadata(k), "ProcessUnits")
 
-      # Add associated files attached to each output tensor
+      # Add associated files attached to each output tensor.
       for k in range(subgraph.OutputTensorMetadataLength()):
-        tensor = subgraph.OutputTensorMetadata(k)
-        self._get_associated_files_from_metadata_struct(tensor, recorded_files)
+        recorded_files += self._get_associated_files_from_table(
+            subgraph.OutputTensorMetadata(k), "AssociatedFiles")
+        recorded_files += self._get_associated_files_from_process_units(
+            subgraph.OutputTensorMetadata(k), "ProcessUnits")
+
+      # Add associated files attached to the input_process_units.
+      recorded_files += self._get_associated_files_from_process_units(
+          subgraph, "InputProcessUnits")
+
+      # Add associated files attached to the output_process_units.
+      recorded_files += self._get_associated_files_from_process_units(
+          subgraph, "OutputProcessUnits")
 
     return recorded_files
 
@@ -317,9 +331,69 @@ def _copy_archived_files(self, src_zip, dst_zip, file_list):
         file_buffer = src_zf.read(f)
         dst_zf.writestr(f, file_buffer)
 
-  def _get_associated_files_from_metadata_struct(self, file_holder, file_list):
-    for j in range(file_holder.AssociatedFilesLength()):
-      file_list.append(file_holder.AssociatedFiles(j).Name().decode("utf-8"))
+  def _get_associated_files_from_process_units(self, table, field_name):
+    """Gets the files that are attached the process units field of a table.
+
+    Args:
+      table: a Flatbuffers table object that contains fields of an array of
+        ProcessUnit, such as TensorMetadata and SubGraphMetadata.
+      field_name: the name of the field in the table that represents an array of
+        ProcessUnit. If the table is TensorMetadata, field_name can be
+        "ProcessUnits". If the table is SubGraphMetadata, field_name can be
+        either "InputProcessUnits" or "OutputProcessUnits".
+
+    Returns:
+      the associated files list.
+    """
+
+    if table is None:
+      return
+
+    file_list = []
+    length_method = getattr(table, field_name + "Length")
+    member_method = getattr(table, field_name)
+    for k in range(length_method()):
+      process_unit = member_method(k)
+      tokenizer = process_unit.Options()
+      if (process_unit.OptionsType() is
+          _metadata_fb.ProcessUnitOptions.BertTokenizerOptions):
+        bert_tokenizer = _metadata_fb.BertTokenizerOptions()
+        bert_tokenizer.Init(tokenizer.Bytes, tokenizer.Pos)
+        file_list += self._get_associated_files_from_table(
+            bert_tokenizer, "VocabFile")
+      elif (process_unit.OptionsType() is
+            _metadata_fb.ProcessUnitOptions.SentencePieceTokenizerOptions):
+        sentence_piece = _metadata_fb.SentencePieceTokenizerOptions()
+        sentence_piece.Init(tokenizer.Bytes, tokenizer.Pos)
+        file_list += self._get_associated_files_from_table(
+            sentence_piece, "SentencePieceModel")
+        file_list += self._get_associated_files_from_table(
+            sentence_piece, "VocabFile")
+    return file_list
+
+  def _get_associated_files_from_table(self, table, field_name):
+    """Gets the associated files that are attached a table directly.
+
+    Args:
+      table: a Flatbuffers table object that contains fields of an array of
+        AssociatedFile, such as TensorMetadata and BertTokenizerOptions.
+      field_name: the name of the field in the table that represents an array of
+        ProcessUnit. If the table is TensorMetadata, field_name can be
+        "AssociatedFiles". If the table is BertTokenizerOptions, field_name can
+        be "VocabFile".
+
+    Returns:
+      the associated files list.
+    """
+
+    if table is None:
+      return
+    file_list = []
+    length_method = getattr(table, field_name + "Length")
+    member_method = getattr(table, field_name)
+    for j in range(length_method()):
+      file_list.append(member_method(j).Name().decode("utf-8"))
+    return file_list
 
   def _populate_associated_files(self):
     """Concatenates associated files after TensorFlow Lite model file.

tensorflow_lite_support/metadata/metadata_parser_test.py

@@ -19,20 +19,19 @@
 from __future__ import print_function
 
 import re
+import tensorflow as tf
 
-from third_party.tensorflow.python.framework import test_util
-from third_party.tensorflow.python.platform import test
 from tensorflow_lite_support.metadata import metadata_parser
 
 
-class MetadataParserTest(test_util.TensorFlowTestCase):
+class MetadataParserTest(tf.test.TestCase):
 
-  def test_version_wellFormedSemanticVersion(self):
+  def testVersionWellFormedSemanticVersion(self):
     # Validates that the version is well-formed (x.y.z).
     self.assertTrue(
         re.match('[0-9]+\\.[0-9]+\\.[0-9]+',
                  metadata_parser.MetadataParser.VERSION))
 
 
 if __name__ == '__main__':
-  test.main()
+  tf.test.main()