Integrate the UltraBERT architecture

@@ -931,6 +931,8 @@         title: T5v1.1       - local: model_doc/ul2         title: UL2+      - local: model_doc/ultrabert+        title: UltraBERT       - local: model_doc/umt5         title: UMT5       - local: model_doc/vaultgemma

@@ -0,0 +1,175 @@+<!--Copyright 2026 The HuggingFace Team. All rights reserved.++Licensed under the Apache License, Version 2.0 (the "License").+You may not use this file except in compliance with the License.+You may obtain a copy of the License at++http://www.apache.org/licenses/LICENSE-2.0++Unless required by applicable law or agreed to in writing, software+distributed under the License is distributed on an "AS IS" BASIS,+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+See the License for the specific language governing permissions and+limitations under the License.+-->+*This model was published in HF papers on 2026-10-10 and contributed to Hugging Face Transformers on 2026-10-14.*++# UltraBERT++## Overview++UltraBERT was introduced in [UltraBERT: Beyond the Limit of BERT Pre-Training](https://huggingface.co/papers/2610.15028) by Jungwoo Park, Taewhoo Lee, and Jaewoo Kang. It is a family of bidirectional text encoders trained with a combination of masking ratios, pruning, and distillation for classification and retrieval tasks.++The released encoders support sequences of up to 32,768 tokens. They use rotary position embeddings (RoPE), RMSNorm, SwiGLU feed-forward layers, and alternating local and global attention. Every third layer uses global attention. The other layers attend to the current token and up to 128 tokens on either side. The RoPE base is 10,000 for local layers and 2,560,000 for global layers.++| Checkpoint | Layers | Hidden size | Attention heads | Use |+|---|---|---|---|---|+| [UltraBERT-150M](https://huggingface.co/dmis-lab/UltraBERT-150M) | 24 | 768 | 6 | Masked language modeling and downstream fine-tuning |+| [UltraBERT-380M](https://huggingface.co/dmis-lab/UltraBERT-380M) | 24 | 1,024 | 8 | Masked language modeling and downstream fine-tuning |+| [UltraBERT-Embed-150M](https://huggingface.co/dmis-lab/UltraBERT-Embed-150M) | 24 | 768 | 6 | Text embeddings |++The masked language modeling checkpoints contain the prediction head trained with the low masking ratio. Its output weights are untied from the input embeddings. The tokenizer uses the SmolLM2 vocabulary of 49,152 tokens.++## Usage examples++### Masked language modeling++Use [`Pipeline`] or [`AutoModelForMaskedLM`] to predict a masked token. The mask token is `<|mask|>`.++<hfoptions id="usage">+<hfoption id="Pipeline">++```python+from transformers import pipeline++fill_mask = pipeline(+    task="fill-mask",+    model="dmis-lab/UltraBERT-150M",+    device_map="auto",+)+print(fill_mask("The capital of France is <|mask|>."))+# [+#     {'score': 0.9612365961074829, 'token': 7042, 'token_str': ' Paris',+#      'sequence': 'The capital of France is Paris.'},+#     {'score': 0.00930546224117279, 'token': 38034, 'token_str': ' Lyon',+#      'sequence': 'The capital of France is Lyon.'},+#     {'score': 0.008486305363476276, 'token': 44558, 'token_str': ' Nice',+#      'sequence': 'The capital of France is Nice.'},+#     {'score': 0.0034405170008540154, 'token': 37628, 'token_str': ' Versailles',+#      'sequence': 'The capital of France is Versailles.'},+#     {'score': 0.0014691735850647092, 'token': 4649, 'token_str': ' France',+#      'sequence': 'The capital of France is France.'}+# ]+```++</hfoption>+<hfoption id="AutoModel">++```python+import torch+from transformers import AutoModelForMaskedLM, AutoTokenizer++checkpoint = "dmis-lab/UltraBERT-150M"+tokenizer = AutoTokenizer.from_pretrained(checkpoint)+model = AutoModelForMaskedLM.from_pretrained(checkpoint, device_map="auto")++inputs = tokenizer("The capital of France is <|mask|>.", return_tensors="pt").to(model.device)+with torch.no_grad():+    logits = model(**inputs).logits++masked_logits = logits[inputs.input_ids == tokenizer.mask_token_id]+print(tokenizer.batch_decode(masked_logits.argmax(dim=-1)[:, None]))+# [' Paris']+```++</hfoption>+</hfoptions>++### Text embeddings++UltraBERT-Embed-150M uses the first token's hidden state as the text embedding. For retrieval, prepend the task instruction to queries and encode documents without an instruction. Use right padding to preserve the first token's position. See the [model card](https://huggingface.co/dmis-lab/UltraBERT-Embed-150M) for instructions for other tasks.++The examples below select the first token's hidden state and normalize the embeddings so their dot product gives cosine similarity.++<hfoptions id="embeddings">+<hfoption id="Pipeline">++```python+import torch+import torch.nn.functional as F+from transformers import pipeline++extractor = pipeline(+    task="feature-extraction",+    model="dmis-lab/UltraBERT-Embed-150M",+    device_map="auto",+)++texts = [+    "Instruct: Given a question, retrieve passages that answer the question\nQuery: What is the capital of France?",+    "Paris is the capital of France.",+    "Berlin is the capital of Germany.",+]+features = extractor(texts, return_tensors=True, truncation=True, tokenize_kwargs={"max_length": 32768})+embeddings = F.normalize(torch.cat([output[:, 0] for output in features]), p=2, dim=-1)++print(embeddings[:1] @ embeddings[1:].T)+# tensor([[0.6527, 0.3981]])+```++</hfoption>+<hfoption id="AutoModel">++```python+import torch+import torch.nn.functional as F+from transformers import AutoModel, AutoTokenizer++checkpoint = "dmis-lab/UltraBERT-Embed-150M"+tokenizer = AutoTokenizer.from_pretrained(checkpoint, padding_side="right")+model = AutoModel.from_pretrained(checkpoint, device_map="auto")++texts = [+    "Instruct: Given a question, retrieve passages that answer the question\nQuery: What is the capital of France?",+    "Paris is the capital of France.",+    "Berlin is the capital of Germany.",+]+inputs = tokenizer(texts, padding=True, truncation=True, max_length=32768, return_tensors="pt").to(model.device)+with torch.no_grad():+    embeddings = F.normalize(model(**inputs).last_hidden_state[:, 0], p=2, dim=-1)++print(embeddings[:1] @ embeddings[1:].T)+# tensor([[0.6527, 0.3981]], device='cuda:0')+```++</hfoption>+</hfoptions>++## UltrabertConfig++[[autodoc]] UltrabertConfig++## UltrabertModel++[[autodoc]] UltrabertModel+    - forward++## UltrabertForMaskedLM++[[autodoc]] UltrabertForMaskedLM+    - forward++## UltrabertForSequenceClassification++[[autodoc]] UltrabertForSequenceClassification+    - forward++## UltrabertForTokenClassification++[[autodoc]] UltrabertForTokenClassification+    - forward++## UltrabertForQuestionAnswering++[[autodoc]] UltrabertForQuestionAnswering+    - forward

