chore: import upstream snapshot with attribution

tensorflow/lite/tools/signature/BUILD

@@ -0,0 +1,118 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "pybind_extension", "replace_with_portable_tf_lib_when_required")
+
+# Utilities for signature_defs in TFLite
+load("//tensorflow/lite:build_def.bzl", "tflite_copts")
+load("//tensorflow/lite:special_rules.bzl", "tflite_portable_test_suite")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "signature_def_util",
+    srcs = ["signature_def_util.cc"],
+    hdrs = ["signature_def_util.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    features = select({
+        "//tensorflow:android": ["-layering_check"],
+        "//conditions:default": [],
+    }),
+    deps = replace_with_portable_tf_lib_when_required([
+        "//tensorflow/core:lib_proto_parsing",
+        "//tensorflow/core:protos_all_cc",
+    ]) + [
+        "//tensorflow/core:protos_all_cc_impl",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_protobuf//:protobuf",
+        "@flatbuffers",
+    ],
+)
+
+cc_test(
+    name = "signature_def_util_test",
+    size = "small",
+    srcs = ["signature_def_util_test.cc"],
+    data = [
+        "//tensorflow/lite:testdata/add.bin",
+    ],
+    tags = [
+        "no_oss",
+        "tflite_not_portable",
+    ],
+    deps = [
+        ":signature_def_util",
+        "//tensorflow/cc/saved_model:signature_constants",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers//:runtime_cc",
+    ],
+)
+
+pybind_extension(
+    name = "_pywrap_signature_def_util_wrapper",
+    srcs = [
+        "signature_def_util_wrapper_pybind11.cc",
+    ],
+    data = [
+        "_pywrap_signature_def_util_wrapper.pyi",
+    ],
+    enable_stub_generation = True,
+    deps = [
+        ":signature_def_util",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "@com_google_absl//absl/status",
+        "@pybind11",
+    ],
+)
+
+py_library(
+    name = "signature_def_utils",
+    srcs = ["signature_def_utils.py"],
+    strict_deps = True,
+    deps = [
+        ":_pywrap_signature_def_util_wrapper",
+        "//tensorflow/core:protos_all_py",
+    ],
+)
+
+py_test(
+    name = "signature_def_utils_test",
+    srcs = ["signature_def_utils_test.py"],
+    data = ["//tensorflow/lite:testdata/add.bin"],
+    strict_deps = True,
+    tags = [
+        "no_mac",
+        "no_oss",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":signature_def_utils",
+        #internal proto upb dep
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/core:protos_all_py",
+    ],
+)
+
+tflite_portable_test_suite()

tensorflow/lite/tools/signature/_pywrap_signature_def_util_wrapper.pyi

@@ -0,0 +1,18 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+def ClearSignatureDefs(arg0: list[int]) -> bytes: ...
+def GetSignatureDefMap(arg0: list[int]) -> dict[str, bytes]: ...
+def SetSignatureDefMap(arg0: list[int], arg1: dict[str, str]) -> bytes: ...

tensorflow/lite/tools/signature/signature_def_util.cc

