examples: add modernBERT (zml#192)

Modern bert implementation from Victor: @vctrmn

Tested with https://huggingface.co/answerdotai/ModernBERT-base
I started from his PR zml#149 and cleaned up a bit the code to make it more
what I consider "idiomatic ZML" even if it's still a bit in flux since
we are so young.

---------

Co-authored-by: vctrmn <[email protected]>
Co-authored-by: Rene Schallner <[email protected]>

Author: 3 people

examples/MODULE.bazel

@@ -139,6 +139,53 @@ http_file(
     url = "https://github.com/karpathy/llama2.c/raw/c02865df300f3bd9e567ce061000dc23bf785a17/tokenizer.bin",
 )
 
+# ModernBERT
+huggingface.model(
+    name = "ModernBERT-base",
+    build_file_content = """\
+package(default_visibility = ["//visibility:public"])
+filegroup(
+    name = "model",
+    srcs = ["model.safetensors"],
+)
+
+filegroup(
+    name = "tokenizer",
+    srcs = ["tokenizer.json"],
+)
+""",
+    commit = "94032bb66234a691cf6248265170006a7ced4970",
+    includes = [
+        "model.safetensors",
+        "tokenizer.json",
+    ],
+    model = "answerdotai/ModernBERT-base",
+)
+use_repo(huggingface, "ModernBERT-base")
+
+huggingface.model(
+    name = "ModernBERT-large",
+    build_file_content = """\
+package(default_visibility = ["//visibility:public"])
+filegroup(
+    name = "model",
+    srcs = ["model.safetensors"],
+)
+
+filegroup(
+    name = "tokenizer",
+    srcs = ["tokenizer.json"],
+)
+""",
+    commit = "4bbcbf40bed02ce487125bcb3c897ea9bdc88340",
+    includes = [
+        "model.safetensors",
+        "tokenizer.json",
+    ],
+    model = "answerdotai/ModernBERT-large",
+)
+use_repo(huggingface, "ModernBERT-large")
+
 bazel_dep(name = "rules_rust", version = "0.57.1")
 rust = use_extension("@rules_rust//rust:extensions.bzl", "rust")
 rust.toolchain(

examples/modernbert/BUILD.bazel

@@ -0,0 +1,63 @@
+load("@zml//bazel:zig.bzl", "zig_cc_binary")
+
+zig_cc_binary(
+    name = "modernbert",
+    srcs = ["modernbert.zig"],
+    main = "main.zig",
+    deps = [
+        "@com_github_hejsil_clap//:clap",
+        "@zml//async",
+        "@zml//stdx",
+        "@zml//zml",
+    ],
+)
+
+cc_binary(
+    name = "ModernBERT-base",
+    args = [
+        "--model=$(location @ModernBERT-base//:model.safetensors)",
+        "--tokenizer=$(location @ModernBERT-base//:tokenizer)",
+        "--num-attention-heads=12",
+        "--tie-word-embeddings=true",
+    ],
+    data = [
+        "@ModernBERT-base//:model.safetensors",
+        "@ModernBERT-base//:tokenizer",
+    ],
+    deps = [":modernbert_lib"],
+)
+
+cc_binary(
+    name = "ModernBERT-large",
+    args = [
+        "--model=$(location @ModernBERT-large//:model.safetensors)",
+        "--tokenizer=$(location @ModernBERT-large//:tokenizer)",
+        "--num-attention-heads=16",
+        "--tie-word-embeddings=true",
+    ],
+    data = [
+        "@ModernBERT-large//:model.safetensors",
+        "@ModernBERT-large//:tokenizer",
+    ],
+    deps = [":modernbert_lib"],
+)
+
+zig_cc_binary(
+    name = "test-implementation",
+    srcs = ["modernbert.zig"],
+    args = [
+        "--model=$(location @ModernBERT-base//:model.safetensors)",
+    ],
+    data = [
+        "@ModernBERT-base//:model.safetensors",
+    ],
+    main = "test.zig",
+    tags = [
+        "no_ci",
+    ],
+    deps = [
+        "@com_github_hejsil_clap//:clap",
+        "@zml//async",
+        "@zml//zml",
+    ],
+)

examples/modernbert/activations/activations.py

@@ -0,0 +1,55 @@
+import logging
+import torch
+from transformers import pipeline
+from tools.zml_utils import ActivationCollector
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s"
+)
+log = logging.getLogger(__name__)
+
+MODEL_NAME: str = "answerdotai/ModernBERT-base"
+
+
+def main() -> None:
+    try:
+        log.info("Start running main()")
+
+        log.info(f"CPU capability : `{torch.backends.cpu.get_cpu_capability()}`")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        log.info(f"Loading model : `{MODEL_NAME}`")
+
+        fill_mask_pipeline = pipeline(
+            "fill-mask",
+            model=MODEL_NAME,
+            device_map=device,
+        )
+        model, tokenizer = fill_mask_pipeline.model, fill_mask_pipeline.tokenizer
+        log.info(
+            f"Model loaded successfully {model.config.architectures} - `{model.config.torch_dtype}` - {tokenizer.model_max_length} max tokens"  # noqa: E501
+        )
+
+        # Wrap the pipeline, and extract activations.
+        # Activations files can be huge for big models,
+        # so let's stop collecting after 1000 layers.
+        zml_pipeline = ActivationCollector(
+            fill_mask_pipeline, max_layers=1000, stop_after_first_step=True
+        )
+
+        input_text = "Paris is the [MASK] of France."
+        outputs, activations = zml_pipeline(input_text)
+        log.info(f"ouputs : {outputs}")
+
+        filename = MODEL_NAME.split("/")[-1] + ".activations.pt"
+        torch.save(activations, filename)
+        log.info(f"Saved {len(activations)} activations to {filename}")
+
+        log.info("End running main()")
+    except Exception as exception:
+        log.error(exception)
+        raise
+
+
+if __name__ == "__main__":
+    main()

examples/modernbert/activations/requirements.in

@@ -0,0 +1,4 @@
+torch
+transformers==4.48.1
+accelerate
+numpy==1.26.4