Test fixing tensorflow-text nightly builds.

WORKSPACE

@@ -4,34 +4,34 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
 http_archive(
     name = "icu",
-    strip_prefix = "icu-release-64-2",
+    build_file = "//third_party/icu:BUILD.bzl",
+    patch_args = ["-p1"],
+    patches = ["//third_party/icu:udata.patch"],
     sha256 = "dfc62618aa4bd3ca14a3df548cd65fe393155edd213e49c39f3a30ccd618fc27",
+    strip_prefix = "icu-release-64-2",
     urls = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/unicode-org/icu/archive/release-64-2.zip",
         "https://github.com/unicode-org/icu/archive/release-64-2.zip",
     ],
-    build_file = "//third_party/icu:BUILD.bzl",
-    patches = ["//third_party/icu:udata.patch"],
-    patch_args = ["-p1"],
 )
 
 http_archive(
     name = "com_google_sentencepiece",
-    strip_prefix = "sentencepiece-0.1.96",
+    build_file = "//third_party/sentencepiece:BUILD",
+    patch_args = ["-p1"],
+    patches = ["//third_party/sentencepiece:sp.patch"],
     sha256 = "8409b0126ebd62b256c685d5757150cf7fcb2b92a2f2b98efb3f38fc36719754",
+    strip_prefix = "sentencepiece-0.1.96",
     urls = [
-        "https://github.com/google/sentencepiece/archive/refs/tags/v0.1.96.zip"
+        "https://github.com/google/sentencepiece/archive/refs/tags/v0.1.96.zip",
     ],
-    build_file = "//third_party/sentencepiece:BUILD",
-    patches = ["//third_party/sentencepiece:sp.patch"],
-    patch_args = ["-p1"],
 )
 
 http_archive(
     name = "cppitertools",
-    urls = ["https://github.com/ryanhaining/cppitertools/archive/refs/tags/v2.0.zip"],
     sha256 = "e56741b108d6baced98c4ccd83fd0d5a545937f2845978799c28d0312c0dee3d",
     strip_prefix = "cppitertools-2.0",
+    urls = ["https://github.com/ryanhaining/cppitertools/archive/refs/tags/v2.0.zip"],
 )
 
 http_archive(
@@ -56,10 +56,10 @@ http_archive(
 
 http_archive(
     name = "org_tensorflow",
-    strip_prefix = "tensorflow-40998f44c0c500ce0f6e3b1658dfbc54f838a82a",
-    sha256 = "5a5bc4599964c71277dcac0d687435291e5810d2ac2f6283cc96736febf73aaf",
+    sha256 = "c75a6fde64ba60d435f795a1a63f5caec0db8b352f544e222b8f51a1f6fec63e",
+    strip_prefix = "tensorflow-735467e89ccfd7ace190363412bb5698164628b5",
     urls = [
-        "https://github.com/tensorflow/tensorflow/archive/40998f44c0c500ce0f6e3b1658dfbc54f838a82a.zip"
+        "https://github.com/tensorflow/tensorflow/archive/735467e89ccfd7ace190363412bb5698164628b5.zip",
     ],
 )
 
@@ -74,13 +74,13 @@ http_archive(
 
 http_archive(
     name = "pybind11",
+    build_file = "//third_party/pybind11:BUILD.bzl",
+    sha256 = "efc901aa0aab439a3fea6efeaf930b5a349fb06394bf845c64ce15a9cf8f0240",
+    strip_prefix = "pybind11-2.13.4",
     urls = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/pybind/pybind11/archive/v2.13.4.tar.gz",
         "https://github.com/pybind/pybind11/archive/v2.13.4.tar.gz",
     ],
-    sha256 = "efc901aa0aab439a3fea6efeaf930b5a349fb06394bf845c64ce15a9cf8f0240",
-    strip_prefix = "pybind11-2.13.4",
-    build_file = "//third_party/pybind11:BUILD.bzl",
 )
 
 http_archive(
@@ -95,31 +95,35 @@ load("@org_tensorflow//third_party/py:python_init_rules.bzl", "python_init_rules
 
 python_init_rules()
 
-load("//tensorflow_text:tftext.bzl", "py_deps_profile")
+load("//tensorflow_text:py_deps_profile.bzl", "py_deps_profile")
 
 py_deps_profile(
     name = "release_or_nightly",
-    requirements_in = "//oss_scripts/pip_package:requirements.in",
-    pip_repo_name = "pypi",
     deps_map = {
-        "tensorflow": ["tf-nightly", "tf_header_lib", "libtensorflow_framework"],
-        "tf-keras": ["tf-keras-nightly"]
+        "tensorflow": [
+            "tf-nightly",
+            "tf_header_lib",
+            "libtensorflow_framework",
+        ],
+        "tf-keras": ["tf-keras-nightly"],
     },
+    pip_repo_name = "pypi",
+    requirements_in = "//oss_scripts/pip_package:requirements.in",
     switch = {
-        "IS_NIGHTLY": "nightly"
-    }
+        "IS_NIGHTLY": "nightly",
+    },
 )
 
 load("@org_tensorflow//third_party/py:python_init_repositories.bzl", "python_init_repositories")
 
 python_init_repositories(
+    default_python_version = "system",
     requirements = {
         "3.9": "//oss_scripts/pip_package:requirements_lock_3_9.txt",
         "3.10": "//oss_scripts/pip_package:requirements_lock_3_10.txt",
         "3.11": "//oss_scripts/pip_package:requirements_lock_3_11.txt",
         "3.12": "//oss_scripts/pip_package:requirements_lock_3_12.txt",
     },
-    default_python_version = "system",
 )
 
 load("@org_tensorflow//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
@@ -136,18 +140,28 @@ install_deps()
 
 # Initialize TensorFlow dependencies.
 load("@org_tensorflow//tensorflow:workspace3.bzl", "tf_workspace3")
+
 tf_workspace3()
+
 load("@org_tensorflow//tensorflow:workspace2.bzl", "tf_workspace2")
+
 tf_workspace2()
+
 load("@org_tensorflow//tensorflow:workspace1.bzl", "tf_workspace1")
+
 tf_workspace1()
+
 load("@org_tensorflow//tensorflow:workspace0.bzl", "tf_workspace0")
+
 tf_workspace0()
 
 # Set up Android.
 load("@org_tensorflow//third_party/android:android_configure.bzl", "android_configure")
-android_configure(name="local_config_android")
+
+android_configure(name = "local_config_android")
+
 load("@local_config_android//:android.bzl", "android_workspace")
+
 android_workspace()
 
 load(
@@ -158,7 +172,22 @@ load(
 python_wheel_version_suffix_repository(name = "tf_wheel_version_suffix")
 
 load(
-    "@local_xla//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
+    "@rules_ml_toolchain//cc/deps:cc_toolchain_deps.bzl",
+    "cc_toolchain_deps",
+)
+
+cc_toolchain_deps()
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_cuda")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64_cuda")
+
+load(
+    "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
     "cuda_json_init_repository",
 )
 
@@ -170,7 +199,7 @@ load(
     "CUDNN_REDISTRIBUTIONS",
 )
 load(
-    "@local_xla//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
+    "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
     "cuda_redist_init_repositories",
     "cudnn_redist_init_repository",
 )
@@ -184,21 +213,21 @@ cudnn_redist_init_repository(
 )
 
 load(
-    "@local_xla//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
+    "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
     "cuda_configure",
 )
 
 cuda_configure(name = "local_config_cuda")
 
 load(
-    "@local_xla//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
+    "@rules_ml_toolchain//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
     "nccl_redist_init_repository",
 )
 
 nccl_redist_init_repository()
 
 load(
-    "@local_xla//third_party/nccl/hermetic:nccl_configure.bzl",
+    "@rules_ml_toolchain//third_party/nccl/hermetic:nccl_configure.bzl",
     "nccl_configure",
 )
 

oss_scripts/configure.sh

@@ -54,6 +54,12 @@ curl https://raw.githubusercontent.com/tensorflow/tensorflow/master/.bazelrc -o
 sed -i -e 's/build --noincompatible_remove_legacy_whole_archive//' .bazelrc
 
 write_to_bazelrc "build:manylinux2014 --config=release_cpu_linux"
+# JUST FOR TESTING
+git clone https://github.com/tensorflow/tensorflow.git org_tensorflow
+cd org_tensorflow
+git checkout 735467e89ccfd7ace190363412bb5698164628b5
+cd ..
+write_to_bazelrc "common:linux --override_repository=org_tensorflow=./org_tensorflow"
 
 if (which python3) | grep -q "python3"; then
   installed_python="python3"
@@ -76,8 +82,16 @@ if [[ "$TF_VERSION" == *"rc"* ]]; then
   REQUIREMENTS_EXTRA_FLAGS="$REQUIREMENTS_EXTRA_FLAGS --pre"
 fi
 
+# JUST FOR TESTING
+wget https://storage.googleapis.com/ml-sysroot-testing/tensorflow_metadata-1.18.0.dev0-py3-none-any.whl
+mv tensorflow_metadata-1.18.0.dev0-py3-none-any.whl /github/tensorflow_metadata-1.18.0.dev0-py3-none-any.whl
+echo "tensorflow-metadata @ file:///github/tensorflow_metadata-1.18.0.dev0-py3-none-any.whl" >> oss_scripts/pip_package/requirements.in
+
 bazel run //oss_scripts/pip_package:requirements.update -- $REQUIREMENTS_EXTRA_FLAGS
 
+# JUST FOR TESTING
+sed -i 's#file://tensorflow_metadata-1.18.0.dev0-py3-none-any.whl#file:///github/tensorflow_metadata-1.18.0.dev0-py3-none-any.whl#g' oss_scripts/pip_package/requirements_lock_3_9.txt
+
 TF_ABIFLAG=$(bazel run //oss_scripts/pip_package:tensorflow_build_info -- abi)
 SHARED_LIBRARY_NAME="libtensorflow_framework.so.2"
 if is_macos; then

oss_scripts/pip_package/BUILD

@@ -1,6 +1,7 @@
 # Tools for building the TF.Text pip package.
-load("@python//:defs.bzl", "compile_pip_requirements")
+
 load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
+load("@rules_python//python:pip.bzl", "compile_pip_requirements")
 
 package(default_visibility = ["//visibility:private"])
 

oss_scripts/pip_package/requirements.in

@@ -1,8 +1,7 @@
 setuptools==70.0.0
 dm-tree==0.1.8  # Limit for macos support.
 numpy
-protobuf==4.25.3  # b/397977335 - Fix crash on python 3.9, 3.10.
 tensorflow
 tf-keras
 tensorflow-datasets
-tensorflow-metadata
+#tensorflow-metadata

tensorflow_text/BUILD

@@ -1,7 +1,8 @@
-# Placeholder: load py_library
-# Placeholder: load py_test
+load("@org_tensorflow//tensorflow:tensorflow.bzl", "if_oss")
 load("@org_tensorflow//tensorflow/lite:build_def.bzl", "tflite_cc_shared_object")
 load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_python//python:py_library.bzl", "py_library")
+load("@rules_python//python:py_test.bzl", "py_test")
 load("//tensorflow_text:tftext.bzl", "extra_py_deps", "py_tf_text_library")
 # [internal] load build_test.bzl
 
@@ -123,9 +124,11 @@ py_library(
         ":wordpiece_tokenizer",
         ":wordshape_ops",
         # python/util:all_util tensorflow dep,
-        "//tensorflow_text/core/pybinds:tflite_registrar",
         "//tensorflow_text/tools/wordpiece_vocab",
-    ] + extra_py_deps(),
+    ] + if_oss(
+        ["//tensorflow_text/core/pybinds:tflite_registrar_library"],
+        ["//tensorflow_text/core/pybinds:tflite_registrar"],
+    ) + extra_py_deps(),
 )
 
 # build_test

tensorflow_text/core/kernels/BUILD

@@ -337,6 +337,7 @@ cc_library(
     deps = [
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@darts_clone",
     ],
 )

tensorflow_text/core/kernels/darts_clone_trie_builder.cc

@@ -19,7 +19,11 @@
 #include <numeric>
 
 #include "absl/container/flat_hash_set.h"
+#include "absl/strings/str_cat.h"
 #include "include/darts.h"
+#define REGISTER_UNARY_VARIANT_DEVICE_COPY_HANDLER \
+  1  // TODO(b/210217594): Remove once this is the default.
+#include "tensorflow/core/framework/variant_op_registry.h"
 
 namespace tensorflow {
 namespace text {

tensorflow_text/core/pybinds/BUILD

@@ -1,7 +1,10 @@
 # Code that exposes C++ libraries to Python via pybind11.
 
-# Placeholder: load py_test
-load("@org_tensorflow//tensorflow:tensorflow.bzl", "pybind_extension")
+load("@rules_python//python:py_test.bzl", "py_test")
+
+# This comment separates the load statements, so copybara reverse transformation works.
+
+load("@org_tensorflow//tensorflow:tensorflow.default.bzl", "pybind_extension", "pywrap_library")
 
 licenses(["notice"])
 
@@ -11,6 +14,11 @@ package(default_visibility = [
     "//tensorflow_text:__subpackages__",
 ])
 
+pywrap_library(
+    name = "tflite_registrar_library",
+    deps = [":tflite_registrar"],
+)
+
 pybind_extension(
     name = "tflite_registrar",
     srcs = [

tensorflow_text/py_deps_profile.bzl

@@ -0,0 +1,49 @@
+"""
+Repository rule for open source dependencies.
+"""
+
+def _py_deps_profile_impl(ctx):
+    is_switch = False
+    for var_name, var_val in ctx.attr.switch.items():
+        is_switch = is_switch or ctx.os.environ.get(var_name, "") == var_val
+
+    prefix = ctx.attr.pip_repo_name
+    reqirements_name = ctx.attr.requirements_in.name
+    requirements_in_substitutions = {}
+    build_content = ['exports_files(["{}"])'.format(reqirements_name)]
+    for k, v in ctx.attr.deps_map.items():
+        repo_name = v[0] if is_switch else k
+        requirements_in_substitutions[k + "\n"] = repo_name + "\n"
+        requirements_in_substitutions[k + "\r\n"] = repo_name + "\r\n"
+        aliased_targets = ["pkg"] + v[1:]
+        norm_repo_name = repo_name.replace("-", "_")
+        norm_alas_name = k.replace("-", "_")
+        for target in aliased_targets:
+            alias_name = "{}_{}".format(norm_alas_name, target)
+            alias_value = "@{}_{}//:{}".format(prefix, norm_repo_name, target)
+            build_content.append("""
+alias(
+    name = "{}",
+    actual = "{}",
+    visibility = ["//visibility:public"]
+)
+""".format(alias_name, alias_value))
+
+    ctx.file("BUILD", "".join(build_content))
+    ctx.template(
+        reqirements_name,
+        ctx.attr.requirements_in,
+        executable = False,
+        substitutions = requirements_in_substitutions,
+    )
+
+py_deps_profile = repository_rule(
+    implementation = _py_deps_profile_impl,
+    attrs = {
+        "requirements_in": attr.label(mandatory = True),
+        "deps_map": attr.string_list_dict(mandatory = True),
+        "pip_repo_name": attr.string(mandatory = True),
+        "switch": attr.string_dict(mandatory = True),
+    },
+    local = True,
+)