@@ -0,0 +1,214 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/tools/signature/signature_def_util.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "flatbuffers/vector.h"  // from @flatbuffers
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using SerializedSignatureDefMap = absl::flat_hash_map<std::string, std::string>;
+using SignatureDefMap = std::map<std::string, tensorflow::SignatureDef>;
+
+const Metadata* GetSignatureDefMetadata(const Model* model) {
+  if (!model || !model->metadata()) {
+    return nullptr;
+  }
+  for (const Metadata* metadata : *model->metadata()) {
+    if (metadata && metadata->name() != nullptr &&
+        absl::string_view(metadata->name()->c_str(),
+                          metadata->name()->size()) ==
+            kSignatureDefsMetadataName) {
+      return metadata;
+    }
+  }
+  return nullptr;
+}
+
+absl::Status ReadSignatureDefMap(const Model* model, const Metadata* metadata,
+                                 SerializedSignatureDefMap* map) {
+  if (!model || !metadata || !map) {
+    return absl::InvalidArgumentError("Arguments must not be nullptr");
+  }
+  if (!model->buffers()) {
+    return absl::InvalidArgumentError("Missing buffers vector in model");
+  }
+  if (metadata->buffer() >= model->buffers()->size()) {
+    return absl::InternalError("Invalid buffer index in metadata");
+  }
+  const flatbuffers::Vector<uint8_t>* flatbuffer_data =
+      model->buffers()->Get(metadata->buffer())->data();
+  if (!flatbuffer_data || flatbuffer_data->size() < 3 ||
+      !flexbuffers::VerifyBuffer(flatbuffer_data->data(),
+                                 flatbuffer_data->size())) {
+    return absl::InvalidArgumentError("Invalid flexbuffer data");
+  }
+  const auto signature_defs =
+      flexbuffers::GetRoot(flatbuffer_data->data(), flatbuffer_data->size())
+          .AsMap();
+  for (size_t i = 0; i < signature_defs.size(); ++i) {
+    const std::string key = signature_defs.Keys()[i].AsString().str();
+    (*map)[key] = signature_defs.Values()[i].AsString().str();
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status SetSignatureDefMap(const Model* model,
+                                const SignatureDefMap& signature_def_map,
+                                std::string* model_data_with_signature_defs) {
+  if (!model || !model_data_with_signature_defs) {
+    return absl::InvalidArgumentError("Arguments must not be nullptr");
+  }
+  if (signature_def_map.empty()) {
+    return absl::InvalidArgumentError("signature_def_map should not be empty");
+  }
+  if (!model->buffers()) {
+    return absl::InvalidArgumentError("Missing buffers vector in model");
+  }
+  flexbuffers::Builder fbb;
+  const size_t start_map = fbb.StartMap();
+  auto mutable_model = std::make_unique<ModelT>();
+  model->UnPackTo(mutable_model.get(), nullptr);
+  uint32_t buffer_id = mutable_model->buffers.size();
+  const Metadata* metadata = GetSignatureDefMetadata(model);
+  if (metadata) {
+    buffer_id = metadata->buffer();
+    if (buffer_id >= mutable_model->buffers.size()) {
+      return absl::InternalError("Invalid buffer index in metadata");
+    }
+  } else {
+    auto buffer = std::make_unique<BufferT>();
+    mutable_model->buffers.emplace_back(std::move(buffer));
+    auto sigdef_metadata = std::make_unique<MetadataT>();
+    sigdef_metadata->buffer = buffer_id;
+    sigdef_metadata->name = kSignatureDefsMetadataName;
+    mutable_model->metadata.emplace_back(std::move(sigdef_metadata));
+  }
+  for (const auto& [key, signature_def] : signature_def_map) {
+    fbb.String(key.c_str(), signature_def.SerializeAsString());
+  }
+  fbb.EndMap(start_map);
+  fbb.Finish();
+  mutable_model->buffers[buffer_id]->data = fbb.GetBuffer();
+  flatbuffers::FlatBufferBuilder builder;
+  flatbuffers::Offset<Model> packed_model =
+      Model::Pack(builder, mutable_model.get());
+  FinishModelBuffer(builder, packed_model);
+  model_data_with_signature_defs->assign(
+      reinterpret_cast<const char*>(builder.GetBufferPointer()),
+      builder.GetSize());
+  return absl::OkStatus();
+}
+
+bool HasSignatureDef(const Model* model, absl::string_view signature_key) {
+  if (!model) {
+    return false;
+  }
+  const Metadata* metadata = GetSignatureDefMetadata(model);
+  if (!metadata) {
+    return false;
+  }
+  if (!model->buffers()) {
+    return false;
+  }
+  if (metadata->buffer() >= model->buffers()->size()) {
+    return false;
+  }
+  const flatbuffers::Vector<uint8_t>* flatbuffer_data =
+      model->buffers()->Get(metadata->buffer())->data();
+  if (!flatbuffer_data || flatbuffer_data->size() < 3 ||
+      !flexbuffers::VerifyBuffer(flatbuffer_data->data(),
+                                 flatbuffer_data->size())) {
+    return false;
+  }
+  const auto signature_defs =
+      flexbuffers::GetRoot(flatbuffer_data->data(), flatbuffer_data->size())
+          .AsMap();
+  return !signature_defs[std::string(signature_key)].IsNull();
+}
+
+absl::Status GetSignatureDefMap(const Model* model,
+                                SignatureDefMap* signature_def_map) {
+  if (!model || !signature_def_map) {
+    return absl::InvalidArgumentError("Arguments must not be nullptr");
+  }
+  SignatureDefMap retrieved_signature_def_map;
+  const Metadata* metadata = GetSignatureDefMetadata(model);
+  if (metadata) {
+    SerializedSignatureDefMap signature_defs;
+    absl::Status status = ReadSignatureDefMap(model, metadata, &signature_defs);
+    if (!status.ok()) {
+      return absl::Status(
+          status.code(),
+          absl::StrCat("Error reading signature def map: ", status.message()));
+    }
+    tensorflow::SignatureDef signature_def;
+    for (const auto& [key, serialized_def] : signature_defs) {
+      if (!signature_def.ParseFromString(serialized_def)) {
+        return absl::InternalError(
+            "Cannot parse signature def found in flatbuffer.");
+      }
+      retrieved_signature_def_map[key] = std::move(signature_def);
+    }
+    *signature_def_map = std::move(retrieved_signature_def_map);
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ClearSignatureDefMap(const Model* model, std::string* model_data) {
+  if (!model || !model_data) {
+    return absl::InvalidArgumentError("Arguments must not be nullptr");
+  }
+  auto mutable_model = std::make_unique<ModelT>();
+  model->UnPackTo(mutable_model.get(), nullptr);
+  auto it = absl::c_find_if(mutable_model->metadata, [](const auto& m) {
+    return m->name == kSignatureDefsMetadataName;
+  });
+  if (it != mutable_model->metadata.end()) {
+    if ((*it)->buffer >= mutable_model->buffers.size()) {
+      return absl::InternalError("Invalid buffer index in metadata");
+    }
+    mutable_model->buffers[(*it)->buffer]->data.clear();
+    mutable_model->metadata.erase(it);
+  }
+  flatbuffers::FlatBufferBuilder builder;
+  flatbuffers::Offset<Model> packed_model =
+      Model::Pack(builder, mutable_model.get());
+  FinishModelBuffer(builder, packed_model);
+  model_data->assign(reinterpret_cast<const char*>(builder.GetBufferPointer()),
+                     builder.GetSize());
+  return absl::OkStatus();
+}
+
+}  // namespace tflite

tensorflow/lite/tools/signature/signature_def_util.h

@@ -0,0 +1,78 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_TOOLS_SIGNATURE_SIGNATURE_DEF_UTIL_H_
+#define TENSORFLOW_LITE_TOOLS_SIGNATURE_SIGNATURE_DEF_UTIL_H_
+
+#include <map>
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+
+// Constant for name of the Metadata entry associated with SignatureDefs.
+inline constexpr absl::string_view kSignatureDefsMetadataName =
+    "signature_defs_metadata";
+
+// The function `SetSignatureDefMap()` results in
+// `model_data_with_signature_defs` containing a serialized TFLite model
+// identical to `model` with a metadata and associated buffer containing
+// a FlexBuffer::Map with `signature_def_map` keys and values serialized to
+// String.
+//
+// If a Metadata entry containing a SignatureDef map exists, it will be
+//   overwritten.
+//
+// Returns error if `model_data_with_signature_defs` is null or
+//   `signature_def_map` is empty.
+//
+// On success, returns absl::OkStatus() or error otherwise.
+// On error, `model_data_with_signature_defs` is unchanged.
+absl::Status SetSignatureDefMap(
+    const Model* model,
+    const std::map<std::string, tensorflow::SignatureDef>& signature_def_map,
+    std::string* model_data_with_signature_defs);
+
+// The function `HasSignatureDef()` returns true if `model` contains a Metadata
+// table pointing to a buffer containing a FlexBuffer::Map and the map has
+// `signature_key` as a key, or false otherwise.
+bool HasSignatureDef(const Model* model, absl::string_view signature_key);
+
+// The function `GetSignatureDefMap()` results in `signature_def_map`
+// pointing to a map<std::string, tensorflow::SignatureDef>
+// parsed from `model`'s metadata buffer.
+//
+// If the Metadata entry does not exist, `signature_def_map` is unchanged.
+//
+// Returns error if `model` or `signature_def_map` is null.
+// If the Metadata entry exists but cannot be parsed, returns an error.
+absl::Status GetSignatureDefMap(
+    const Model* model,
+    std::map<std::string, tensorflow::SignatureDef>* signature_def_map);
+
+// The function `ClearSignatureDefMap` results in `model_data`
+// containing a serialized Model identical to `model` omitting any
+// SignatureDef-related metadata or buffers.
+//
+// Returns error if `model` or `model_data` is null.
+// On success, returns absl::OkStatus() or error otherwise.
+absl::Status ClearSignatureDefMap(const Model* model, std::string* model_data);
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_SIGNATURE_SIGNATURE_DEF_UTIL_H_

tensorflow/lite/tools/signature/signature_def_util_test.cc

@@ -0,0 +1,195 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/tools/signature/signature_def_util.h"
+
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/cc/saved_model/signature_constants.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+#include "tensorflow/lite/core/model_builder.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using ::tensorflow::kClassifyMethodName;
+using ::tensorflow::kDefaultServingSignatureDefKey;
+using ::tensorflow::kPredictMethodName;
+using ::tensorflow::SignatureDef;
+using ::testing::ElementsAre;
+using ::testing::EqualsProto;
+using ::testing::Pair;
+using ::testing::status::StatusIs;
+
+constexpr absl::string_view kSignatureInput = "input";
+constexpr absl::string_view kSignatureOutput = "output";
+constexpr absl::string_view kTestFilePath =
+    "tensorflow/lite/testdata/add.bin";
+
+SignatureDef GetTestSignatureDef() {
+  SignatureDef signature_def;
+  tensorflow::TensorInfo input_tensor;
+  tensorflow::TensorInfo output_tensor;
+  input_tensor.set_name(kSignatureInput);
+  output_tensor.set_name(kSignatureOutput);
+  signature_def.set_method_name(kClassifyMethodName);
+  (*signature_def.mutable_inputs())[kSignatureInput] = std::move(input_tensor);
+  (*signature_def.mutable_outputs())[kSignatureOutput] =
+      std::move(output_tensor);
+  return signature_def;
+}
+
+class SimpleSignatureDefUtilTest : public testing::Test {
+ protected:
+  void SetUp() override {
+    flatbuffer_model_ =
+        FlatBufferModel::BuildFromFile(std::string(kTestFilePath).c_str());
+    if (!flatbuffer_model_) {
+      GTEST_SKIP() << "Failed to load model";
+    }
+    model_ = flatbuffer_model_->GetModel();
+    if (!model_) {
+      GTEST_SKIP() << "Failed to get model";
+    }
+  }
+
+  std::unique_ptr<FlatBufferModel> flatbuffer_model_;
+  const Model* model_;
+};
+
+TEST_F(SimpleSignatureDefUtilTest, SetSignatureDefTest) {
+  SignatureDef expected_signature_def = GetTestSignatureDef();
+  std::string model_output;
+  const std::map<std::string, SignatureDef> expected_signature_def_map = {
+      {std::string(kDefaultServingSignatureDefKey), expected_signature_def}};
+  ASSERT_OK(
+      SetSignatureDefMap(model_, expected_signature_def_map, &model_output));
+  const Model* add_model = flatbuffers::GetRoot<Model>(model_output.data());
+  EXPECT_TRUE(HasSignatureDef(add_model, kDefaultServingSignatureDefKey));
+  std::map<std::string, SignatureDef> test_signature_def_map;
+  ASSERT_OK(GetSignatureDefMap(add_model, &test_signature_def_map));
+  EXPECT_THAT(test_signature_def_map,
+              ElementsAre(Pair(std::string(kDefaultServingSignatureDefKey),
+                               EqualsProto(expected_signature_def))));
+}
+
+TEST_F(SimpleSignatureDefUtilTest, OverwriteSignatureDefTest) {
+  SignatureDef expected_signature_def = GetTestSignatureDef();
+  std::string model_output;
+  std::map<std::string, SignatureDef> expected_signature_def_map = {
+      {std::string(kDefaultServingSignatureDefKey), expected_signature_def}};
+  ASSERT_OK(
+      SetSignatureDefMap(model_, expected_signature_def_map, &model_output));
+  const Model* add_model = flatbuffers::GetRoot<Model>(model_output.data());
+  EXPECT_TRUE(HasSignatureDef(add_model, kDefaultServingSignatureDefKey));
+  std::map<std::string, SignatureDef> test_signature_def_map;
+  ASSERT_OK(GetSignatureDefMap(add_model, &test_signature_def_map));
+  EXPECT_THAT(test_signature_def_map,
+              ElementsAre(Pair(std::string(kDefaultServingSignatureDefKey),
+                               EqualsProto(expected_signature_def))));
+  expected_signature_def.set_method_name(std::string(kPredictMethodName));
+  expected_signature_def_map.erase(std::string(kDefaultServingSignatureDefKey));
+  static constexpr absl::string_view kTestSignatureDefKey = "ServingTest";
+  expected_signature_def_map[std::string(kTestSignatureDefKey)] =
+      expected_signature_def;
+  ASSERT_OK(
+      SetSignatureDefMap(add_model, expected_signature_def_map, &model_output));
+  const Model* final_model = flatbuffers::GetRoot<Model>(model_output.data());
+  EXPECT_FALSE(HasSignatureDef(final_model, kDefaultServingSignatureDefKey));
+  EXPECT_TRUE(HasSignatureDef(final_model, kTestSignatureDefKey));
+  ASSERT_OK(GetSignatureDefMap(final_model, &test_signature_def_map));
+  EXPECT_THAT(test_signature_def_map,
+              ElementsAre(Pair(std::string(kTestSignatureDefKey),
+                               EqualsProto(expected_signature_def))));
+}
+
+TEST_F(SimpleSignatureDefUtilTest, GetSignatureDefTest) {
+  std::map<std::string, SignatureDef> test_signature_def_map;
+  EXPECT_OK(GetSignatureDefMap(model_, &test_signature_def_map));
+  EXPECT_TRUE(test_signature_def_map.empty());
+  EXPECT_FALSE(HasSignatureDef(model_, kDefaultServingSignatureDefKey));
+}
+
+TEST_F(SimpleSignatureDefUtilTest, ClearSignatureDefTest) {
+  const uint32_t expected_num_buffers = model_->buffers()->size();
+  SignatureDef expected_signature_def = GetTestSignatureDef();
+  std::string model_output;
+  std::map<std::string, SignatureDef> expected_signature_def_map = {
+      {std::string(kDefaultServingSignatureDefKey), expected_signature_def}};
+  ASSERT_OK(
+      SetSignatureDefMap(model_, expected_signature_def_map, &model_output));
+  const Model* add_model = flatbuffers::GetRoot<Model>(model_output.data());
+  EXPECT_TRUE(HasSignatureDef(add_model, kDefaultServingSignatureDefKey));
+  std::map<std::string, SignatureDef> test_signature_def_map;
+  ASSERT_OK(GetSignatureDefMap(add_model, &test_signature_def_map));
+  SignatureDef test_signature_def =
+      test_signature_def_map[std::string(kDefaultServingSignatureDefKey)];
+  EXPECT_THAT(test_signature_def, EqualsProto(expected_signature_def));
+  ASSERT_OK(ClearSignatureDefMap(add_model, &model_output));
+  const Model* clear_model = flatbuffers::GetRoot<Model>(model_output.data());
+  EXPECT_FALSE(HasSignatureDef(clear_model, kDefaultServingSignatureDefKey));
+  EXPECT_EQ(expected_num_buffers + 1, clear_model->buffers()->size());
+}
+
+TEST_F(SimpleSignatureDefUtilTest, SetSignatureDefErrorsTest) {
+  std::map<std::string, SignatureDef> test_signature_def_map;
+  std::string model_output;
+  EXPECT_THAT(SetSignatureDefMap(model_, test_signature_def_map, &model_output),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+  SignatureDef test_signature_def;
+  test_signature_def_map[std::string(kDefaultServingSignatureDefKey)] =
+      test_signature_def;
+  EXPECT_THAT(SetSignatureDefMap(model_, test_signature_def_map, nullptr),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST_F(SimpleSignatureDefUtilTest, GetSignatureDefErrorsTest) {
+  auto mutable_model = std::make_unique<ModelT>();
+  model_->UnPackTo(mutable_model.get(), nullptr);
+  uint32_t buffer_id = mutable_model->buffers.size();
+  auto buffer = std::make_unique<BufferT>();
+  buffer->data = {0, 1};
+  mutable_model->buffers.emplace_back(std::move(buffer));
+  auto sigdef_metadata = std::make_unique<MetadataT>();
+  sigdef_metadata->buffer = buffer_id;
+  sigdef_metadata->name = kSignatureDefsMetadataName;
+  mutable_model->metadata.emplace_back(std::move(sigdef_metadata));
+
+  flatbuffers::FlatBufferBuilder builder;
+  flatbuffers::Offset<Model> packed_model =
+      Model::Pack(builder, mutable_model.get());
+  FinishModelBuffer(builder, packed_model);
+
+  const Model* invalid_model =
+      flatbuffers::GetRoot<Model>(builder.GetBufferPointer());
+
+  EXPECT_FALSE(HasSignatureDef(invalid_model, kDefaultServingSignatureDefKey));
+  std::map<std::string, SignatureDef> test_signature_def_map;
+  EXPECT_THAT(GetSignatureDefMap(invalid_model, &test_signature_def_map),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/tools/signature/signature_def_util_wrapper_pybind11.cc

@@ -0,0 +1,103 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstdint>
+#include <map>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+#include "tensorflow/lite/core/model_builder.h"
+#include "tensorflow/lite/tools/signature/signature_def_util.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+
+py::bytes WrappedSetSignatureDefMap(
+    const std::vector<uint8_t>& model_buffer,
+    const std::map<std::string, std::string>& serialized_signature_def_map) {
+  auto flatbuffer_model = tflite::FlatBufferModel::BuildFromBuffer(
+      reinterpret_cast<const char*>(model_buffer.data()), model_buffer.size());
+  auto* model = flatbuffer_model->GetModel();
+  if (!model) {
+    throw std::invalid_argument("Invalid model");
+  }
+  std::string data;
+  std::map<std::string, tensorflow::SignatureDef> signature_def_map;
+  for (const auto& entry : serialized_signature_def_map) {
+    tensorflow::SignatureDef signature_def;
+    if (!signature_def.ParseFromString(entry.second)) {
+      throw std::invalid_argument("Cannot parse signature def");
+    }
+    signature_def_map[entry.first] = signature_def;
+  }
+  auto status = tflite::SetSignatureDefMap(model, signature_def_map, &data);
+  if (status != absl::OkStatus()) {
+    throw std::invalid_argument(std::string(status.message()));
+  }
+  return py::bytes(data);
+}
+
+std::map<std::string, py::bytes> WrappedGetSignatureDefMap(
+    const std::vector<uint8_t>& model_buffer) {
+  auto flatbuffer_model = tflite::FlatBufferModel::BuildFromBuffer(
+      reinterpret_cast<const char*>(model_buffer.data()), model_buffer.size());
+  auto* model = flatbuffer_model->GetModel();
+  if (!model) {
+    throw std::invalid_argument("Invalid model");
+  }
+  std::string content;
+  std::map<std::string, tensorflow::SignatureDef> signature_def_map;
+  auto status = tflite::GetSignatureDefMap(model, &signature_def_map);
+  if (status != absl::OkStatus()) {
+    throw std::invalid_argument("Cannot parse signature def");
+  }
+  std::map<std::string, py::bytes> serialized_signature_def_map;
+  for (const auto& entry : signature_def_map) {
+    serialized_signature_def_map[entry.first] =
+        py::bytes(entry.second.SerializeAsString());
+  }
+  return serialized_signature_def_map;
+}
+
+py::bytes WrappedClearSignatureDefs(const std::vector<uint8_t>& model_buffer) {
+  auto flatbuffer_model = tflite::FlatBufferModel::BuildFromBuffer(
+      reinterpret_cast<const char*>(model_buffer.data()), model_buffer.size());
+  auto* model = flatbuffer_model->GetModel();
+  if (!model) {
+    throw std::invalid_argument("Invalid model");
+  }
+  std::string content;
+  auto status = tflite::ClearSignatureDefMap(model, &content);
+  if (status != absl::OkStatus()) {
+    throw std::invalid_argument("An unknown error occurred");
+  }
+  return py::bytes(content);
+}
+
+PYBIND11_MODULE(_pywrap_signature_def_util_wrapper, m) {
+  m.doc() = R"pbdoc(
+    _pywrap_signature_def_util_wrapper
+    -----
+  )pbdoc";
+
+  m.def("SetSignatureDefMap", &WrappedSetSignatureDefMap);
+
+  m.def("GetSignatureDefMap", &WrappedGetSignatureDefMap);
+
+  m.def("ClearSignatureDefs", &WrappedClearSignatureDefs);
+}

tensorflow/lite/tools/signature/signature_def_utils.py

@@ -0,0 +1,94 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utility functions related to SignatureDefs.
+
+Note: This utility is not compatiable with tensorflow.org/lite/guide/signatures.
+"""
+
+from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.lite.tools.signature import _pywrap_signature_def_util_wrapper as signature_def_util
+
+
+def set_signature_defs(tflite_model, signature_def_map):
+  """Sets SignatureDefs to the Metadata of a TfLite flatbuffer buffer.
+
+  Args:
+    tflite_model: Binary TFLite model (bytes or bytes-like object) to which to
+      add signature_def.
+    signature_def_map: dict containing SignatureDefs to store in metadata.
+  Returns:
+    buffer: A TFLite model binary identical to model buffer with
+      metadata field containing SignatureDef.
+
+  Raises:
+    ValueError:
+      tflite_model buffer does not contain a valid TFLite model.
+      signature_def_map is empty or does not contain a SignatureDef.
+  """
+  model = tflite_model
+  if not isinstance(tflite_model, bytearray):
+    model = bytearray(tflite_model)
+  serialized_signature_def_map = {
+      k: v.SerializeToString() for k, v in signature_def_map.items()}
+  model_buffer = signature_def_util.SetSignatureDefMap(
+      model, serialized_signature_def_map)
+  return model_buffer
+
+
+def get_signature_defs(tflite_model):
+  """Get SignatureDef dict from the Metadata of a TfLite flatbuffer buffer.
+
+  Args:
+    tflite_model: TFLite model buffer to get the signature_def.
+
+  Returns:
+    dict containing serving names to SignatureDefs if exists, otherwise, empty
+      dict.
+
+  Raises:
+    ValueError:
+      tflite_model buffer does not contain a valid TFLite model.
+    DecodeError:
+      SignatureDef cannot be parsed from TfLite SignatureDef metadata.
+  """
+  model = tflite_model
+  if not isinstance(tflite_model, bytearray):
+    model = bytearray(tflite_model)
+  serialized_signature_def_map = signature_def_util.GetSignatureDefMap(model)
+  def _deserialize(serialized):
+    signature_def = meta_graph_pb2.SignatureDef()
+    signature_def.ParseFromString(serialized)
+    return signature_def
+  return {k: _deserialize(v) for k, v in serialized_signature_def_map.items()}
+
+
+def clear_signature_defs(tflite_model):
+  """Clears SignatureDefs from the Metadata of a TfLite flatbuffer buffer.
+
+  Args:
+    tflite_model: TFLite model buffer to remove signature_defs.
+
+  Returns:
+    buffer: A TFLite model binary identical to model buffer with
+      no SignatureDef metadata.
+
+  Raises:
+    ValueError:
+      tflite_model buffer does not contain a valid TFLite model.
+  """
+  model = tflite_model
+  if not isinstance(tflite_model, bytearray):
+    model = bytearray(tflite_model)
+  return signature_def_util.ClearSignatureDefs(model)

tensorflow/lite/tools/signature/signature_def_utils_test.py

@@ -0,0 +1,71 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for signature_def_util.py.
+
+   - Tests adding a SignatureDef to TFLite metadata.
+"""
+
+import tensorflow as tf
+from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.lite.tools.signature import signature_def_utils
+
+
+class SignatureDefUtilsTest(tf.test.TestCase):
+
+  def testAddSignatureDefToFlatbufferMetadata(self):
+    """Test a SavedModel conversion has correct Metadata."""
+    filename = tf.compat.v1.resource_loader.get_path_to_datafile(
+        '../../testdata/add.bin')
+    if not tf.io.gfile.exists(filename):
+      raise IOError('File "{0}" does not exist in {1}.'.format(
+          filename,
+          tf.compat.v1.resource_loader.get_root_dir_with_all_resources()))
+
+    with tf.io.gfile.GFile(filename, 'rb') as fp:
+      tflite_model = bytearray(fp.read())
+
+    self.assertIsNotNone(tflite_model, 'TFLite model is none')
+    sig_input_tensor = meta_graph_pb2.TensorInfo(
+        dtype=tf.as_dtype(tf.float32).as_datatype_enum,
+        tensor_shape=tf.TensorShape([1, 8, 8, 3]).as_proto())
+    sig_input_tensor_signature = {'x': sig_input_tensor}
+    sig_output_tensor = meta_graph_pb2.TensorInfo(
+        dtype=tf.as_dtype(tf.float32).as_datatype_enum,
+        tensor_shape=tf.TensorShape([1, 8, 8, 3]).as_proto())
+    sig_output_tensor_signature = {'y': sig_output_tensor}
+    predict_signature_def = (
+        tf.compat.v1.saved_model.build_signature_def(
+            sig_input_tensor_signature, sig_output_tensor_signature,
+            tf.saved_model.PREDICT_METHOD_NAME))
+    serving_key = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
+    signature_def_map = {serving_key: predict_signature_def}
+    tflite_model = signature_def_utils.set_signature_defs(
+        tflite_model, signature_def_map)
+    saved_signature_def_map = signature_def_utils.get_signature_defs(
+        tflite_model)
+    signature_def = saved_signature_def_map.get(serving_key)
+    self.assertIsNotNone(signature_def, 'SignatureDef not found')
+    self.assertEqual(signature_def.SerializeToString(),
+                     predict_signature_def.SerializeToString())
+    remove_tflite_model = (
+        signature_def_utils.clear_signature_defs(tflite_model))
+    signature_def_map = signature_def_utils.get_signature_defs(
+        remove_tflite_model)
+    self.assertIsNone(signature_def_map.get(serving_key),
+                      'SignatureDef found, but should be missing')
+
+
+if __name__ == '__main__':
+  tf.test.main()