Fix pip_build.py to build shim package

keras_nlp/keras_nlp/__init__.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 import os
 
+# Add everything in /api/ to the module search path.
+import keras_hub
+
 # Import everything from /api/ into keras.
 from keras_hub.api import *  # noqa: F403
 from keras_hub.api import __version__  # Import * ignores names start with "_".
 
-# Add everything in /api/ to the module search path.
-import keras_hub
 __path__.extend(keras_hub.__path__)  # noqa: F405
 # Don't pollute namespace.
 del keras_hub

keras_nlp/setup.py

@@ -38,10 +38,7 @@ def get_version(rel_path):
 HERE = pathlib.Path(__file__).parent
 README = (HERE / "README.md").read_text()
 PARENT = HERE.parent
-if os.path.exists(PARENT / "keras_hub" / "version_utils.py"):
-    VERSION = get_version(PARENT / "keras_hub" / "version_utils.py")
-else:
-    VERSION = get_version(PARENT / "keras_hub" / "src" / "version_utils.py")
+VERSION = get_version(PARENT / "keras_hub" / "src" / "version_utils.py")
 
 setup(
     name="keras-nlp",

pip_build.py

@@ -30,13 +30,13 @@
 
 import argparse
 import datetime
-import glob
 import os
 import pathlib
 import re
 import shutil
 
-package = "keras_hub"
+hub_package = "keras_hub"
+nlp_package = "keras_nlp"
 build_directory = "tmp_build_dir"
 dist_directory = "dist"
 to_copy = ["setup.py", "setup.cfg", "README.md"]
@@ -46,90 +46,110 @@ def ignore_files(_, filenames):
     return [f for f in filenames if "_test" in f]
 
 
-def export_version_string(version, is_nightly=False):
+def update_version(build_path, package, version, is_nightly=False):
     """Export Version and Package Name."""
     if is_nightly:
         date = datetime.datetime.now()
         version += f".dev{date.strftime('%Y%m%d%H')}"
-        # Replaces `name="keras-hub"` in `setup.py` with `keras-hub-nightly`
-        with open("setup.py") as f:
-            setup_contents = f.read()
-        with open("setup.py", "w") as f:
-            setup_contents = setup_contents.replace(
-                'name="keras-hub"', 'name="keras-hub-nightly"'
-            )
-            f.write(setup_contents)
+        package_name = f"{package}-nightly"
+    else:
+        package_name = package
 
-    # Make sure to export the __version__ string
-    with open(os.path.join(package, "src", "version_utils.py")) as f:
-        init_contents = f.read()
-    with open(os.path.join(package, "src", "version_utils.py"), "w") as f:
-        init_contents = re.sub(
-            "\n__version__ = .*\n",
-            f'\n__version__ = "{version}"\n',
-            init_contents,
+    with open(build_path / "setup.py") as f:
+        setup_contents = f.read()
+    with open(build_path / "setup.py", "w") as f:
+        setup_contents = setup_contents.replace(
+            "name=", f'name="{package_name}", # '
+        )
+        setup_contents = setup_contents.replace(
+            "VERSION = ", f'VERSION = "{version}" # '
         )
-        f.write(init_contents)
+        f.write(setup_contents)
 
+    # Make sure to export the __version__ string
+    version_utils = build_path / package / "src" / "version_utils.py"
+    if os.path.exists(version_utils):
+        with open(version_utils) as f:
+            contents = f.read()
+        with open(version_utils, "w") as f:
+            contents = re.sub(
+                "\n__version__ = .*\n",
+                f'\n__version__ = "{version}"\n',
+                contents,
+            )
+            f.write(contents)
 
-def copy_source_to_build_directory(root_path):
+
+def copy_source_to_build_directory(root_path, package):
     # Copy sources (`keras_hub/` directory and setup files) to build
     # directory
-    os.chdir(root_path)
-    os.mkdir(build_directory)
     shutil.copytree(
-        package, os.path.join(build_directory, package), ignore=ignore_files
+        root_path / package,
+        root_path / build_directory / package,
+        ignore=ignore_files,
     )
     for fname in to_copy:
-        shutil.copy(fname, os.path.join(f"{build_directory}", fname))
-    os.chdir(build_directory)
+        shutil.copy(root_path / fname, root_path / build_directory / fname)
+
+
+def build_wheel(build_path, dist_path, __version__):
+    # Build the package
+    os.chdir(build_path)
+    os.system("python3 -m build")
+
+    # Save the dist files generated by the build process
+    if not os.path.exists(dist_path):
+        os.mkdir(dist_path)
+    for fpath in (build_path / dist_directory).glob("*.*"):
+        shutil.copy(fpath, dist_path)
+
+    # Find the .whl file path
+    for fname in os.listdir(dist_path):
+        if __version__ in fname and fname.endswith(".whl"):
+            whl_path = dist_path / fname
+            print(f"Build successful. Wheel file available at {whl_path}")
+            return whl_path
+    print("Build failed.")
+    return None
 
 
-def build(root_path, is_nightly=False):
+def build(root_path, is_nightly=False, keras_nlp=True):
     if os.path.exists(build_directory):
         raise ValueError(f"Directory already exists: {build_directory}")
 
     try:
-        copy_source_to_build_directory(root_path)
-        print(os.getcwd())
+        whls = []
+        build_path = root_path / build_directory
+        dist_path = root_path / dist_directory
+        os.mkdir(build_path)
 
         from keras_hub.src.version_utils import __version__  # noqa: E402
 
-        export_version_string(__version__, is_nightly)
-        return build_and_save_output(root_path, __version__)
-    finally:
-        # Clean up: remove the build directory (no longer needed)
-        shutil.rmtree(build_directory)
-
+        copy_source_to_build_directory(root_path, hub_package)
+        update_version(build_path, hub_package, __version__, is_nightly)
+        whl = build_wheel(build_path, dist_path, __version__)
+        whls.append(whl)
 
-def build_and_save_output(root_path, __version__):
-    # Build the package
-    os.system("python3 -m build")
+        if keras_nlp:
+            build_path = root_path / build_directory / nlp_package
+            dist_path = root_path / nlp_package / dist_directory
 
-    # Save the dist files generated by the build process
-    os.chdir(root_path)
-    if not os.path.exists(dist_directory):
-        os.mkdir(dist_directory)
-    for fpath in glob.glob(
-        os.path.join(build_directory, dist_directory, "*.*")
-    ):
-        shutil.copy(fpath, dist_directory)
+            copy_source_to_build_directory(root_path, nlp_package)
+            update_version(build_path, nlp_package, __version__, is_nightly)
+            whl = build_wheel(build_path, dist_path, __version__)
+            whls.append(whl)
 
-    # Find the .whl file path
-    whl_path = None
-    for fname in os.listdir(dist_directory):
-        if __version__ in fname and fname.endswith(".whl"):
-            whl_path = os.path.abspath(os.path.join(dist_directory, fname))
-    if whl_path:
-        print(f"Build successful. Wheel file available at {whl_path}")
-    else:
-        print("Build failed.")
-    return whl_path
+        return whls
+    finally:
+        # Clean up: remove the build directory (no longer needed)
+        os.chdir(root_path)
+        shutil.rmtree(root_path / build_directory)
 
 
-def install_whl(whl_fpath):
-    print(f"Installing wheel file: {whl_fpath}")
-    os.system(f"pip3 install {whl_fpath} --force-reinstall --no-dependencies")
+def install_whl(whls):
+    for path in whls:
+        print(f"Installing wheel file: {path}")
+        os.system(f"pip3 install {path} --force-reinstall --no-dependencies")
 
 
 if __name__ == "__main__":
@@ -138,14 +158,22 @@ def install_whl(whl_fpath):
         "--install",
         action="store_true",
         help="Whether to install the generated wheel file.",
+        default=False,
     )
     parser.add_argument(
         "--nightly",
         action="store_true",
         help="Whether to generate nightly wheel file.",
+        default=False,
+    )
+    parser.add_argument(
+        "--keras_nlp",
+        action="store_true",
+        help="Whether to build the keras-nlp shim package.",
+        default=True,
     )
     args = parser.parse_args()
     root_path = pathlib.Path(__file__).parent.resolve()
-    whl_path = build(root_path, args.nightly)
-    if whl_path and args.install:
-        install_whl(whl_path)
+    whls = build(root_path, args.nightly, args.keras_nlp)
+    if whls and args.install:
+        install_whl(whls)

setup.py

@@ -37,10 +37,7 @@ def get_version(rel_path):
 
 HERE = pathlib.Path(__file__).parent
 README = (HERE / "README.md").read_text()
-if os.path.exists("keras_hub/version_utils.py"):
-    VERSION = get_version("keras_hub/version_utils.py")
-else:
-    VERSION = get_version("keras_hub/src/version_utils.py")
+VERSION = get_version("keras_hub/src/version_utils.py")
 
 setup(
     name="keras-hub",