Extend MLX backend coverage to ModernBERT and Longformer architectures

docs/export-matrix.md

@@ -34,8 +34,8 @@ when a new family is added to the code.
 | RoBERTa | `roberta` | ✅ (via bert) | ✅ | ⚠️ INT4-if-recall-holds | ✅ | ✅ | ✅ | ✅ | ✅ | ⚠️ backbone only |
 | XLM-RoBERTa | `xlm-roberta`, `xlm_roberta` | ✅ (via bert) | ✅ | ⚠️ INT4-if-recall-holds | ✅ | ✅ | ✅ | ✅ | ✅ | ⚠️ backbone only |
 | DeBERTa-v2 | `deberta`, `deberta-v2` | ✅ | ✅ | ⚠️ INT4-if-recall-holds | ✅ | ✅ | ✅ | ✅ | ✅ | ⚠️ backbone only |
-| ModernBERT | — | ❌ family not in MLX allowlist | ❌ | ❌ | ❌ family not in CoreML allowlist | ❌ | ⚠️ upstream ONNX support required | ⚠️ same as ONNX (default) | ⚠️ same as ONNX (default) | ❌ |
-| Longformer | — | ❌ family not in MLX allowlist | ❌ | ❌ | ❌ family not in CoreML allowlist | ❌ | ⚠️ upstream ONNX support required | ⚠️ same as ONNX (default) | ⚠️ same as ONNX (default) | ❌ |
+| ModernBERT | `modernbert`, `modern-bert` | ✅ | ✅ | ⚠️ INT4-if-recall-holds | ❌ family not in CoreML allowlist | ❌ | ⚠️ upstream ONNX support required | ⚠️ same as ONNX (default) | ⚠️ same as ONNX (default) | ❌ |
+| Longformer | `longformer` | ✅ | ✅ | ⚠️ INT4-if-recall-holds | ❌ family not in CoreML allowlist | ❌ | ⚠️ upstream ONNX support required | ⚠️ same as ONNX (default) | ⚠️ same as ONNX (default) | ❌ |
 | OpenAI Privacy Filter | `openai-privacy-filter`, `privacy-filter`, `privacy-filter-nemotron`, `nemotron-privacy-filter`, `privacy-filter-multilingual`, `multilingual-privacy-filter` | ✅ | ✅ | ⚠️ INT4-if-recall-holds | ❌ family not in CoreML allowlist | ❌ | ❌ no ONNX converter path today | ❌ | ❌ | ❌ |
 | GLiNER (span) | `gliner-uni-encoder-span` | ✅ (MLX custom family) | ✅ | ⚠️ INT4-if-recall-holds | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | GLiClass | `gliclass-uni-encoder` | ✅ (MLX custom family) | ✅ | ⚠️ INT4-if-recall-holds | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
@@ -59,10 +59,11 @@ truth is the backend code, but the essentials are:
   Any other architecture fails before tracing with an error that lists the
   supported families. See [CoreML Packaging](coreml-export.md).
 - **MLX enforces a family allowlist.** The MLX dispatcher resolves architectures
-  through `_SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES` (BERT-family and
-  DeBERTa-v2) plus the MLX-custom set (`gliner-uni-encoder-span`,
-  `gliclass-uni-encoder`, `gliner-uni-encoder-token-relex`) and the privacy
-  filter family. Architectures outside the allowlist raise a clear
+  through `_SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES` (BERT-family,
+  DeBERTa-v2, ModernBERT, and Longformer) plus the MLX-custom set
+  (`gliner-uni-encoder-span`, `gliclass-uni-encoder`,
+  `gliner-uni-encoder-token-relex`) and the privacy filter family.
+  Architectures outside the allowlist raise a clear
   "unsupported family" error rather than falling through silently. See
   [MLX Backend](mlx-backend.md).
 - **INT4 is certified only when recall holds.** `openmed.mlx.convert`

openmed/mlx/convert.py

@@ -110,6 +110,31 @@
     ("classifier.", "classifier."),
 ]
 