@@ -1887,6 +1887,22 @@ def _build_checkpoint_conversion_mapping():             WeightRenaming(r"depth_head\.project", "decoder.project"),             WeightRenaming(r"depth_head\.depth_head\.", "decoder.head."),         ],+        "UltrabertModel": [+            PrefixChange(prefix_to_remove="ultrabert"),+        ],+        "UltrabertForMaskedLM": [+            WeightRenaming(r"^ultrabert\.", "model."),+            WeightRenaming(r"^cls\.", "head.decoder."),+        ],+        "UltrabertForSequenceClassification": [+            WeightRenaming(r"^ultrabert\.", "model."),+        ],+        "UltrabertForTokenClassification": [+            WeightRenaming(r"^ultrabert\.", "model."),+        ],+        "UltrabertForQuestionAnswering": [+            WeightRenaming(r"^ultrabert\.", "model."),+        ],     }      mapping["Tipsv2DptForNormalEstimation"] = [

@@ -492,6 +492,7 @@     from .trocr import *     from .tvp import *     from .udop import *+    from .ultrabert import *     from .umt5 import *     from .unispeech import *     from .unispeech_sat import *

@@ -710,6 +710,7 @@         ("trocr", "TrOCRConfig"),         ("tvp", "TvpConfig"),         ("udop", "UdopConfig"),+        ("ultrabert", "UltrabertConfig"),         ("umt5", "UMT5Config"),         ("unispeech", "UniSpeechConfig"),         ("unispeech-sat", "UniSpeechSatConfig"),

@@ -551,6 +551,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):         ("tipsv2_vision_model", "Tipsv2VisionModel"),         ("tvp", "TvpModel"),         ("udop", "UdopModel"),+        ("ultrabert", "UltrabertModel"),         ("umt5", "UMT5Model"),         ("unispeech", "UniSpeechModel"),         ("unispeech-sat", "UniSpeechSatModel"),@@ -1265,6 +1266,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):         ("roformer", "RoFormerForMaskedLM"),         ("squeezebert", "SqueezeBertForMaskedLM"),         ("tapas", "TapasForMaskedLM"),+        ("ultrabert", "UltrabertForMaskedLM"),         ("xlm", "XLMWithLMHeadModel"),         ("xlm-roberta", "XLMRobertaForMaskedLM"),         ("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"),@@ -1529,6 +1531,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):         ("t5gemma", "T5GemmaForSequenceClassification"),         ("t5gemma2", "T5Gemma2ForSequenceClassification"),         ("tapas", "TapasForSequenceClassification"),+        ("ultrabert", "UltrabertForSequenceClassification"),         ("umt5", "UMT5ForSequenceClassification"),         ("xlm", "XLMForSequenceClassification"),         ("xlm-roberta", "XLMRobertaForSequenceClassification"),@@ -1614,6 +1617,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):         ("splinter", "SplinterForQuestionAnswering"),         ("squeezebert", "SqueezeBertForQuestionAnswering"),         ("t5", "T5ForQuestionAnswering"),+        ("ultrabert", "UltrabertForQuestionAnswering"),         ("umt5", "UMT5ForQuestionAnswering"),         ("xlm", "XLMForQuestionAnsweringSimple"),         ("xlm-roberta", "XLMRobertaForQuestionAnswering"),@@ -1742,6 +1746,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):         ("t5", "T5ForTokenClassification"),         ("t5gemma", "T5GemmaForTokenClassification"),         ("t5gemma2", "T5Gemma2ForTokenClassification"),+        ("ultrabert", "UltrabertForTokenClassification"),         ("umt5", "UMT5ForTokenClassification"),         ("xlm", "XLMForTokenClassification"),         ("xlm-roberta", "XLMRobertaForTokenClassification"),

@@ -0,0 +1,33 @@+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+#           This file was automatically generated from the modules in `src/transformers/models/ultrabert`.+#           Do NOT edit this file manually as any edits will be overwritten by auto-generation of the file.+#                       A module is picked up once it exports a name through `__all__`.+#                   Regenerate the file with: `python utils/check_inits.py --fix_and_overwrite`+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+# Copyright 2026 the HuggingFace Team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License").+# You may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.+from typing import TYPE_CHECKING++from ...utils import _LazyModule+from ...utils.import_utils import define_import_structure+++if TYPE_CHECKING:+    from .configuration_ultrabert import *+    from .modeling_ultrabert import *+else:+    import sys++    _file = globals()["__file__"]+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

@@ -0,0 +1,128 @@+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+#           This file was automatically generated from src/transformers/models/ultrabert/modular_ultrabert.py.+#               Do NOT edit this file manually as any edits will be overwritten by the generation of+#             the file from the modular. If any change should be done, please apply the change to the+#                          modular_ultrabert.py file directly. One of our CI enforces this.+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+# Copyright 2026 the HuggingFace Team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License").+# You may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++from typing import Literal++from ...configuration_utils import PreTrainedConfig, strict+from ...modeling_rope_utils import RopeParameters+from ...utils import auto_docstring+++@auto_docstring(checkpoint="dmis-lab/UltraBERT-150M")+@strict+class UltrabertConfig(PreTrainedConfig):+    r"""+    sliding_window (`int`, *optional*, defaults to 128):+        Inclusive attention radius on either side of each token in sliding attention layers.+    prediction_head_size (`int`, *optional*, defaults to 768):+        Intermediate dimension of the two-layer masked language modeling prediction head.+    classifier_pooling (`str`, *optional*, defaults to `"cls"`):+        Pooling method for sequence classification. Either `"cls"` or `"mean"`.+    """++    model_type = "ultrabert"+    keys_to_ignore_at_inference = ["past_key_values"]+    # Default tensor parallel plan for base model `UltrabertModel`+    base_model_tp_plan = {+        "layers.*.self_attn.q_proj": "colwise",+        "layers.*.self_attn.k_proj": "colwise",+        "layers.*.self_attn.v_proj": "colwise",+        "layers.*.self_attn.o_proj": "rowwise",+        "layers.*.mlp.gate_proj": "colwise",+        "layers.*.mlp.up_proj": "colwise",+        "layers.*.mlp.down_proj": "rowwise",+    }+    base_model_pp_plan = {+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),+        "norm": (["hidden_states"], ["hidden_states"]),+    }++    vocab_size: int = 32000+    hidden_size: int = 4096+    intermediate_size: int = 11008+    num_hidden_layers: int = 32+    num_attention_heads: int = 32+    num_key_value_heads: int | None = None+    hidden_act: str = "silu"+    max_position_embeddings: int = 32768+    initializer_range: float = 0.02+    rms_norm_eps: float = 1e-6+    use_cache: bool = False+    pad_token_id: int | None = 0+    bos_token_id: int | None = 0+    eos_token_id: int | list[int] | None = 0+    tie_word_embeddings: bool = False+    rope_parameters: RopeParameters | dict | None = None+    attention_bias: bool = False+    attention_dropout: int | float = 0.0+    mlp_bias: bool = False+    head_dim: int | None = None+    classifier_pooling: Literal["cls", "mean"] = "cls"+    sliding_window: int = 128+    layer_types: list[str] | None = None+    prediction_head_size: int = 768+    classifier_dropout: float = 0.0++    def __post_init__(self, **kwargs):+        # BC -> existing checkpoints store attention windows per layer.+        if "attention_windows" in kwargs:+            attention_windows = kwargs.pop("attention_windows")+            self.sliding_window = next((window for window in attention_windows if window), self.sliding_window)+            if self.layer_types is None:+                self.layer_types = [+                    "sliding_attention" if window else "full_attention" for window in attention_windows+                ]+        if self.layer_types is None:+            self.layer_types = [+                "sliding_attention" if (i + 1) % 3 else "full_attention" for i in range(self.num_hidden_layers)+            ]++        # BC -> existing checkpoints store RoPE bases in rope_theta.+        rope_theta = kwargs.pop("rope_theta", None)+        if rope_theta is not None and self.rope_parameters is None:+            self.rope_parameters = {+                kind: {"rope_type": "default", "rope_theta": theta}+                for kind, theta in zip(self.layer_types, rope_theta)+            }+        if self.rope_parameters is None:+            self.rope_parameters = {+                "sliding_attention": {"rope_type": "default", "rope_theta": 10_000.0},+                "full_attention": {"rope_type": "default", "rope_theta": 2_560_000.0},+            }+        if self.num_key_value_heads is None:+            self.num_key_value_heads = self.num_attention_heads+        if self.head_dim is None:+            self.head_dim = self.hidden_size // self.num_attention_heads+        if self.num_key_value_heads is None:+            self.num_key_value_heads = self.num_attention_heads++        super().__post_init__(**kwargs)++    def validate_architecture(self):+        """Part of `@strict`-powered validation. Validates the architecture of the config."""+        if self.hidden_size % self.num_attention_heads != 0:+            raise ValueError(+                f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention "+                f"heads ({self.num_attention_heads})."+            )+++__all__ = ["UltrabertConfig"]

