Use tf-keras-nightly for tf-hub-nightly.

PiperOrigin-RevId: 604578087

Author: TensorFlow Hub Authors

tensorflow_hub/pip_package/build_pip_package.sh

@@ -58,15 +58,6 @@ function main() {
   cp "tensorflow_hub/LICENSE" "${TMPDIR}/LICENSE.txt"
   cp -R "${RUNFILES}/tensorflow_hub" "${TMPDIR}"
 
-  # Remove this once the library supports Keras 3.
-  # Install tf-keras-nightly instead if tf-nightly is installed.
-  if [[ $TENSORFLOW_VERSION == *tf-nightly ]]
-  then
-    sed -i "s|tf-keras|tf-keras-nightly|g" "${TMPDIR}/setup.py"
-  else
-    echo "Did not replace dependency on tf-keras"
-  fi
-
   pushd ${TMPDIR}
   rm -f MANIFEST
 

tensorflow_hub/pip_package/setup.py

@@ -15,10 +15,12 @@
 """Setup for pip package."""
 
 from datetime import datetime
+import os
+import sys
+
 from setuptools import find_packages
 from setuptools import setup
 
-import sys
 
 # Can't import the module during setup.py.
 # Use execfile to find __version__.
@@ -28,7 +30,6 @@
 REQUIRED_PACKAGES = [
     'numpy >= 1.12.0',
     'protobuf >= 3.19.6',  # No less than what ../WORKSPACE uses.
-    'tf-keras >= 2.14.1',
 ]
 
 project_name = 'tensorflow-hub'
@@ -41,9 +42,19 @@
 # If we're dealing with a nightly build we need to make sure that the
 # version changes for every release.
 version = __version__
+is_nightly = False
 if project_name == 'tf-hub-nightly':
+  is_nightly = True
   version += datetime.now().strftime('%Y%m%d%H%M')
 
+if os.environ.get('TENSORFLOW_VERSION', '').endswith('tf-nightly'):
+  is_nightly = True
+
+if is_nightly:
+  REQUIRED_PACKAGES += ['tf-keras-nightly']
+else:
+  REQUIRED_PACKAGES += ['tf-keras >= 2.14.1']
+
 setup(
     name=project_name,  # Automatic: tensorflow_hub, etc. Case insensitive.
     version=version.replace('-', ''),
@@ -83,4 +94,4 @@
         'tensorflow machine learning share module subgraph component hub '
         'embedding retraining transfer'
     ),
-)
+)