Stop statically linking TensorFlow into every OSS TF.Text op library.

WORKSPACE

@@ -72,7 +72,7 @@ http_archive(
     name = "org_tensorflow",
     patch_args = ["-p1"],
     patches = [
-        "//third_party/tensorflow:tensorflow.xla_if_static.bzl.patch",
+        "//third_party/tensorflow:tensorflow.shim_sources.patch",
     ],
     sha256 = "b7fc9646df01c14968e1546ad430a72f6ecc26867c085266f71f43968a6d5f5b",
     strip_prefix = "tensorflow-2.22.0-rc0",

tensorflow_text/core/kernels/fast_wordpiece_tokenizer_model_builder.cc

@@ -22,6 +22,7 @@
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"

tensorflow_text/core/kernels/sentence_fragmenter_v2.cc

@@ -21,7 +21,6 @@
 #include "absl/strings/string_view.h"
 #include "icu4c/source/common/unicode/uchar.h"
 #include "icu4c/source/common/unicode/utf8.h"
-#include "tensorflow/core/lib/core/status.h"
 
 namespace tensorflow {
 namespace text {

tensorflow_text/core/kernels/sentence_fragmenter_v2.h

@@ -34,12 +34,12 @@
 #ifndef TENSORFLOW_TEXT_CORE_KERNELS_SENTENCE_FRAGMENTER_V2_H_
 #define TENSORFLOW_TEXT_CORE_KERNELS_SENTENCE_FRAGMENTER_V2_H_
 
+#include <cstdint>
 #include <vector>
 
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
 #include "icu4c/source/common/unicode/utypes.h"
-#include "tensorflow/core/lib/core/status.h"
 
 namespace tensorflow {
 namespace text {
@@ -140,7 +140,7 @@ struct SentenceFragment {
     HAS_SENTENTIAL_CLOSE_PAREN = 0x0008,  // e.g.: (Mushrooms are fungi!)
   };
   // A mask of the above listed properties.
-  uint32 properties = 0;
+  uint32_t properties = 0;
   int terminal_punc_token = -1;
 };
 

tensorflow_text/core/kernels/sentence_fragmenter_v2_test.cc

@@ -736,7 +736,7 @@ TEST_F(IsEmoticonTest, IsEmoticon) {
 TEST(SentenceFragmenterTest, Basic) {
   //                             1
   //                   012345678901234
-  string test_input = "Hello. Foo bar!";
+  std::string test_input = "Hello. Foo bar!";
   SentenceFragmenterV2 fragmenter(test_input);
   std::vector<SentenceFragment> fragments;
   EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok());
@@ -749,7 +749,7 @@ TEST(SentenceFragmenterTest, Basic) {
 TEST(SentenceFragmenterTest, BasicEllipsis) {
   //                             1
   //                   012345678901234
-  string test_input = "Hello...foo bar";
+  std::string test_input = "Hello...foo bar";
   SentenceFragmenterV2 fragmenter(test_input);
   std::vector<SentenceFragment> fragments;
   EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok());
@@ -763,7 +763,7 @@ TEST(SentenceFragmenterTest, BasicEllipsis) {
 TEST(SentenceFragmenterTest, Parentheses) {
   //                             1         2
   //                   012345678901234567890123456789
-  string test_input = "Hello (who are you...) foo bar";
+  std::string test_input = "Hello (who are you...) foo bar";
   SentenceFragmenterV2 fragmenter(test_input);
   std::vector<SentenceFragment> fragments;
   EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok());
@@ -776,7 +776,7 @@ TEST(SentenceFragmenterTest, Parentheses) {
 TEST(SentenceFragmenterTest, MidFragmentParentheses) {
   //                             1         2
   //                   012345678901234567890123456789
-  string test_input = "Hello (who are you) world? Foo bar";
+  std::string test_input = "Hello (who are you) world? Foo bar";
   SentenceFragmenterV2 fragmenter(test_input);
   std::vector<SentenceFragment> fragments;
   EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok());
@@ -789,7 +789,7 @@ TEST(SentenceFragmenterTest, MidFragmentParentheses) {
 TEST(SentenceFragmenterTest, PunctuationAfterParentheses) {
   //                             1         2
   //                   01234567890123456789012345678
-  string test_input = "Hello (who are you)? Foo bar!";
+  std::string test_input = "Hello (who are you)? Foo bar!";
   SentenceFragmenterV2 fragmenter(test_input);
   std::vector<SentenceFragment> fragments;
   EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok());
@@ -802,7 +802,7 @@ TEST(SentenceFragmenterTest, PunctuationAfterParentheses) {
 TEST(SentenceFragmenterTest, ManyFinalPunctuations) {
   //                             1         2
   //                   0123456789012345678901234
-  string test_input = "Hello!!!!! Who are you??";
+  std::string test_input = "Hello!!!!! Who are you??";
   SentenceFragmenterV2 fragmenter(test_input);
   std::vector<SentenceFragment> fragments;
   EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok());
@@ -815,7 +815,7 @@ TEST(SentenceFragmenterTest, ManyFinalPunctuations) {
 TEST(SentenceFragmenterTest, NewLine) {
   //                             1         2             3
   //                   012345678901234567890 1 23456 7 89012 3 45678
-  string test_input = "Who let the dogs out?\r\nWho?\r\nWho?\r\nWho?";
+  std::string test_input = "Who let the dogs out?\r\nWho?\r\nWho?\r\nWho?";
   SentenceFragmenterV2 fragmenter(test_input);
   std::vector<SentenceFragment> fragments;
   EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok());
@@ -832,7 +832,7 @@ TEST(SentenceFragmenterTest, NewLine) {
 TEST(SentenceFragmenterTest, WhiteSpaceInPunctuation) {
   //                             1         2
   //                   0123456789012345678901234
-  string test_input = "Hello?? !!! Who are you??";
+  std::string test_input = "Hello?? !!! Who are you??";
   SentenceFragmenterV2 fragmenter(test_input);
   std::vector<SentenceFragment> fragments;
   EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok());
@@ -850,7 +850,7 @@ TEST(FragmentBoundaryMatchTest, NoStateChange) {
   FragmentBoundaryMatch f;
   //                   ||
   //                   012345678901234
-  string test_input = "Hello...foo bar";
+  std::string test_input = "Hello...foo bar";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_FALSE(f.GotTerminalPunc());
@@ -864,7 +864,7 @@ TEST(FragmentBoundaryMatchTest, BasicEllipsis) {
   FragmentBoundaryMatch f;
   //                   |  |
   //                   0123456789
-  string test_input = "...foo bar";
+  std::string test_input = "...foo bar";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_TRUE(f.GotTerminalPunc());
@@ -878,7 +878,7 @@ TEST(FragmentBoundaryMatchTest, BasicPeriod) {
   FragmentBoundaryMatch f;
   //                   ||
   //                   0123456789
-  string test_input = ". Foo bar";
+  std::string test_input = ". Foo bar";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_TRUE(f.GotTerminalPunc());
@@ -892,7 +892,7 @@ TEST(FragmentBoundaryMatchTest, BasicAcronym) {
   FragmentBoundaryMatch f;
   //                   |  |
   //                   0123456789
-  string test_input = "A.B. xyz";
+  std::string test_input = "A.B. xyz";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_TRUE(f.GotTerminalPunc());
@@ -906,7 +906,7 @@ TEST(FragmentBoundaryMatchTest, LongerAcronym) {
   FragmentBoundaryMatch f;
   //                   |    |
   //                   0123456789
-  string test_input = "I.B.M. yo";
+  std::string test_input = "I.B.M. yo";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_TRUE(f.GotTerminalPunc());
@@ -920,7 +920,7 @@ TEST(FragmentBoundaryMatchTest, Emoticon) {
   FragmentBoundaryMatch f;
   //                   |   |
   //                   0123456789012
-  string test_input = ">:-( hello...";
+  std::string test_input = ">:-( hello...";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_TRUE(f.GotTerminalPunc());
@@ -934,7 +934,7 @@ TEST(FragmentBoundaryMatchTest, ParensWithEllipsis) {
   FragmentBoundaryMatch f;
   //                   ||
   //                   0123456789012345
-  string test_input = ".foo...) foo bar";
+  std::string test_input = ".foo...) foo bar";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_TRUE(f.GotTerminalPunc());
@@ -948,7 +948,7 @@ TEST(FragmentBoundaryMatchTest, ClosingParenWithEllipsis) {
   FragmentBoundaryMatch f;
   //                   |  |
   //                   012345678901
-  string test_input = "...) foo bar";
+  std::string test_input = "...) foo bar";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_TRUE(f.GotTerminalPunc());
@@ -962,7 +962,7 @@ TEST(FragmentBoundaryMatchTest, BeginAndEndParenWithEllipsis) {
   FragmentBoundaryMatch f;
   //                   ||
   //                   0123456789012
-  string test_input = "(...) foo bar";
+  std::string test_input = "(...) foo bar";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_FALSE(f.GotTerminalPunc());
@@ -986,7 +986,7 @@ TEST(FragmentBoundaryMatchTest, AcronymInSentence) {
   FragmentBoundaryMatch f;
   //                   |   |
   //                   0123456789012
-  string test_input = "U.S. don't be surprised.";
+  std::string test_input = "U.S. don't be surprised.";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_TRUE(f.GotTerminalPunc());
@@ -1000,7 +1000,7 @@ TEST(FragmentBoundaryMatchTest, HelloWithEllipsis) {
   FragmentBoundaryMatch f;
   //                   ||
   //                   01234567890
-  string test_input = "o...foo bar";
+  std::string test_input = "o...foo bar";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_FALSE(f.GotTerminalPunc());
@@ -1024,7 +1024,7 @@ TEST(FragmentBoundaryMatchTest, ThreeStatesWithClosigParen) {
   FragmentBoundaryMatch f;
   //                   ||
   //                   0123456789012
-  string test_input = "w...) foo bar";
+  std::string test_input = "w...) foo bar";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_FALSE(f.GotTerminalPunc());
@@ -1068,7 +1068,7 @@ TEST(FragmentBoundaryMatchTest, NoTransition) {
   FragmentBoundaryMatch f;
   //                   |  |
   //                   0123456789012
-  string test_input = "...foo bar";
+  std::string test_input = "...foo bar";
   int index = 0;
   EXPECT_TRUE(f.Advance(index, test_input));
   EXPECT_TRUE(f.GotTerminalPunc());

tensorflow_text/core/kernels/whitespace_tokenizer_kernel_template.h

@@ -20,7 +20,6 @@
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/lite/kernels/shim/op_kernel.h"
 #include "tensorflow/lite/kernels/shim/shape.h"
 #include "tensorflow/lite/kernels/shim/status_macros.h"

tensorflow_text/tftext.bzl

@@ -9,6 +9,31 @@ load(
     _pywrap_library = "pywrap_library",
 )
 
+def _local_tf_op_shim():
+    """Returns a package-local `tf_op_shim` built against the TF wheel headers.
+
+    Only the three shim sources are compiled here. They are linked dynamically
+    against libtensorflow_framework.so (like any TF custom op), instead of
+    depending on @org_tensorflow//tensorflow/lite/kernels/shim:tf_op_shim, which
+    statically pulls TensorFlow implementation targets (protos, protobuf
+    runtime, lib_internal, ...) into every _<op>.so.
+    """
+    name = "tf_op_shim_oss"
+    if not native.existing_rule(name):
+        tf_cc_library(
+            name = name,
+            srcs = [
+                "@org_tensorflow//tensorflow/lite/kernels/shim:shape.cc",
+                "@org_tensorflow//tensorflow/lite/kernels/shim:tf_op_shim.cc",
+                "@org_tensorflow//tensorflow/lite/kernels/shim:tf_tensor_view.cc",
+            ],
+            deps = [
+                "@com_google_absl//absl/algorithm:container",
+                "@com_google_absl//absl/types:span",
+            ],
+        )
+    return ":" + name
+
 def py_tf_text_library(
         name,
         srcs = [],
@@ -55,7 +80,7 @@ def py_tf_text_library(
             }),
             alwayslink = 1,
             deps = cc_op_kernels +
-                   ["@org_tensorflow//tensorflow/lite/kernels/shim:tf_op_shim"] +
+                   [_local_tf_op_shim()] +
                    select({
                        "@org_tensorflow//tensorflow:mobile": [
                            "@org_tensorflow//tensorflow/core:portable_tensorflow_lib_lite",

third_party/tensorflow/tensorflow.shim_sources.patch

@@ -0,0 +1,33 @@
+diff --git a/tensorflow/lite/kernels/shim/BUILD b/tensorflow/lite/kernels/shim/BUILD
+--- a/tensorflow/lite/kernels/shim/BUILD
++++ b/tensorflow/lite/kernels/shim/BUILD
+@@ -13,6 +13,15 @@
+     licenses = ["notice"],
+ )
+
++exports_files(
++    [
++        "shape.cc",
++        "tf_op_shim.cc",
++        "tf_tensor_view.cc",
++    ],
++    visibility = ["//visibility:public"],
++)
++
+ filegroup(
+     name = "tf_headers",
+     srcs = [
+diff --git a/tensorflow/lite/kernels/shim/tf_op_shim.cc b/tensorflow/lite/kernels/shim/tf_op_shim.cc
+--- a/tensorflow/lite/kernels/shim/tf_op_shim.cc
++++ b/tensorflow/lite/kernels/shim/tf_op_shim.cc
+@@ -79,8 +79,8 @@ absl::StatusOr<AttrValue> TfInitContext::GetAttr(
+     return absl::InvalidArgumentError(
+         absl::StrCat("Non-existent attribute: ", attr_name, "\nop def:\n",
+                      context_->def().DebugString()));
+-  const auto& attr_value = context_->def().attr().at(attr_name);
+-  return TfAttrValueToShimAttrValue(attr_value);
++  const auto* attr_value = ::tensorflow::AttrSlice(context_->def()).Find(attr_name);
++  return TfAttrValueToShimAttrValue(*attr_value);
+ }
+
+ TfInvokeContext::TfInvokeContext(::tensorflow::OpKernelContext* context)