+_MODERNBERT_KEY_REPLACEMENTS: list[Tuple[str, str]] = [
+    ("model.layers.", "model.encoder.layers."),
+    ("model.embeddings.tok_embeddings.", "model.embeddings.word_embeddings."),
+    (".attn.Wqkv.", ".attention.qkv_proj."),
+    (".attn.Wo.", ".attention.out_proj."),
+    (".mlp.Wi.", ".mlp.wi_proj."),
+    (".mlp.Wo.", ".mlp.wo_proj."),
+]
+
+_LONGFORMER_KEY_REPLACEMENTS: list[Tuple[str, str]] = [
+    ("longformer.encoder.layer.", "longformer.encoder.layers."),
+    ("longformer.embeddings.LayerNorm.", "longformer.embeddings.norm."),
+    (".attention.self.query_global.", ".attention.query_global_proj."),
+    (".attention.self.key_global.", ".attention.key_global_proj."),
+    (".attention.self.value_global.", ".attention.value_global_proj."),
+    (".attention.self.query.", ".attention.query_proj."),
+    (".attention.self.key.", ".attention.key_proj."),
+    (".attention.self.value.", ".attention.value_proj."),
+    (".attention.output.dense.", ".attention.out_proj."),
+    (".attention.output.LayerNorm.", ".ln1."),
+    (".intermediate.dense.", ".linear1."),
+    (".output.dense.", ".linear2."),
+    (".output.LayerNorm.", ".ln2."),
+]
+
 _OPF_MODEL_TYPES = {
     "openai-privacy-filter",
     "privacy-filter",
@@ -317,6 +342,15 @@ def _infer_source_model_type(key: str, model_type: str | None) -> str:
         return "roberta"
     if key.startswith("electra."):
         return "electra"
+    if key.startswith("model."):
+        if (
+            ".layers." in key
+            or key.startswith("model.embeddings.")
+            or key.startswith("model.final_norm.")
+        ):
+            return "modernbert"
+    if key.startswith("longformer."):
+        return "longformer"
     return "bert"
 
 
@@ -325,7 +359,11 @@ def remap_key(key: str, model_type: str | None = None) -> str:
     source_model_type = _infer_source_model_type(key, model_type)
     resolved_model_type = resolve_model_type(source_model_type)
 
-    if resolved_model_type == "deberta-v2":
+    if source_model_type in {"modernbert", "modern-bert"}:
+        replacements = _MODERNBERT_KEY_REPLACEMENTS
+    elif source_model_type == "longformer":
+        replacements = _LONGFORMER_KEY_REPLACEMENTS
+    elif resolved_model_type == "deberta-v2":
         replacements = _DEBERTA_V2_KEY_REPLACEMENTS
     elif source_model_type == "distilbert":
         replacements = _DISTILBERT_KEY_REPLACEMENTS
@@ -427,10 +465,14 @@ def save_mlx_model(
         logger.info("Quantizing to %d bits ...", quantize_bits)
         model = build_model(config)
         model.load_weights(list(mlx_weights.items()))
-        quantize_kwargs = {"bits": quantize_bits}
-        if quantize_group_size is not None:
-            quantize_kwargs["group_size"] = quantize_group_size
-        nn.quantize(model, **quantize_kwargs)
+        if quantize_group_size is None:
+            nn.quantize(model, bits=quantize_bits)
+        else:
+            nn.quantize(
+                model,
+                group_size=quantize_group_size,
+                bits=quantize_bits,
+            )
         mlx_weights = dict(tree_flatten(model.parameters()))
         config_to_save["_mlx_quantization"] = {
             "bits": quantize_bits,

openmed/mlx/models/__init__.py

@@ -12,6 +12,9 @@
     "bert": "bert",
     "distilbert": "bert",
     "electra": "bert",
+    "modernbert": "modernbert",
+    "modern-bert": "modernbert",
+    "longformer": "longformer",
     "openai-privacy-filter": "openai-privacy-filter",
     "privacy-filter": "openai-privacy-filter",
     # ``OpenMed/privacy-filter-nemotron*`` is the OpenAI Privacy Filter
@@ -83,7 +86,11 @@ def resolve_artifact_family(
         family = normalize_model_type(manifest.get("family"))
         if family in _CUSTOM_FAMILIES:
             return family
-        resolved_family = _SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(family)
+        resolved_family = (
+            _SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(family)
+            if family is not None
+            else None
+        )
         if resolved_family is not None:
             return resolved_family
 
@@ -92,7 +99,11 @@ def resolve_artifact_family(
         family = normalize_model_type(config.get("_mlx_family"))
         if family in _CUSTOM_FAMILIES:
             return family
-        resolved_family = _SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(family)
+        resolved_family = (
+            _SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(family)
+            if family is not None
+            else None
+        )
         if resolved_family is not None:
             return resolved_family
 
@@ -107,7 +118,11 @@ def resolve_artifact_family(
         model_type = config
 
     normalized = normalize_model_type(model_type)
-    resolved = _SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(normalized)
+    resolved = (
+        _SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(normalized)
+        if normalized is not None
+        else None
+    )
     if resolved is None:
         supported = ", ".join(
             sorted(_CUSTOM_FAMILIES | set(_SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES))
@@ -234,6 +249,81 @@ def normalize_model_config(
             "bidirectional_right_context",
             normalized.get("bidirectional_right_context", default_context),
         )
+    elif family == "modernbert":
+        normalized.setdefault("type_vocab_size", 0)
+        normalized.setdefault("_mlx_position_offset", 0)
+        normalized.setdefault(
+            "norm_eps",
+            normalized.get("norm_eps", normalized.get("layer_norm_eps", 1e-5)),
+        )
+        normalized.setdefault("norm_bias", normalized.get("norm_bias", False))
+        normalized.setdefault(
+            "hidden_activation",
+            normalized.get("hidden_activation", normalized.get("hidden_act", "gelu")),
+        )
+        normalized.setdefault(
+            "embedding_dropout",
+            normalized.get("embedding_dropout", normalized.get("dropout", 0.0)),
+        )
+        normalized.setdefault(
+            "attention_dropout",
+            normalized.get("attention_dropout", 0.0),
+        )
+        normalized.setdefault("mlp_dropout", normalized.get("mlp_dropout", 0.0))
+        normalized.setdefault(
+            "classifier_dropout",
+            normalized.get("classifier_dropout", 0.0),
+        )
+        normalized.setdefault("attention_bias", normalized.get("attention_bias", False))
+        normalized.setdefault("mlp_bias", normalized.get("mlp_bias", False))
+        normalized.setdefault(
+            "classifier_bias", normalized.get("classifier_bias", False)
+        )
+        normalized.setdefault(
+            "classifier_activation",
+            normalized.get("classifier_activation", "gelu"),
+        )
+        normalized.setdefault("local_attention", normalized.get("local_attention", 128))
+        normalized.setdefault("sliding_window", int(normalized["local_attention"]) // 2)
+        if not normalized.get("layer_types"):
+            every_n_layers = int(normalized.get("global_attn_every_n_layers", 3))
+            every_n_layers = max(1, every_n_layers)
+            normalized["layer_types"] = [
+                "full_attention" if index % every_n_layers == 0 else "sliding_attention"
+                for index in range(int(normalized["num_hidden_layers"]))
+            ]
+        rope_parameters = normalized.get("rope_parameters")
+        if not isinstance(rope_parameters, dict) or not any(
+            key in rope_parameters
+            for key in ("full_attention", "sliding_attention", "global", "local")
+        ):
+            shared_theta = (
+                rope_parameters.get("rope_theta")
+                if isinstance(rope_parameters, dict)
+                else None
+            )
+            global_rope_theta = normalized.get(
+                "global_rope_theta", shared_theta or 160000.0
+            )
+            if global_rope_theta is None:
+                global_rope_theta = shared_theta or 160000.0
+            local_rope_theta = normalized.get(
+                "local_rope_theta", shared_theta or 10000.0
+            )
+            if local_rope_theta is None:
+                local_rope_theta = shared_theta or 10000.0
+            normalized["rope_parameters"] = {
+                "full_attention": {"rope_theta": float(global_rope_theta)},
+                "sliding_attention": {"rope_theta": float(local_rope_theta)},
+            }
+    elif family == "longformer":
+        normalized.setdefault("type_vocab_size", 2)
+        normalized.setdefault("_mlx_position_offset", 0)
+        attention_window = normalized.get("attention_window", 512)
+        if isinstance(attention_window, (list, tuple)):
+            normalized["attention_window"] = [int(value) for value in attention_window]
+        else:
+            normalized["attention_window"] = int(attention_window)
     else:
         normalized.setdefault("type_vocab_size", normalized.get("type_vocab_size", 2))
         normalized.setdefault("_mlx_position_offset", 0)
@@ -262,6 +352,16 @@ def build_model(
 
         return DebertaV2ForTokenClassification(config)
 
+    if family == "modernbert":
+        from openmed.mlx.models.modernbert_tc import ModernBertForTokenClassification
+
+        return ModernBertForTokenClassification(config)
+
+    if family == "longformer":
+        from openmed.mlx.models.longformer_tc import LongformerForTokenClassification
+
+        return LongformerForTokenClassification(config)
+
     if family == "openai-privacy-filter":
         from openmed.mlx.models.privacy_filter import (
             OpenAIPrivacyFilterForTokenClassification,