repo_description: Pretrained model hub for Keras 3.
repo_topics: cv deep-learning jax keras llm machine-learning natural-language-processing nlp python pytorch tensorflow
code: """Script to create (and optionally install) a `.whl` archive for KerasHub.
  
  By default this will also create a shim package for `keras-nlp` (the old
  package name) that provides a backwards compatible namespace.
  
  Usage:
  
  1. Create `.whl` files in `dist/` and `keras_nlp/dist/`:
  
  ```
  python3 pip_build.py
  ```
  
  2. Also install the new packages immediately after:
  
  ```
  python3 pip_build.py --install
  ```
  
  3. Only build keras-hub:
  
  ```
  python3 pip_build.py --install --skip_keras_nlp
  ```
  """
  
  import argparse
  import datetime
  import os
  import pathlib
  import re
  import shutil
  import subprocess
  
  from keras_hub.src.version import __version__
  
  
  def ignore_files(_, filenames):
      return [f for f in filenames if "_test" in f]
  
  
  hub_package = "keras_hub"
  nlp_package = "keras_nlp"
  build_directory = "tmp_build_dir"
  dist_directory = "dist"
  to_copy = ["pyproject.toml", "README.md"]
  
  
  def update_nightly_version(build_path, version):
      """Rewrite library version with the nightly package version."""
      date = datetime.datetime.now()
      new_version = re.sub(
          r"([0-9]+\.[0-9]+\.[0-9]+).*",  # Match version without suffix.
          r"\1.dev" + date.strftime("%Y%m%d%H%M"),  # Add dev{date} suffix.
          version,
      )
  
      version_file = build_path / hub_package / "src" / "version.py"
      version_contents = version_file.read_text()
      version_contents = re.sub(
          "\n__version__ = .*\n",
          f'\n__version__ = "{new_version}"\n',
          version_contents,
      )
      version_file.write_text(version_contents)
      return new_version
  
  
  def update_nightly_name(build_path, pkg_name):
      """Rewrite library name with the nightly package name."""
      new_pkg_name = f"{pkg_name}-nightly"
      pyproj_file = build_path / "pyproject.toml"
      pyproj_contents = pyproj_file.read_text()
      pyproj_contents = pyproj_contents.replace(
          f'name = "{pkg_name}"', f'name = "{new_pkg_name}"'
      )
      pyproj_file.write_text(pyproj_contents)
      return new_pkg_name
  
  
  def pin_keras_nlp_version(build_path, pkg_name, version):
      """Pin keras-nlp version and dependency to the keras-hub version."""
      pyproj_file = build_path / "pyproject.toml"
      pyproj_contents = pyproj_file.read_text()
      pyproj_contents = re.sub(
          "version = .*\n",
          f'version = "{version}"\n',
          pyproj_contents,
      )
  
      pyproj_contents = re.sub(
          "dependencies = .*\n",
          f'dependencies = ["{pkg_name}=={version}"]\n',
          pyproj_contents,
      )
      pyproj_file.write_text(pyproj_contents)
  
  
  def copy_source_to_build_directory(src, dst, package):
      # Copy sources (`keras_hub/` directory and setup files) to build
      # directory
      shutil.copytree(src / package, dst / package, ignore=ignore_files)
      for fname in to_copy:
          shutil.copy(src / fname, dst / fname)
  
  
  def build_wheel(build_path, dist_path, name, 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)
  
      # Check for the expected .whl file path
      name = name.replace("-", "_")
      whl_path = dist_path / f"{name}-{version}-py3-none-any.whl"
      if not os.path.exists(whl_path):
          raise ValueError(f"Could not find whl {whl_path}")
      print(f"Build successful. Wheel file available at {whl_path}")
      return whl_path
  
  
  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:
          whls = []
          build_path = root_path / build_directory
          dist_path = root_path / dist_directory
          os.mkdir(build_path)
          copy_source_to_build_directory(root_path, build_path, hub_package)
  
          version = __version__
          pkg_name = hub_package.replace("_", "-")
          if is_nightly:
              version = update_nightly_version(build_path, version)
              pkg_name = update_nightly_name(build_path, pkg_name)
              assert "dev" in version, "Version should contain dev"
              assert "nightly" in pkg_name, "Name should contain nightly"
  
          whl = build_wheel(build_path, dist_path, pkg_name, version)
          whls.append(whl)
  
          if keras_nlp:
              build_path = root_path / build_directory / nlp_package
              dist_path = root_path / nlp_package / dist_directory
              copy_source_to_build_directory(
                  root_path / nlp_package, build_path, nlp_package
              )
  
              pin_keras_nlp_version(build_path, pkg_name, version)
              nlp_pkg_name = nlp_package.replace("_", "-")
              if is_nightly:
                  nlp_pkg_name = update_nightly_name(build_path, nlp_pkg_name)
                  assert "dev" in version, "Version should contain dev"
                  assert "nightly" in nlp_pkg_name, "Name should contain nightly"
  
              whl = build_wheel(build_path, dist_path, nlp_pkg_name, version)
              whls.append(whl)
  
          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(whls):
      for path in whls:
          print(f"Installing wheel file: {path}")
          subprocess.run(
              [
                  "pip3",
                  "install",
                  path,
                  "--force-reinstall",
                  "--no-dependencies",
              ],
              check=True,
          )
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser()
      parser.add_argument(
          "--install",
          action="store_true",
          help="Whether to install the generated wheel file.",
      )
      parser.add_argument(
          "--nightly",
          action="store_true",
          help="Whether to generate nightly wheel file.",
      )
      parser.add_argument(
          "--skip_keras_nlp",
          action="store_true",
          help="Whether to build the keras-nlp shim package.",
      )
      args = parser.parse_args()
      root_path = pathlib.Path(__file__).parent.resolve()
      whls = build(root_path, args.nightly, not args.skip_keras_nlp)
      if whls and args.install:
          install_whl(whls)
  