module(
    name = "zml",
)

bazel_dep(name = "abseil-cpp", version = "20250814.0")
bazel_dep(name = "aspect_bazel_lib", version = "2.20.0")
bazel_dep(name = "aspect_rules_py", version = "1.6.3")
bazel_dep(name = "bazel_skylib", version = "1.8.1")
bazel_dep(name = "libxev", version = "0.0.0-20251014.0-9f785d2")
bazel_dep(name = "patchelf", version = "0.18.0")
bazel_dep(name = "pcre2", version = "10.45")
bazel_dep(name = "platforms", version = "1.0.0")
bazel_dep(name = "protobuf", version = "32.0", repo_name = "com_google_protobuf")
bazel_dep(name = "rules_cc", version = "0.2.0")
bazel_dep(name = "rules_distroless", version = "0.5.1")
bazel_dep(name = "rules_license", version = "1.0.0")
bazel_dep(name = "rules_oci", version = "2.2.6")
bazel_dep(name = "rules_proto", version = "7.1.0")
bazel_dep(name = "rules_python", version = "1.5.3")
bazel_dep(name = "rules_rust", version = "0.63.0")
bazel_dep(name = "rules_zig", version = "0.12.1")
bazel_dep(name = "toolchains_llvm_bootstrapped", version = "0.2.4")
bazel_dep(name = "with_cfg.bzl", version = "0.11.0")

bazel_dep(name = "buildifier_prebuilt", version = "8.2.0.2", dev_dependency = True)

# Remove this when rules_zig@0.12.2 is released.
archive_override(
    module_name = "rules_zig",
    urls = ["https://github.com/aherrmann/rules_zig/archive/cbf03e45c9ffa7d23a50790929649805a08afbf9.tar.gz"],
    strip_prefix = "rules_zig-cbf03e45c9ffa7d23a50790929649805a08afbf9",
    integrity = "sha256-lu6gNjulEIqG68I7DRyJzWInDVk3lL+mCj2J14vYKwc=",
)

zig = use_extension("@rules_zig//zig:extensions.bzl", "zig")
zig.index(file = "//bazel:zig_index.json")
zig.toolchain(zig_version = "0.15.1")
zig.mirrors(urls = [
    "https://mirror.zml.ai/zig",
    "https://ziglang.org/builds/",
])
use_repo(zig, "zig_toolchains")

register_toolchains("@rules_zig//zig/target:all")

register_toolchains("@zig_toolchains//:all")

register_toolchains(
    "@toolchains_llvm_bootstrapped//toolchain:linux_amd64",
    "@toolchains_llvm_bootstrapped//toolchain:linux_aarch64",
)

uv = use_extension("@rules_python//python/uv:uv.bzl", "uv", dev_dependency = True)
uv.configure(version = "0.6.5")

python = use_extension("@rules_python//python/extensions:python.bzl", "python")
python.toolchain(python_version = "3.11")
use_repo(python, "python_versions")

pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
pip.parse(
    download_only = True,
    extra_pip_args = [
        "--abi=cp311",
        "--implementation=cp",
        "--python-version=311",
        "--platform=linux_x86_64",
        "--platform=manylinux2014_x86_64",
        "--platform=manylinux_2_27_x86_64",
        "--platform=manylinux_2_28_x86_64",
        "--platform=manylinux_2_29_x86_64",
    ],
    hub_name = "neuron_py_deps",
    python_version = "3.11",
    requirements_lock = "//runtimes/neuron:requirements.lock.txt",
)
use_repo(pip, "neuron_py_deps")

# For downloading models
pip.parse(
    download_only = True,
    hub_name = "huggingface_hub",
    python_version = "3.11",
    requirements_lock = "//tools/hf:requirements.lock.txt",
)
use_repo(pip, "huggingface_hub")

cpu = use_extension("//runtimes/cpu:cpu.bzl", "cpu_pjrt_plugin")
use_repo(cpu, "libpjrt_cpu_darwin_amd64", "libpjrt_cpu_darwin_arm64", "libpjrt_cpu_linux_amd64")

cuda = use_extension("//runtimes/cuda:cuda.bzl", "cuda_packages")
use_repo(cuda, "libpjrt_cuda")

rocm = use_extension("//runtimes/rocm:rocm.bzl", "rocm_packages")
use_repo(rocm, "hipblaslt", "libpjrt_rocm", "rocblas")

tpu = use_extension("//runtimes/tpu:tpu.bzl", "tpu_packages")
use_repo(tpu, "libpjrt_tpu")

