fix: fix tf.text nightly, also provide structure for smooth switching…

… between tensorflow (release) and tf-nightly dependencies (nightly and continuous jobs)

PiperOrigin-RevId: 641331580

WORKSPACE

@@ -88,6 +88,21 @@ load("@org_tensorflow//third_party/py:python_init_rules.bzl", "python_init_rules
 
 python_init_rules()
 
+load("//tensorflow_text:tftext.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"]
+    },
+    switch = {
+        "IS_NIGHTLY": "nightly"
+    }
+)
+
 load("@org_tensorflow//third_party/py:python_init_repositories.bzl", "python_init_repositories")
 
 python_init_repositories(

oss_scripts/configure.sh

@@ -41,12 +41,12 @@ else
   if is_macos; then
     #  Only Apple Silicon will be installed with tensorflow-macos.
     if [[ x"$(arch)" == x"arm64" ]]; then
-      pip install tensorflow-macos==2.13.0
+      pip install tensorflow-macos==2.16.1
     else
-      pip install tensorflow==2.13.0
+      pip install tensorflow==2.16.1
     fi
   else
-    pip install tensorflow==2.13.0
+    pip install tensorflow==2.16.1
   fi
 fi
 

oss_scripts/pip_package/BUILD

@@ -14,14 +14,14 @@ compile_pip_requirements(
         "--rebuild",
     ],
     generate_hashes = True,
-    requirements_in = "requirements.in",
+    requirements_in = "@release_or_nightly//:requirements.in",
     requirements_txt = REQUIREMENTS,
 )
 
 py_binary(
     name = "tensorflow_build_info",
     srcs = ["tensorflow_build_info.py"],
-    deps = ["@pypi_tf_nightly//:pkg"],
+    deps = ["@release_or_nightly//:tensorflow_pkg"],
 )
 
 sh_binary(

oss_scripts/pip_package/requirements.in

@@ -1,4 +1,4 @@
-tf-nightly
-tf_keras-nightly
-tensorflow_datasets
+tensorflow
+tf-keras
+tensorflow-datasets
 tensorflow-metadata

oss_scripts/pip_package/setup.py

@@ -32,7 +32,7 @@
 from setuptools.dist import Distribution
 
 project_name = 'tensorflow-text'
-project_version = '2.13.0'
+project_version = '2.16.1'
 
 
 class BinaryDistribution(Distribution):

oss_scripts/prepare_tf_dep.sh

@@ -16,8 +16,17 @@ fi
 HERMETIC_PYTHON_VERSION=$($installed_python  -c "import sys; print('.'.join(map(str, sys.version_info[:2])))")
 export HERMETIC_PYTHON_VERSION
 
+
+echo "TF_VERSION=$TF_VERSION"
+REQUIREMENTS_EXTRA_FLAGS="--upgrade"
+if [[ "$TF_VERSION" == *"rc"* ]]; then
+  REQUIREMENTS_EXTRA_FLAGS="$REQUIREMENTS_EXTRA_FLAGS --pre"
+fi
+
+bazel run //oss_scripts/pip_package:requirements.update -- $REQUIREMENTS_EXTRA_FLAGS
+
 # Update setup.nightly.py with current tf version.
-tf_version=$(bazel run  //oss_scripts/pip_package:tensorflow_build_info -- version)
+tf_version=$(bazel run //oss_scripts/pip_package:tensorflow_build_info -- version)
 echo "Updating setup.nightly.py to version $tf_version"
 sed -i $ext "s/project_version = '.*'/project_version = '${tf_version}'/" oss_scripts/pip_package/setup.nightly.py
 # Update __version__.
@@ -47,5 +56,3 @@ sed -E -i $ext "s/strip_prefix = \"tensorflow-.+\",/strip_prefix = \"tensorflow-
 sed -E -i $ext "s|\"https://github.com/tensorflow/tensorflow/archive/.+\.zip\"|\"https://github.com/tensorflow/tensorflow/archive/${commit_slug}.zip\"|" WORKSPACE
 prev_shasum=$(grep -A 1 -e "strip_prefix.*tensorflow-" WORKSPACE | tail -1 | awk -F '"' '{print $2}')
 sed -i $ext "s/sha256 = \"${prev_shasum}\",//" WORKSPACE
-
-bazel run //oss_scripts/pip_package:requirements.update -- --upgrade

oss_scripts/run_build.sh

@@ -13,20 +13,10 @@ fi
 # Run configure.
 source oss_scripts/configure.sh
 
-# Verify correct version of Bazel
-installed_bazel_version=$(bazel version | grep label | sed -e 's/.*: //')
-tf_bazel_version=$(head -n 1 .bazelversion)
-if [ "$installed_bazel_version" != "$tf_bazel_version" ]; then
-  echo "Incorrect version of Bazel installed."
-  echo "Version $tf_bazel_version should be installed, but found version ${installed_bazel_version}."
-  echo "Run oss_scripts/install_bazel.sh or manually install the correct version."
-  exit 1
-fi
-
 # Set tensorflow version
 if [[ $osname != "darwin" ]] || [[ ! $(sysctl -n machdep.cpu.brand_string) =~ "Apple" ]]; then
   source oss_scripts/prepare_tf_dep.sh
 fi
 
 # Build the pip package.
-bazel run --enable_runfiles oss_scripts/pip_package:build_pip_package "$(realpath .)"
+bazel run --enable_runfiles //oss_scripts/pip_package:build_pip_package -- "$(realpath .)"

tensorflow_text/core/kernels/sentencepiece/native.bzl.oss

@@ -80,8 +80,8 @@ def micore_tf_deps():
             # Standard references for Tensorflow when building for Linux. We use
             # an indirection via the alias targets below, to facilitate whitelisting
             # these deps in the mobile license presubmit checks.
-            "@pypi_tf_nightly//:libtensorflow_framework",
-            "@pypi_tf_nightly//:tf_header_lib",
+            "@release_or_nightly//:tensorflow_libtensorflow_framework",
+            "@release_or_nightly//:tensorflow_tf_header_lib",
         ],
 
     )

tensorflow_text/tftext.bzl

@@ -48,13 +48,13 @@ def py_tf_text_library(
             }),
             alwayslink = 1,
             deps = cc_op_kernels +