@@ -0,0 +1,655 @@+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+#           This file was automatically generated from src/transformers/models/ultrabert/modular_ultrabert.py.+#               Do NOT edit this file manually as any edits will be overwritten by the generation of+#             the file from the modular. If any change should be done, please apply the change to the+#                          modular_ultrabert.py file directly. One of our CI enforces this.+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨+# Copyright 2026 the HuggingFace Team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License").+# You may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++from collections.abc import Callable++import torch+from torch import nn+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss++from ... import initialization as init+from ...activations import ACT2FN+from ...cache_utils import Cache+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func+from ...masking_utils import create_bidirectional_mask, create_bidirectional_sliding_window_mask+from ...modeling_layers import GenericForQuestionAnswering, GradientCheckpointingLayer+from ...modeling_outputs import BaseModelOutput, MaskedLMOutput, SequenceClassifierOutput, TokenClassifierOutput+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel+from ...processing_utils import Unpack+from ...utils import TransformersKwargs, auto_docstring+from ...utils.generic import can_return_tuple, merge_with_config_defaults+from ...utils.output_capturing import capture_outputs+from .configuration_ultrabert import UltrabertConfig+++@use_kernel_forward_from_hub("RMSNorm")+class UltrabertRMSNorm(nn.Module):+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:+        """+        UltrabertRMSNorm is equivalent to T5LayerNorm+        """+        super().__init__()+        self.weight = nn.Parameter(torch.ones(hidden_size))+        self.variance_epsilon = eps++    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:+        input_dtype = hidden_states.dtype+        hidden_states = hidden_states.to(torch.float32)+        variance = hidden_states.pow(2).mean(-1, keepdim=True)+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)+        return self.weight * hidden_states.to(input_dtype)++    def extra_repr(self):+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"+++class UltrabertRotaryEmbedding(nn.Module):+    def __init__(self, config: UltrabertConfig, device=None):+        super().__init__()+        self.max_seq_len_cached = config.max_position_embeddings+        self.original_max_seq_len = config.max_position_embeddings+        self.config = config+        self.layer_types = sorted(set(config.layer_types))+        self.rope_type = {}+        for layer_type in self.layer_types:+            rope_params = self.config.rope_parameters[layer_type]+            if rope_params is None:+                continue++            self.rope_type[layer_type] = rope_params["rope_type"]+            rope_init_fn: Callable = self.compute_default_rope_parameters+            if self.rope_type[layer_type] != "default":+                rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type[layer_type]]+            curr_inv_freq, curr_attention_scaling = rope_init_fn(self.config, device, layer_type=layer_type)+            setattr(self, f"{layer_type}_inv_freq", nn.Buffer(curr_inv_freq, persistent=False))+            setattr(self, f"{layer_type}_original_inv_freq", nn.Buffer(curr_inv_freq.clone(), persistent=False))+            setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling)++    @staticmethod+    def compute_default_rope_parameters(+        config: UltrabertConfig, device=None, layer_type: str | None = None, **kwargs+    ) -> tuple[torch.Tensor, float]:+        """+        Computes the inverse frequencies according to the original RoPE implementation+        Args:+            config ([`~transformers.PreTrainedConfig`]):+                The model configuration.+            layer_type (`str`):+                The current layer type if the model has different RoPE parameters per type.+                Should not be used unless `config.layer_types is not None`++        Returns:+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).+        """+        # For backward compatibility standardize the `rope_parameters_dict` if it uses old format+        base = config.rope_parameters[layer_type]["rope_theta"]+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads++        attention_factor = 1.0  # Unused in this type of RoPE+        # Compute the inverse frequencies+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))+        return inv_freq.to(device), attention_factor++    @torch.no_grad()+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)+    def forward(self, x, position_ids, layer_type):+        inv_freq = getattr(self, f"{layer_type}_inv_freq")+        attention_scaling = getattr(self, f"{layer_type}_attention_scaling")++        freqs = position_ids[..., None].float() * inv_freq.to(device=x.device, dtype=torch.float)+        emb = torch.cat((freqs, freqs), dim=-1)+        cos = emb.cos() * attention_scaling+        sin = emb.sin() * attention_scaling++        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)+++def rotate_half(x):+    """Rotates half the hidden dims of the input."""+    x1 = x[..., : x.shape[-1] // 2]+    x2 = x[..., x.shape[-1] // 2 :]+    return torch.cat((-x2, x1), dim=-1)+++@use_kernel_forward_from_hub("rotary_pos_emb")+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):+    """Applies Rotary Position Embedding to the query and key tensors.++    Args:+        q (`torch.Tensor`): The query tensor.+        k (`torch.Tensor`): The key tensor.+        cos (`torch.Tensor`): The cosine part of the rotary embedding.+        sin (`torch.Tensor`): The sine part of the rotary embedding.+        unsqueeze_dim (`int`, *optional*, defaults to 1):+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.+    Returns:+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.+    """+    cos = cos.unsqueeze(unsqueeze_dim)+    sin = sin.unsqueeze(unsqueeze_dim)+    q_embed = (q * cos) + (rotate_half(q) * sin)+    k_embed = (k * cos) + (rotate_half(k) * sin)+    return q_embed, k_embed+++def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:+    """+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)+    """+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape+    if n_rep == 1:+        return hidden_states+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)+++def eager_attention_forward(+    module: nn.Module,+    query: torch.Tensor,+    key: torch.Tensor,+    value: torch.Tensor,+    attention_mask: torch.Tensor | None,+    scaling: float,+    dropout: float = 0.0,+    **kwargs: Unpack[TransformersKwargs],+):+    key_states = repeat_kv(key, module.num_key_value_groups)+    value_states = repeat_kv(value, module.num_key_value_groups)++    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling+    if attention_mask is not None:+        attn_weights = attn_weights + attention_mask++    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)+    attn_output = torch.matmul(attn_weights, value_states)+    attn_output = attn_output.transpose(1, 2).contiguous()++    return attn_output, attn_weights+++@use_kernelized_func(apply_rotary_pos_emb)+class UltrabertAttention(nn.Module):+    """Multi-headed attention from 'Attention Is All You Need' paper"""++    def __init__(self, config: UltrabertConfig, layer_idx: int):+        super().__init__()+        self.config = config+        self.layer_idx = layer_idx+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads+        self.scaling = self.head_dim**-0.5+        self.attention_dropout = config.attention_dropout+        self.is_causal = False++        self.q_proj = nn.Linear(+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias+        )+        self.k_proj = nn.Linear(+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias+        )+        self.v_proj = nn.Linear(+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias+        )+        self.o_proj = nn.Linear(+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias+        )+        # +1 is needed because flash attention sets inclusive boundaries (see modeling_flash_attention_utils.py)+        self.sliding_window = (+            config.sliding_window + 1 if config.layer_types[layer_idx] == "sliding_attention" else None+        )++    def forward(+        self,+        hidden_states: torch.Tensor,+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,+        attention_mask: torch.Tensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> tuple[torch.Tensor, torch.Tensor]:+        # Unlike EuroBERT, this has no KV cache and it uses a per-layer attention window.+        input_shape = hidden_states.shape[:-1]+        hidden_shape = (*input_shape, -1, self.head_dim)++        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)++        cos, sin = position_embeddings+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)++        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(+            self.config._attn_implementation, eager_attention_forward+        )++        attn_output, attn_weights = attention_interface(+            self,+            query_states,+            key_states,+            value_states,+            attention_mask,+            dropout=0.0 if not self.training else self.attention_dropout,+            scaling=self.scaling,+            sliding_window=self.sliding_window,+            **kwargs,+        )++        attn_output = attn_output.reshape(*input_shape, -1).contiguous()+        attn_output = self.o_proj(attn_output)+        return attn_output, attn_weights+++class UltrabertMLP(nn.Module):+    def __init__(self, config):+        super().__init__()+        self.config = config+        self.hidden_size = config.hidden_size+        self.intermediate_size = config.intermediate_size+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)+        self.act_fn = ACT2FN[config.hidden_act]++    def forward(self, x):+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))+        return down_proj+++class UltrabertDecoderLayer(GradientCheckpointingLayer):+    def __init__(self, config: UltrabertConfig, layer_idx: int):+        super().__init__()+        self.hidden_size = config.hidden_size++        self.self_attn = UltrabertAttention(config=config, layer_idx=layer_idx)++        self.mlp = UltrabertMLP(config)+        self.input_layernorm = UltrabertRMSNorm(config.hidden_size, eps=config.rms_norm_eps)+        self.post_attention_layernorm = UltrabertRMSNorm(config.hidden_size, eps=config.rms_norm_eps)++    def forward(+        self,+        hidden_states: torch.Tensor,+        attention_mask: torch.Tensor | None = None,+        position_ids: torch.LongTensor | None = None,+        past_key_values: Cache | None = None,+        use_cache: bool | None = False,+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> torch.Tensor:+        residual = hidden_states+        hidden_states = self.input_layernorm(hidden_states)+        # Self Attention+        hidden_states, _ = self.self_attn(+            hidden_states=hidden_states,+            attention_mask=attention_mask,+            position_ids=position_ids,+            past_key_values=past_key_values,+            use_cache=use_cache,+            position_embeddings=position_embeddings,+            **kwargs,+        )+        hidden_states = residual + hidden_states++        # Fully Connected+        residual = hidden_states+        hidden_states = self.post_attention_layernorm(hidden_states)+        hidden_states = self.mlp(hidden_states)+        hidden_states = residual + hidden_states+        return hidden_states+++@auto_docstring+class UltrabertPreTrainedModel(PreTrainedModel):+    config: UltrabertConfig+    base_model_prefix = "model"+    supports_gradient_checkpointing = True+    _no_split_modules = ["UltrabertDecoderLayer"]+    _skip_keys_device_placement = ["past_key_values"]+    _supports_flash_attn = True+    _supports_sdpa = True+    _supports_flex_attn = True++    _can_compile_fullgraph = True+    _supports_attention_backend = True+    _can_record_outputs = {+        "hidden_states": UltrabertDecoderLayer,+        "attentions": UltrabertAttention,+    }++    def _init_weights(self, module):+        super()._init_weights(module)+        if isinstance(module, UltrabertRotaryEmbedding):+            for layer_type in module.layer_types:+                rope_init_fn = module.compute_default_rope_parameters+                if module.rope_type[layer_type] != "default":+                    rope_init_fn = ROPE_INIT_FUNCTIONS[module.rope_type[layer_type]]+                curr_inv_freq, _ = rope_init_fn(module.config, layer_type=layer_type)+                init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)+                init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)+++@auto_docstring+class UltrabertModel(UltrabertPreTrainedModel):+    def __init__(self, config: UltrabertConfig):+        super().__init__(config)+        self.padding_idx = config.pad_token_id+        self.vocab_size = config.vocab_size++        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)+        self.layers = nn.ModuleList(+            [UltrabertDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]+        )+        self.norm = UltrabertRMSNorm(config.hidden_size, eps=config.rms_norm_eps)+        self.rotary_emb = UltrabertRotaryEmbedding(config=config)+        self.gradient_checkpointing = False++        # Initialize weights and apply final processing+        self.post_init()++    @merge_with_config_defaults+    @capture_outputs+    @auto_docstring+    def forward(+        self,+        input_ids: torch.Tensor | None = None,+        attention_mask: torch.Tensor | None = None,+        position_ids: torch.Tensor | None = None,+        inputs_embeds: torch.Tensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> BaseModelOutput:+        if (input_ids is None) ^ (inputs_embeds is not None):+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")+        if inputs_embeds is None:+            inputs_embeds = self.embed_tokens(input_ids)+        if position_ids is None:+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0)++        # Unlike EuroBERT, this selects attention masks and RoPE parameters by layer type.+        mask_kwargs = {"config": self.config, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask}+        attention_masks = {+            "full_attention": create_bidirectional_mask(**mask_kwargs),+            "sliding_attention": create_bidirectional_sliding_window_mask(**mask_kwargs),+        }++        position_embeddings = {+            layer_type: self.rotary_emb(inputs_embeds, position_ids, layer_type)+            for layer_type in set(self.config.layer_types)+        }+        hidden_states = inputs_embeds+        for i, encoder_layer in enumerate(self.layers):+            layer_type = self.config.layer_types[i]+            hidden_states = encoder_layer(+                hidden_states,+                attention_mask=attention_masks[layer_type],+                position_ids=position_ids,+                position_embeddings=position_embeddings[layer_type],+                **kwargs,+            )+        hidden_states = self.norm(hidden_states)+        return BaseModelOutput(last_hidden_state=hidden_states)+++class UltrabertPredictionHead(nn.Module):+    def __init__(self, config: UltrabertConfig):+        super().__init__()+        self.linear1 = nn.Linear(config.hidden_size, config.prediction_head_size, bias=False)+        self.linear2 = nn.Linear(config.prediction_head_size, config.hidden_size, bias=False)+        self.act_fn = ACT2FN["gelu_pytorch_tanh"]+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)++    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:+        hidden_states = self.linear1(hidden_states)+        hidden_states = self.act_fn(hidden_states)+        hidden_states = self.linear2(hidden_states)+        hidden_states = self.act_fn(hidden_states)+        return self.decoder(hidden_states)+++@auto_docstring(+    custom_intro="""+    The Ultrabert Model with a decoder head on top that is used for masked language modeling.+    """+)+class UltrabertForMaskedLM(UltrabertPreTrainedModel):+    _tied_weights_keys = {"head.decoder.weight": "model.embed_tokens.weight"}++    def __init__(self, config: UltrabertConfig):+        super().__init__(config)+        self.config = config+        self.model = UltrabertModel(config)+        self.head = UltrabertPredictionHead(config)+        self.decoder = nn.Identity()+        self.sparse_prediction = False+        self.sparse_pred_ignore_index = -100++        # Initialize weights and apply final processing+        self.post_init()++    def get_output_embeddings(self):+        # UltraBERT doesn't normally tie embeddings, so output embeddings differ from input embeddings+        return self.head.decoder++    def set_output_embeddings(self, new_embeddings):+        # UltraBERT doesn't normally tie embeddings, so output embeddings differ from input embeddings+        self.head.decoder = new_embeddings++    @can_return_tuple+    @auto_docstring+    def forward(+        self,+        input_ids: torch.LongTensor | None = None,+        attention_mask: torch.Tensor | None = None,+        position_ids: torch.Tensor | None = None,+        inputs_embeds: torch.Tensor | None = None,+        labels: torch.Tensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> tuple[torch.Tensor] | MaskedLMOutput:+        outputs = self.model(+            input_ids=input_ids,+            attention_mask=attention_mask,+            position_ids=position_ids,+            inputs_embeds=inputs_embeds,+            **kwargs,+        )+        last_hidden_state = outputs[0]++        if self.sparse_prediction and labels is not None:+            # flatten labels and output first+            labels = labels.view(-1)+            last_hidden_state = last_hidden_state.view(labels.shape[0], -1)++            # then filter out the non-masked tokens+            mask_tokens = labels != self.sparse_pred_ignore_index+            last_hidden_state = last_hidden_state[mask_tokens]+            labels = labels[mask_tokens]++        logits = self.decoder(self.head(last_hidden_state))++        loss = None+        if labels is not None:+            loss = self.loss_function(logits, labels, vocab_size=self.config.vocab_size, **kwargs)++        return MaskedLMOutput(+            loss=loss,+            logits=logits,+            hidden_states=outputs.hidden_states,+            attentions=outputs.attentions,+        )+++@auto_docstring(+    custom_intro="""+    The Ultrabert Model with a sequence classification head on top that performs pooling.+    """+)+class UltrabertForSequenceClassification(UltrabertPreTrainedModel):+    def __init__(self, config: UltrabertConfig):+        super().__init__(config)+        self.num_labels = config.num_labels+        self.config = config++        self.model = UltrabertModel(config)+        self.head = nn.Identity()+        self.drop = torch.nn.Dropout(config.classifier_dropout)+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)++        # Initialize weights and apply final processing+        self.post_init()++    @can_return_tuple+    @auto_docstring+    def forward(+        self,+        input_ids: torch.LongTensor | None = None,+        attention_mask: torch.Tensor | None = None,+        position_ids: torch.Tensor | None = None,+        inputs_embeds: torch.Tensor | None = None,+        labels: torch.Tensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> tuple[torch.Tensor] | SequenceClassifierOutput:+        outputs = self.model(+            input_ids=input_ids,+            attention_mask=attention_mask,+            position_ids=position_ids,+            inputs_embeds=inputs_embeds,+            **kwargs,+        )+        last_hidden_state = outputs[0]++        if self.config.classifier_pooling == "cls":+            last_hidden_state = last_hidden_state[:, 0]+        elif self.config.classifier_pooling == "mean":+            if attention_mask is None:+                attention_mask = torch.ones(+                    last_hidden_state.shape[:2], device=last_hidden_state.device, dtype=torch.bool+                )+            last_hidden_state = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1) / attention_mask.sum(+                dim=1, keepdim=True+            )++        pooled_output = self.head(last_hidden_state)+        pooled_output = self.drop(pooled_output)+        logits = self.classifier(pooled_output)++        loss = None+        if labels is not None:+            if self.config.problem_type is None:+                if self.num_labels == 1:+                    self.config.problem_type = "regression"+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):+                    self.config.problem_type = "single_label_classification"+                else:+                    self.config.problem_type = "multi_label_classification"++            if self.config.problem_type == "regression":+                loss_fct = MSELoss()+                if self.num_labels == 1:+                    loss = loss_fct(logits.squeeze(), labels.squeeze())+                else:+                    loss = loss_fct(logits, labels)+            elif self.config.problem_type == "single_label_classification":+                loss_fct = CrossEntropyLoss()+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))+            elif self.config.problem_type == "multi_label_classification":+                loss_fct = BCEWithLogitsLoss()+                loss = loss_fct(logits, labels)++        return SequenceClassifierOutput(+            loss=loss,+            logits=logits,+            hidden_states=outputs.hidden_states,+            attentions=outputs.attentions,+        )+++@auto_docstring(+    custom_intro="""+    The Ultrabert Model with a token classification head on top, e.g. for Named Entity Recognition (NER) tasks.+    """+)+class UltrabertForTokenClassification(UltrabertPreTrainedModel):+    def __init__(self, config: UltrabertConfig):+        super().__init__(config)+        self.num_labels = config.num_labels++        self.model = UltrabertModel(config)+        self.head = nn.Identity()+        self.drop = torch.nn.Dropout(config.classifier_dropout)+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)++        # Initialize weights and apply final processing+        self.post_init()++    @can_return_tuple+    @auto_docstring+    def forward(+        self,+        input_ids: torch.LongTensor | None = None,+        attention_mask: torch.Tensor | None = None,+        position_ids: torch.Tensor | None = None,+        inputs_embeds: torch.Tensor | None = None,+        labels: torch.Tensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> tuple[torch.Tensor] | TokenClassifierOutput:+        outputs = self.model(+            input_ids=input_ids,+            attention_mask=attention_mask,+            position_ids=position_ids,+            inputs_embeds=inputs_embeds,+            **kwargs,+        )+        last_hidden_state = outputs[0]++        last_hidden_state = self.head(last_hidden_state)+        last_hidden_state = self.drop(last_hidden_state)+        logits = self.classifier(last_hidden_state)++        loss = None+        if labels is not None:+            loss_fct = CrossEntropyLoss()+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))++        return TokenClassifierOutput(+            loss=loss,+            logits=logits,+            hidden_states=outputs.hidden_states,+            attentions=outputs.attentions,+        )+++class UltrabertForQuestionAnswering(GenericForQuestionAnswering, UltrabertPreTrainedModel):+    pass+++__all__ = [+    "UltrabertPreTrainedModel",+    "UltrabertModel",+    "UltrabertForMaskedLM",+    "UltrabertForSequenceClassification",+    "UltrabertForTokenClassification",+    "UltrabertForQuestionAnswering",+]

