Add docker build scripts and enable aarch64 pip wheels.

PiperOrigin-RevId: 692239607

oss_scripts/BUILD.oss

@@ -0,0 +1 @@
+licenses(["notice"])

oss_scripts/build.Dockerfile.aarch64

@@ -0,0 +1,37 @@
+# Constructs the environment within which we will build the tensorflow-text
+# pip wheels.
+
+FROM linaro/tensorflow-arm64-build:2.18-multipython
+LABEL maintainer="TensorFlow-Text team <[email protected]>"
+
+ARG PYTHON_VERSION
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHON_BIN_PATH=/usr/bin/python${PYTHON_VERSION}
+
+# Install supplementary Python interpreters
+RUN ln -s ${PYTHON_BIN_PATH} /usr/local/bin/python && \
+    ln -s ${PYTHON_BIN_PATH} /usr/local/bin/python3 && \
+    ln -s ${PYTHON_BIN_PATH} /usr/bin/python
+
+RUN --mount=type=cache,target=/var/cache/apt \
+  apt update && \
+  apt install -yqq \
+    apt-utils \
+    build-essential \
+    checkinstall \
+    libffi-dev
+
+# Install pip dependencies needed for tensorflow-text
+RUN --mount=type=cache,target=/root/.cache \
+  ${PYTHON_BIN_PATH} -m pip install -U pip && \
+  ${PYTHON_BIN_PATH} -m pip install -U \
+    absl-py \
+    auditwheel \
+    etils[epath] \
+    patchelf \
+    setuptools \
+    twine \
+    wheel;
+
+WORKDIR "/tmp/tensorflow_text"

oss_scripts/build.Dockerfile.x86_64

@@ -0,0 +1,32 @@
+# Constructs the environment within which we will build the tensorflow-text
+# pip wheels.
+
+ARG PYTHON_VERSION
+FROM tensorflow/build:2.18-python${PYTHON_VERSION}
+LABEL maintainer="TensorFlow-Text team <[email protected]>"
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install supplementary Python interpreters
+RUN mkdir /tmp/python
+RUN --mount=type=cache,target=/var/cache/apt \
+  apt update && \
+  apt install -yqq \
+    apt-utils \
+    build-essential \
+    checkinstall \
+    libffi-dev
+
+# Install pip dependencies needed for tensorflow text.
+RUN --mount=type=cache,target=/root/.cache \
+  python${PYTHON_VERSION} -m pip install -U pip && \
+  python${PYTHON_VERSION} -m pip install -U \
+    absl-py \
+    auditwheel \
+    etils[epath] \
+    patchelf \
+    setuptools \
+    twine \
+    wheel;
+
+WORKDIR "/tmp/tensorflow_text"

oss_scripts/configure.sh