-                ["@org_tensorflow//tensorflow/lite/kernels/shim:tf_op_shim"] +
-                select({
-                    "@org_tensorflow//tensorflow:mobile": [
-                        "@org_tensorflow//tensorflow/core:portable_tensorflow_lib_lite",
-                ],
-                    "//conditions:default": [],
-            }),
+                   ["@org_tensorflow//tensorflow/lite/kernels/shim:tf_op_shim"] +
+                   select({
+                       "@org_tensorflow//tensorflow:mobile": [
+                           "@org_tensorflow//tensorflow/core:portable_tensorflow_lib_lite",
+                       ],
+                       "//conditions:default": [],
+                   }),
         )
 
         native.cc_binary(
@@ -129,6 +129,7 @@ def tf_cc_library(
     """
     if "kernel" in name:
         alwayslink = 1
+
     # These are "random" deps likely needed by each library (http://b/142433427)
     oss_deps = []
     oss_deps = oss_deps + _dedupe(deps, "@com_google_absl//absl/container:flat_hash_map")
@@ -143,8 +144,8 @@ def tf_cc_library(
             "@org_tensorflow//tensorflow/core:portable_tensorflow_lib_lite",
         ],
         "//conditions:default": [
-            "@pypi_tf_nightly//:libtensorflow_framework",
-            "@pypi_tf_nightly//:tf_header_lib",
+            "@release_or_nightly//:tensorflow_libtensorflow_framework",
+            "@release_or_nightly//:tensorflow_tf_header_lib",
         ] + tf_deps + oss_deps,
     })
     native.cc_library(
@@ -155,8 +156,8 @@ def tf_cc_library(
         copts = copts,
         compatible_with = compatible_with,
         testonly = testonly,
-        alwayslink = alwayslink)
-
+        alwayslink = alwayslink,
+    )
 
 def tflite_cc_library(
         name,
@@ -179,6 +180,7 @@ def tflite_cc_library(
         testonly: If library is only for testing
         alwayslink: If symbols should be exported
     """
+
     # Necessary build deps for tflite ops
     tflite_deps = [
         "@org_tensorflow//tensorflow/core:framework",
@@ -190,6 +192,7 @@ def tflite_cc_library(
         "@org_tensorflow//tensorflow/lite/kernels/shim:tflite_op_shim",
         "@org_tensorflow//tensorflow/lite/kernels/shim:tflite_op_wrapper",
     ]
+
     # These are "random" deps likely needed by each library (http://b/142433427)
     oss_deps = [
         "@com_google_absl//absl/container:flat_hash_map",
@@ -202,8 +205,8 @@ def tflite_cc_library(
             "@org_tensorflow//tensorflow/core:portable_tensorflow_lib_lite",
         ],
         "//conditions:default": [
-            "@pypi_tf_nightly//:libtensorflow_framework",
-            "@pypi_tf_nightly//:tf_header_lib",
+            "@release_or_nightly//:tensorflow_libtensorflow_framework",
+            "@release_or_nightly//:tensorflow_tf_header_lib",
         ] + oss_deps,
     })
     native.cc_library(
@@ -214,13 +217,59 @@ def tflite_cc_library(
         copts = copts,
         compatible_with = compatible_with,
         testonly = testonly,
-        alwayslink = alwayslink)
+        alwayslink = alwayslink,
+    )
 
 def extra_py_deps():
     return [
-        "@pypi_tf_nightly//:pkg",
-        "@pypi_tf_keras_nightly//:pkg",
+        "@release_or_nightly//:tensorflow_pkg",
+        "@release_or_nightly//:tf_keras_pkg",
         "@pypi_tensorflow_datasets//:pkg",
         "@pypi_tensorflow_metadata//:pkg",
     ]
 
+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,
+)

third_party/tensorflow/tf_configure.bzl

@@ -9,6 +9,24 @@ load("@python_version_repo//:py_version.bzl", "REQUIREMENTS_WITH_LOCAL_WHEELS")
 load("@rules_python//python:pip.bzl", "package_annotation", "pip_parse")
 
 def tf_configure():
+    tensorflow_annotation = """
+cc_library(
+    name = "tf_header_lib",
+    hdrs = glob(["site-packages/tensorflow/include/**/*"]),
+    strip_include_prefix="site-packages/tensorflow/include/",
+    visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "libtensorflow_framework",
+    srcs = select({
+        "//conditions:default": ["site-packages/tensorflow/libtensorflow_framework.so.2"],
+        "@bazel_tools//src/conditions:darwin":["site-packages/tensorflow/libtensorflow_framework.2.dylib"],
+        "@bazel_tools//src/conditions:darwin_x86_64": ["site-packages/tensorflow/libtensorflow_framework.2.dylib"],
+    }),
+    visibility = ["//visibility:public"],
+)
+"""
     pip_parse(
         name = "pypi",
         annotations = {
@@ -30,25 +48,11 @@ cc_library(
 )
 """,
             ),
+            "tensorflow": package_annotation(
+                additive_build_content = tensorflow_annotation,
+            ),
             "tf-nightly": package_annotation(
-                additive_build_content = """
-cc_library(
-    name = "tf_header_lib",
-    hdrs = glob(["site-packages/tensorflow/include/**/*"]),
-    strip_include_prefix="site-packages/tensorflow/include/",
-    visibility = ["//visibility:public"],
-)
-
-cc_library(
-    name = "libtensorflow_framework",
-    srcs = select({
-        "//conditions:default": ["site-packages/tensorflow/libtensorflow_framework.so.2"],
-        "@bazel_tools//src/conditions:darwin":["site-packages/tensorflow/libtensorflow_framework.2.dylib"],
-        "@bazel_tools//src/conditions:darwin_x86_64": ["site-packages/tensorflow/libtensorflow_framework.2.dylib"],
-    }),
-    visibility = ["//visibility:public"],
-)
-""",
+                additive_build_content = tensorflow_annotation,
             ),
         },
         python_interpreter_target = interpreter,