@@ -0,0 +1,281 @@+# Copyright 2026 the HuggingFace Team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License").+# You may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++from collections.abc import Callable+from typing import Literal++import torch+from torch import nn++from ...activations import ACT2FN+from ...configuration_utils import strict+from ...masking_utils import create_bidirectional_mask, create_bidirectional_sliding_window_mask+from ...modeling_layers import GenericForQuestionAnswering+from ...modeling_outputs import BaseModelOutput+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS+from ...processing_utils import Unpack+from ...utils import TransformersKwargs, auto_docstring+from ...utils.generic import merge_with_config_defaults+from ...utils.output_capturing import capture_outputs+from ..eurobert.configuration_eurobert import EuroBertConfig+from ..eurobert.modeling_eurobert import (+    EuroBertAttention,+    EuroBertDecoderLayer,+    EuroBertMLP,+    EuroBertModel,+    apply_rotary_pos_emb,+    eager_attention_forward,+)+from ..llama.modeling_llama import LlamaRMSNorm+from ..modernbert.modeling_modernbert import (+    ModernBertForMaskedLM,+    ModernBertForSequenceClassification,+    ModernBertForTokenClassification,+    ModernBertRotaryEmbedding,+)+from ..olmo3.modeling_olmo3 import Olmo3PreTrainedModel+++@auto_docstring(checkpoint="dmis-lab/UltraBERT-150M")+@strict+class UltrabertConfig(EuroBertConfig):+    r"""+    sliding_window (`int`, *optional*, defaults to 128):+        Inclusive attention radius on either side of each token in sliding attention layers.+    prediction_head_size (`int`, *optional*, defaults to 768):+        Intermediate dimension of the two-layer masked language modeling prediction head.+    classifier_pooling (`str`, *optional*, defaults to `"cls"`):+        Pooling method for sequence classification. Either `"cls"` or `"mean"`.+    """++    model_type = "ultrabert"++    vocab_size: int = 32000+    hidden_size: int = 4096+    intermediate_size: int = 11008+    num_hidden_layers: int = 32+    num_attention_heads: int = 32+    rms_norm_eps: float = 1e-6+    pad_token_id: int | None = 0+    bos_token_id: int | None = 0+    eos_token_id: int | list[int] | None = 0+    use_cache: bool = False+    max_position_embeddings: int = 32768+    sliding_window: int = 128+    layer_types: list[str] | None = None+    prediction_head_size: int = 768+    classifier_dropout: float = 0.0+    classifier_pooling: Literal["cls", "mean"] = "cls"+    pretraining_tp = AttributeError()+    mask_token_id = AttributeError()++    def __post_init__(self, **kwargs):+        # BC -> existing checkpoints store attention windows per layer.+        if "attention_windows" in kwargs:+            attention_windows = kwargs.pop("attention_windows")+            self.sliding_window = next((window for window in attention_windows if window), self.sliding_window)+            if self.layer_types is None:+                self.layer_types = [+                    "sliding_attention" if window else "full_attention" for window in attention_windows+                ]+        if self.layer_types is None:+            self.layer_types = [+                "sliding_attention" if (i + 1) % 3 else "full_attention" for i in range(self.num_hidden_layers)+            ]++        # BC -> existing checkpoints store RoPE bases in rope_theta.+        rope_theta = kwargs.pop("rope_theta", None)+        if rope_theta is not None and self.rope_parameters is None:+            self.rope_parameters = {+                kind: {"rope_type": "default", "rope_theta": theta}+                for kind, theta in zip(self.layer_types, rope_theta)+            }+        if self.rope_parameters is None:+            self.rope_parameters = {+                "sliding_attention": {"rope_type": "default", "rope_theta": 10_000.0},+                "full_attention": {"rope_type": "default", "rope_theta": 2_560_000.0},+            }+        super().__post_init__(**kwargs)+++class UltrabertRMSNorm(LlamaRMSNorm):+    pass+++class UltrabertRotaryEmbedding(ModernBertRotaryEmbedding):+    pass+++class UltrabertAttention(EuroBertAttention):+    def __init__(self, config: UltrabertConfig, layer_idx: int):+        super().__init__(config, layer_idx)+        # +1 is needed because flash attention sets inclusive boundaries (see modeling_flash_attention_utils.py)+        self.sliding_window = (+            config.sliding_window + 1 if config.layer_types[layer_idx] == "sliding_attention" else None+        )++    def forward(+        self,+        hidden_states: torch.Tensor,+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,+        attention_mask: torch.Tensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> tuple[torch.Tensor, torch.Tensor]:+        # Unlike EuroBERT, this has no KV cache and it uses a per-layer attention window.+        input_shape = hidden_states.shape[:-1]+        hidden_shape = (*input_shape, -1, self.head_dim)++        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)++        cos, sin = position_embeddings+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)++        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(+            self.config._attn_implementation, eager_attention_forward+        )++        attn_output, attn_weights = attention_interface(+            self,+            query_states,+            key_states,+            value_states,+            attention_mask,+            dropout=0.0 if not self.training else self.attention_dropout,+            scaling=self.scaling,+            sliding_window=self.sliding_window,+            **kwargs,+        )++        attn_output = attn_output.reshape(*input_shape, -1).contiguous()+        attn_output = self.o_proj(attn_output)+        return attn_output, attn_weights+++class UltrabertMLP(EuroBertMLP):+    pass+++class UltrabertDecoderLayer(EuroBertDecoderLayer):+    pass+++class UltrabertPreTrainedModel(Olmo3PreTrainedModel):+    config: UltrabertConfig+++class UltrabertModel(EuroBertModel):+    @merge_with_config_defaults+    @capture_outputs+    @auto_docstring+    def forward(+        self,+        input_ids: torch.Tensor | None = None,+        attention_mask: torch.Tensor | None = None,+        position_ids: torch.Tensor | None = None,+        inputs_embeds: torch.Tensor | None = None,+        **kwargs: Unpack[TransformersKwargs],+    ) -> BaseModelOutput:+        if (input_ids is None) ^ (inputs_embeds is not None):+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")+        if inputs_embeds is None:+            inputs_embeds = self.embed_tokens(input_ids)+        if position_ids is None:+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0)++        # Unlike EuroBERT, this selects attention masks and RoPE parameters by layer type.+        mask_kwargs = {"config": self.config, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask}+        attention_masks = {+            "full_attention": create_bidirectional_mask(**mask_kwargs),+            "sliding_attention": create_bidirectional_sliding_window_mask(**mask_kwargs),+        }++        position_embeddings = {+            layer_type: self.rotary_emb(inputs_embeds, position_ids, layer_type)+            for layer_type in set(self.config.layer_types)+        }+        hidden_states = inputs_embeds+        for i, encoder_layer in enumerate(self.layers):+            layer_type = self.config.layer_types[i]+            hidden_states = encoder_layer(+                hidden_states,+                attention_mask=attention_masks[layer_type],+                position_ids=position_ids,+                position_embeddings=position_embeddings[layer_type],+                **kwargs,+            )+        hidden_states = self.norm(hidden_states)+        return BaseModelOutput(last_hidden_state=hidden_states)+++class UltrabertPredictionHead(nn.Module):+    def __init__(self, config: UltrabertConfig):+        super().__init__()+        self.linear1 = nn.Linear(config.hidden_size, config.prediction_head_size, bias=False)+        self.linear2 = nn.Linear(config.prediction_head_size, config.hidden_size, bias=False)+        self.act_fn = ACT2FN["gelu_pytorch_tanh"]+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)++    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:+        hidden_states = self.linear1(hidden_states)+        hidden_states = self.act_fn(hidden_states)+        hidden_states = self.linear2(hidden_states)+        hidden_states = self.act_fn(hidden_states)+        return self.decoder(hidden_states)+++class UltrabertForMaskedLM(ModernBertForMaskedLM):+    _tied_weights_keys = {"head.decoder.weight": "model.embed_tokens.weight"}++    def __init__(self, config: UltrabertConfig):+        super().__init__(config)+        self.decoder = nn.Identity()+        self.sparse_prediction = False+        self.sparse_pred_ignore_index = -100++    def get_output_embeddings(self):+        # UltraBERT doesn't normally tie embeddings, so output embeddings differ from input embeddings+        return self.head.decoder++    def set_output_embeddings(self, new_embeddings):+        # UltraBERT doesn't normally tie embeddings, so output embeddings differ from input embeddings+        self.head.decoder = new_embeddings+++class UltrabertForSequenceClassification(ModernBertForSequenceClassification):+    def __init__(self, config: UltrabertConfig):+        super().__init__(config)+        self.head = nn.Identity()+++class UltrabertForTokenClassification(ModernBertForTokenClassification):+    def __init__(self, config: UltrabertConfig):+        super().__init__(config)+        self.head = nn.Identity()+++class UltrabertForQuestionAnswering(GenericForQuestionAnswering, UltrabertPreTrainedModel):+    pass+++__all__ = [+    "UltrabertConfig",+    "UltrabertPreTrainedModel",+    "UltrabertModel",+    "UltrabertForMaskedLM",+    "UltrabertForSequenceClassification",+    "UltrabertForTokenClassification",+    "UltrabertForQuestionAnswering",+]