neuron = use_extension("//runtimes/neuron:neuron.bzl", "neuron_packages")
use_repo(neuron, "aws-neuronx-collectives", "aws-neuronx-runtime-lib", "libgomp1", "libpjrt_neuron", "zlib1g")

zls = use_extension("//third_party/zls:zls_toolchain.bzl", "zls_toolchains")
use_repo(zls, "zls_toolchains")

register_toolchains("@zls_toolchains//:all")

non_module_deps = use_extension("//:third_party/non_module_deps.bzl", "non_module_deps")
use_repo(non_module_deps, "com_github_hejsil_clap", "com_google_sentencepiece", "mnist", "org_swig_swig", "xla")

xla = use_extension("//third_party/xla:xla.bzl", "xla")
use_repo(
    xla,
    "com_github_grpc_grpc",
    "llvm-raw",
    "local_config_cuda",
    "local_config_remote_execution",
    "local_config_rocm",
    "local_config_tensorrt",
    "python_version_repo",
    "rules_ml_toolchain",
    "stablehlo",
    "triton",
    "tsl",
)

llvm = use_extension("//third_party/xla:llvm.bzl", "llvm")
llvm.configure()
use_repo(llvm, "llvm-project")

rust = use_extension("@rules_rust//rust:extensions.bzl", "rust")
rust.toolchain(
    edition = "2021",
    extra_target_triples = [
        "aarch64-apple-darwin",
        "x86_64-apple-darwin",
        "aarch64-unknown-linux-gnu",
        "x86_64-unknown-linux-gnu",
    ],
    versions = ["1.89.0"],
)
use_repo(rust, "rust_toolchains")

register_toolchains("@rust_toolchains//:all")

crate = use_extension("@rules_rust//crate_universe:extensions.bzl", "crate")
crate.from_cargo(
    name = "crates",
    cargo_lockfile = "//zml/tokenizer/hftokenizers:Cargo.lock",
    manifests = ["//zml/tokenizer/hftokenizers:Cargo.toml"],
    supported_platform_triples = [
        "aarch64-apple-darwin",
        "x86_64-apple-darwin",
        "aarch64-unknown-linux-gnu",
        "x86_64-unknown-linux-gnu",
    ],
)
use_repo(crate, "crates")

apt = use_extension("@rules_distroless//apt:extensions.bzl", "apt")
apt.install(
    name = "apt_cuda",
    lock = "//runtimes/cuda:packages.lock.json",
    manifest = "//runtimes/cuda:packages.yaml",
)
use_repo(apt, "apt_cuda")

apt.install(
    name = "apt_rocm",
    lock = "//runtimes/rocm:packages.lock.json",
    manifest = "//runtimes/rocm:packages.yaml",
)
use_repo(apt, "apt_rocm")

apt.install(
    name = "apt_neuron",
    lock = "//runtimes/neuron:packages.lock.json",
    manifest = "//runtimes/neuron:packages.yaml",
)
use_repo(apt, "apt_neuron")

oci = use_extension("@rules_oci//oci:extensions.bzl", "oci")
oci.pull(
    name = "distroless_cc_debian12",
    digest = "sha256:eccec5274132c1be0ce5d2c8e6fe41033e64af5e987ccee9007826e4c012069d",
    image = "gcr.io/distroless/cc-debian12",
    platforms = ["linux/amd64"],
)
use_repo(oci, "distroless_cc_debian12", "distroless_cc_debian12_linux_amd64")

huggingface = use_extension("@zml//bazel:huggingface.bzl", "huggingface")

################################################################################
# ModernBERT
huggingface.model(
    name = "ModernBERT-base",
    build_file_content = """exports_files(glob(["**"]), visibility = ["//visibility:public"])""",
    commit = "94032bb66234a691cf6248265170006a7ced4970",
    includes = [
        "model.safetensors",
        "tokenizer.json",
    ],
    model = "answerdotai/ModernBERT-base",
)
huggingface.model(
    name = "ModernBERT-large",
    build_file_content = """exports_files(glob(["**"]), visibility = ["//visibility:public"])""",
    commit = "4bbcbf40bed02ce487125bcb3c897ea9bdc88340",
    includes = [
        "model.safetensors",
        "tokenizer.json",
    ],
    model = "answerdotai/ModernBERT-large",
)
use_repo(huggingface, "ModernBERT-base", "ModernBERT-large")
################################################################################