@@ -32,23 +32,9 @@ function is_macos() {
 # Remove .bazelrc if it already exist
 [ -e .bazelrc ] && rm .bazelrc
 
-if [[ $(pip show tensorflow) == *tensorflow* ]] ||
-   [[ $(pip show tensorflow-macos) == *tensorflow-macos* ]] ||
-   [[ $(pip show tf-nightly) == *tf-nightly* ]]; then
-  echo 'Using installed tensorflow.'
-else
-  echo 'Installing tensorflow.'
-  if is_macos; then
-    #  Only Apple Silicon will be installed with tensorflow-macos.
-    if [[ x"$(arch)" == x"arm64" ]]; then
-      pip install tensorflow-macos==2.16.1
-    else
-      pip install tensorflow==2.16.1
-    fi
-  else
-    pip install tensorflow==2.16.1
-  fi
-fi
+echo 'Installing tensorflow.'
+# For main branch, install nightly.  CHANGE THIS ON RELEASE BRANCHES.
+pip install tf-nightly
 
 # Copy the current bazelversion of TF.
 curl https://raw.githubusercontent.com/tensorflow/tensorflow/master/.bazelversion -o .bazelversion

oss_scripts/docker_builds.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+# This script builds a docker container, and pip wheels for all supported
+# Python versions.  Run from the root tensorflow-text directory:
+#
+# ./oss_scripts/docker_builds.sh
+
+set -e -x
+
+# Clean previous images.
+for PYTHON_VERSION in 3.9 3.10 3.11 3.12
+do
+  docker rmi -f tensorflow_text:${PYTHON_VERSION} || true
+done
+
+ARCH=$(uname -m)
+BUILD_ARGS=()
+if [ "$ARCH" == "x86_64" ]; then
+  BUILD_ARGS+=("[email protected]${PYTHON_VERSION}_config_cuda//crosstool:toolchain")
+  BUILD_ARGS+=("--config=release_cpu_linux")
+  AUDITWHEEL_PLATFORM="manylinux2014_x86_64"
+elif [ "$ARCH" == "aarch64" ]; then
+  BUILD_ARGS+=("--crosstool_top=@ml2014_aarch64_config_aarch64//crosstool:toolchain")
+  AUDITWHEEL_PLATFORM="manylinux2014_aarch64"
+fi
+
+# Build wheel for each Python version.
+for PYTHON_VERSION in 3.9 3.10 3.11
+do
+  DOCKER_BUILDKIT=1 docker build --progress=plain --no-cache \
+    --build-arg PYTHON_VERSION=${PYTHON_VERSION} \
+    -t tensorflow_text:${PYTHON_VERSION} - < "oss_scripts/build.Dockerfile.${ARCH}"
+
+  docker run --rm -a stdin -a stdout -a stderr \
+    --env PYTHON_VERSION=${PYTHON_VERSION} \
+    --env BUILD_ARGS=${BUILD_ARGS} \
+    --env AUDITWHEEL_PLATFORM=${AUDITWHEEL_PLATFORM} \
+    -v $PWD:/tmp/tensorflow_text \
+    --name tensorflow_text tensorflow_text:${PYTHON_VERSION} \
+    bash oss_scripts/run_build.sh
+done

oss_scripts/pip_package/requirements.in

@@ -1,4 +1,5 @@
+ai-edge-litert
 tensorflow
 tf-keras
 tensorflow-datasets
-tensorflow-metadata
+tensorflow-metadata

oss_scripts/requirements/BUILD.oss

@@ -0,0 +1,30 @@
+load("@rules_python//python:pip.bzl", "compile_pip_requirements")
+
+licenses(["notice"])
+
+compile_pip_requirements(
+    # Defines targets which use pip-compile to keep the Python locked
+    # requirements up-to-date:
+    #
+    # :python_requirements.update  bazel run this target to update
+    #                              ./python_requirements.txt by recursively following
+    #                              and locking the dependencies seeded by
+    #                              ./python_requirements.in
+    #
+    # :python_requirements_test    bazel test target which fails if
+    #                              ./python_requirements.txt does not match
+    #                              that generated from ./python_requirements.in
+    name = "python_requirements",
+    extra_args = [
+        "--allow-unsafe",
+        # ^ lets pip-compile include setuptools, recommended by
+        # `pip-compile -h` as future default behavior
+    ],
+    requirements_in = "python_requirements.in",
+    requirements_txt = "python_requirements.txt",
+    tags = [
+        "manual",
+        # ^ exclude .update and _test targets from wildcards in,
+        # e.g., `bazel test ...`
+    ],
+)

oss_scripts/requirements/python_requirements.in

@@ -0,0 +1,29 @@
+# Specify the Python packages available as dependencies to targets in Bazel.
+#
+# When modifying this list, always run
+# the //oss_scripts/requirements:python_requirements.update target:
+#
+#     bazel run //oss_scripts/requirements:python_requirements.update
+#
+# to compile (using pip-compile) this list of direct dependencies into a pinned
+# requirements file---a complete list of direct and transitive dependencies,
+# pinned by version and cryptographic hash. The pinned requirements file is
+# used in @rules_python's pip_parse() in the WORKSPACE file to create the
+# external repositories available as dependencies to py_binary() and
+# py_library() targets.
+#
+# To upgrade dependencies to their latest version, run the update target with
+# the option --upgrade:
+#
+#   bazel run //oss_scripts/requirements:python_requirements.update -- --upgrade
+#
+# Without the --upgrade option, the underlying pip-compile only adds or removes
+# dependencies without upgrading them to the latest versions available in PyPI.
+#
+# Both this input file and the pinned requirements file should be committed to
+# git. Avoid committing changes that break other developers by using an
+# environment that meets the project's recommendations. Dependency resolution
+# is sensitive to the Python environment (interpreter version, etc.) in which
+# it is run.
+
+tf-nightly