GitHub publishes no diff for this file. View it on GitHub ↗

@@ -0,0 +1,369 @@+# Copyright 2026 the HuggingFace Team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License").+# You may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++import unittest++from transformers import AutoTokenizer, UltrabertConfig, is_torch_available+from transformers.testing_utils import Expectations, require_torch, slow, torch_device++from ...test_configuration_common import ConfigTester+from ...test_memory_cleanup_mixin import MemoryCleanupMixin+from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask+from ...test_pipeline_mixin import PipelineTesterMixin+++if is_torch_available():+    import torch++    from transformers import (+        AutoModel,+        AutoModelForMaskedLM,+        AutoModelForSequenceClassification,+        UltrabertForMaskedLM,+        UltrabertForQuestionAnswering,+        UltrabertForSequenceClassification,+        UltrabertForTokenClassification,+        UltrabertModel,+    )+++class UltrabertModelTester:+    if is_torch_available():+        base_model_class = UltrabertModel++    def __init__(+        self,+        parent,+        batch_size=13,+        seq_length=7,+        is_training=True,+        use_input_mask=True,+        vocab_size=99,+        hidden_size=32,+        num_hidden_layers=2,+        num_attention_heads=4,+        num_key_value_heads=2,+        intermediate_size=37,+        hidden_act="silu",+        attention_dropout=0.0,+        max_position_embeddings=512,+        initializer_range=0.02,+        num_labels=3,+        pad_token_id=0,+        sliding_window=2,+        prediction_head_size=32,+        layer_types=None,+    ):+        self.parent = parent+        self.batch_size = batch_size+        self.seq_length = seq_length+        self.is_training = is_training+        self.use_input_mask = use_input_mask+        self.vocab_size = vocab_size+        self.hidden_size = hidden_size+        self.num_hidden_layers = num_hidden_layers+        self.num_attention_heads = num_attention_heads+        self.num_key_value_heads = num_key_value_heads+        self.intermediate_size = intermediate_size+        self.hidden_act = hidden_act+        self.attention_dropout = attention_dropout+        self.max_position_embeddings = max_position_embeddings+        self.initializer_range = initializer_range+        self.num_labels = num_labels+        self.pad_token_id = pad_token_id+        self.sliding_window = sliding_window+        self.prediction_head_size = prediction_head_size+        self.layer_types = (+            layer_types+            if layer_types is not None+            else ["sliding_attention" if i % 2 == 0 else "full_attention" for i in range(num_hidden_layers)]+        )++    def get_config(self, **kwargs):+        options = {+            "vocab_size": self.vocab_size,+            "hidden_size": self.hidden_size,+            "num_hidden_layers": self.num_hidden_layers,+            "num_attention_heads": self.num_attention_heads,+            "num_key_value_heads": self.num_key_value_heads,+            "intermediate_size": self.intermediate_size,+            "hidden_act": self.hidden_act,+            "attention_dropout": self.attention_dropout,+            "max_position_embeddings": self.max_position_embeddings,+            "initializer_range": self.initializer_range,+            "pad_token_id": self.pad_token_id,+            "sliding_window": self.sliding_window,+            "prediction_head_size": self.prediction_head_size,+            "layer_types": self.layer_types,+        }+        options.update(kwargs)+        return UltrabertConfig(**options)++    def prepare_config_and_inputs(self):+        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)+        attention_mask = random_attention_mask([self.batch_size, self.seq_length]) if self.use_input_mask else None+        sequence_labels = ids_tensor([self.batch_size], self.num_labels)+        token_labels = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)+        return self.get_config(), input_ids, attention_mask, sequence_labels, token_labels++    def create_and_check_model(self, config, input_ids, attention_mask, sequence_labels, token_labels):+        model = UltrabertModel(config).to(torch_device).eval()+        result = model(input_ids, attention_mask=attention_mask)+        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))+        result = model(input_ids)+        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))++    def create_and_check_for_masked_lm(self, config, input_ids, attention_mask, sequence_labels, token_labels):+        model = UltrabertForMaskedLM(config).to(torch_device).eval()+        result = model(input_ids, attention_mask=attention_mask, labels=token_labels)+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))++    def create_and_check_for_sequence_classification(+        self, config, input_ids, attention_mask, sequence_labels, token_labels+    ):+        config.num_labels = self.num_labels+        model = UltrabertForSequenceClassification(config).to(torch_device).eval()+        result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels)+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_labels))++    def create_and_check_for_token_classification(+        self, config, input_ids, attention_mask, sequence_labels, token_labels+    ):+        config.num_labels = self.num_labels+        labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels)+        labels[:, -1] = -100+        model = UltrabertForTokenClassification(config).to(torch_device).eval()+        result = model(input_ids, attention_mask=attention_mask, labels=labels)+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels))++    def create_and_check_for_question_answering(+        self, config, input_ids, attention_mask, sequence_labels, token_labels+    ):+        model = UltrabertForQuestionAnswering(config).to(torch_device).eval()+        start_positions = ids_tensor([self.batch_size], self.seq_length)+        end_positions = start_positions + 1+        result = model(+            input_ids, attention_mask=attention_mask, start_positions=start_positions, end_positions=end_positions+        )+        self.parent.assertEqual(result.start_logits.shape, (self.batch_size, self.seq_length))+        self.parent.assertEqual(result.end_logits.shape, (self.batch_size, self.seq_length))++    def prepare_config_and_inputs_for_common(self):+        config, input_ids, attention_mask, _, _ = self.prepare_config_and_inputs()+        return config, {"input_ids": input_ids, "attention_mask": attention_mask}+++@require_torch+class UltrabertModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):+    all_model_classes = (+        (+            UltrabertModel,+            UltrabertForMaskedLM,+            UltrabertForSequenceClassification,+            UltrabertForTokenClassification,+            UltrabertForQuestionAnswering,+        )+        if is_torch_available()+        else ()+    )+    pipeline_model_mapping = (+        {+            "feature-extraction": UltrabertModel,+            "fill-mask": UltrabertForMaskedLM,+            "text-classification": UltrabertForSequenceClassification,+            "token-classification": UltrabertForTokenClassification,+            "zero-shot": UltrabertForSequenceClassification,+        }+        if is_torch_available()+        else {}+    )+    model_tester_class = UltrabertModelTester++    def setUp(self):+        self.model_tester = UltrabertModelTester(self)+        self.config_tester = ConfigTester(self, config_class=UltrabertConfig, hidden_size=32, num_attention_heads=4)++    def test_config(self):+        self.config_tester.run_common_tests()++    @unittest.skip(reason="SDPA's flash backend doesn't support the mandatory sliding-window attention masks.")+    def test_sdpa_can_dispatch_on_flash(self):+        pass++    def test_model(self):+        self.model_tester.create_and_check_model(*self.model_tester.prepare_config_and_inputs())++    def test_for_masked_lm(self):+        self.model_tester.create_and_check_for_masked_lm(*self.model_tester.prepare_config_and_inputs())++    def test_for_sequence_classification(self):+        self.model_tester.create_and_check_for_sequence_classification(*self.model_tester.prepare_config_and_inputs())++    def test_for_token_classification(self):+        self.model_tester.create_and_check_for_token_classification(*self.model_tester.prepare_config_and_inputs())++    def test_for_question_answering(self):+        self.model_tester.create_and_check_for_question_answering(*self.model_tester.prepare_config_and_inputs())++    def test_legacy_config(self):+        attention_windows = [2, 2, None]+        rope_theta = [10000.0, 10000.0, 2560000.0]+        config = self.model_tester.get_config(+            num_hidden_layers=3, layer_types=None, attention_windows=attention_windows, rope_theta=rope_theta+        )+        self.assertEqual(config.sliding_window, 2)+        self.assertEqual(+            config.layer_types, ["sliding_attention" if window else "full_attention" for window in attention_windows]+        )+        for layer_type, base in zip(config.layer_types, rope_theta):+            self.assertEqual(config.rope_parameters[layer_type], {"rope_type": "default", "rope_theta": base})+        self.assertNotIn("attention_windows", config.to_dict())+        self.assertNotIn("rope_theta", config.to_dict())++    def test_rope_parameters_override_legacy_theta(self):+        rope_parameters = {+            "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0},+            "full_attention": {"rope_type": "linear", "rope_theta": 2560000.0, "factor": 2.0},+        }+        config = self.model_tester.get_config(+            rope_parameters=rope_parameters, rope_theta=[12345.0] * self.model_tester.num_hidden_layers+        )+        self.assertEqual(config.rope_parameters, rope_parameters)++    def test_attention_window_boundaries(self):+        config = self.model_tester.get_config(attn_implementation="eager")+        model = UltrabertModel(config).to(torch_device).eval()+        input_ids = ids_tensor([self.model_tester.batch_size, self.model_tester.seq_length], config.vocab_size)+        with torch.no_grad():+            outputs = model(input_ids, output_attentions=True)+        for attention, layer_type in zip(outputs.attentions, model.config.layer_types):+            if layer_type == "sliding_attention":+                window = model.config.sliding_window+                self.assertTrue(torch.all(attention[:, :, 0, window] > 0))+                self.assertTrue(torch.all(attention[:, :, 0, window + 1] == 0))+                self.assertTrue(torch.all(attention[:, :, window, 0] > 0))+                self.assertTrue(torch.all(attention[:, :, window + 1, 0] == 0))+            else:+                self.assertTrue(torch.all(attention[:, :, 0, -1] > 0))+++@require_torch+@slow+class UltrabertIntegrationTest(MemoryCleanupMixin, unittest.TestCase):+    """Check native loading against float32 outputs from the original checkpoint code."""++    def load_checkpoint(self, auto_class, checkpoint):+        model = auto_class.from_pretrained(+            checkpoint,+            dtype=torch.float32,+            attn_implementation="sdpa",+        )+        tokenizer = AutoTokenizer.from_pretrained(checkpoint)+        return model, tokenizer++    def test_inference_embedding(self):+        model, tokenizer = self.load_checkpoint(AutoModel, "dmis-lab/UltraBERT-Embed-150M")+        inputs = tokenizer(+            ["The quick brown fox jumps over the lazy dog. " * 30, "Hello world!"],+            padding=True,+            return_tensors="pt",+        )+        expected_slice = Expectations(+            {+                (None, None): torch.tensor(+                    [+                        [3.815864, 1.319243, -2.833093],+                        [0.084467, -0.133652, -0.315864],+                        [3.821545, -4.447476, -2.473820],+                        [1.938878, -0.788768, 6.423543],+                        [0.796432, -5.030339, -2.032893],+                    ]+                )+            }+        ).get_expectation()+        expected_padded_cls = Expectations(+            {(None, None): torch.tensor([0.821952, -1.553840, 1.073606])}+        ).get_expectation()+        with torch.no_grad():+            output = model(**inputs).last_hidden_state+        self.assertEqual(output.shape, (2, 302, 768))+        torch.testing.assert_close(output[0, [0, 127, 128, 129, 256], :3], expected_slice, rtol=1e-4, atol=1e-4)+        torch.testing.assert_close(output[1, 0, :3], expected_padded_cls, rtol=1e-4, atol=1e-4)++    def test_inference_masked_lm_150m(self):+        model, tokenizer = self.load_checkpoint(AutoModelForMaskedLM, "dmis-lab/UltraBERT-150M")+        inputs = tokenizer(+            ["The capital of France is <|mask|>.", "The capital of Germany is <|mask|>."],+            padding=True,+            return_tensors="pt",+        )+        expected_slice = Expectations(+            {+                (None, None): torch.tensor(+                    [+                        [-173.299500, -172.092453, -172.083084, -166.659607, -168.356415],+                        [-166.776337, -165.526764, -165.573578, -160.336792, -161.950073],+                    ]+                )+            }+        ).get_expectation()+        with torch.no_grad():+            logits = model(**inputs).logits+        self.assertEqual(logits.shape, (2, 8, 49152))+        masked_logits = logits[inputs.input_ids == tokenizer.mask_token_id]+        torch.testing.assert_close(masked_logits[:, :5], expected_slice, rtol=1e-4, atol=1e-4)+        self.assertEqual(tokenizer.batch_decode(masked_logits.argmax(-1)[:, None]), [" Paris", " Berlin"])++    def test_inference_masked_lm_380m(self):+        model, tokenizer = self.load_checkpoint(AutoModelForMaskedLM, "dmis-lab/UltraBERT-380M")+        inputs = tokenizer(+            ["The capital of France is <|mask|>.", "The capital of Germany is <|mask|>."],+            padding=True,+            return_tensors="pt",+        )+        expected_slice = Expectations(+            {+                (None, None): torch.tensor(+                    [+                        [-195.506729, -193.607315, -193.664520, -189.095596, -189.571594],+                        [-189.560608, -187.724716, -187.760635, -183.072098, -183.630432],+                    ]+                )+            }+        ).get_expectation()+        with torch.no_grad():+            logits = model(**inputs).logits+        self.assertEqual(logits.shape, (2, 8, 49152))+        masked_logits = logits[inputs.input_ids == tokenizer.mask_token_id]+        torch.testing.assert_close(masked_logits[:, :5], expected_slice, rtol=1e-4, atol=1e-4)+        self.assertEqual(tokenizer.batch_decode(masked_logits.argmax(-1)[:, None]), [" Paris", " Berlin"])++    def test_inference_sequence_classification(self):+        model, tokenizer = self.load_checkpoint(+            AutoModelForSequenceClassification,+            "cross-encoder-testing/ultrabert-150m-reranker-test",+        )+        inputs = tokenizer(+            [+                ["What is the capital of France?", "Paris is the capital of France."],+                ["What is the capital of France?", "Berlin is the capital of Germany and its largest city."],+            ],+            padding=True,+            return_tensors="pt",+        )+        expected = Expectations({(None, None): torch.tensor([[3.867966175], [-4.749768734]])}).get_expectation()+        with torch.no_grad():+            logits = model(**inputs).logits+        self.assertEqual(logits.shape, (2, 1))+        torch.testing.assert_close(logits, expected, rtol=1e-4, atol=1e-4)