UltraBERT-380M-clairic-extraction

README.md

@@ -0,0 +1,108 @@
+---
+library_name: transformers
+base_model:
+- dmis-lab/UltraBERT-380M
+license: apache-2.0
+language:
+- en
+pipeline_tag: token-classification
+datasets:
+- tasksource/tasksource-jev-typed-decisions
+tags:
+- information-extraction
+- extractive-qa
+- ner
+- span-extraction
+- tasksource
+---
+
+# UltraBERT-380M-clairic-extraction
+
+A 380M-parameter extractor for Clairic-style requests: give it a text and a set of fields, each described by a short instruction, and it returns the matching spans for each field (one answer, a list, or nothing when the text has no answer). Fields can be questions ("What is the total due?") or entity types ("Every company mentioned in the text"), so one model covers extractive question answering and named-entity recognition.
+
+[UltraBERT-380M](https://huggingface.co/dmis-lab/UltraBERT-380M) trained on the `extraction` config of [tasksource-jev-typed-decisions](https://huggingface.co/datasets/tasksource/tasksource-jev-typed-decisions) (release v2026.10.24.1, 1M instruction/text pairs, 60k steps): extractive QA and span-tagging datasets from [tasksource](https://github.com/sileod/tasksource), rewritten as (instruction, text, spans) pairs. The model tags the answer tokens of the text (B/I/O) conditioned on the instruction. Inputs up to 2048 tokens.
+
+It is the larger sibling of [UltraBERT-150M-clairic-extraction](https://huggingface.co/tasksource/UltraBERT-150M-clairic-extraction), with the same interface: more accurate, especially on list fields, and better at leaving a field empty when the text has no answer, for about 2.5x the compute.
+
+Not affiliated with Cloudflare; "Clairic-style" refers to the request format.
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModelForTokenClassification, AutoTokenizer
+
+name = "tasksource/UltraBERT-380M-clairic-extraction"
+tok = AutoTokenizer.from_pretrained(name)
+model = AutoModelForTokenClassification.from_pretrained(name, trust_remote_code=True).eval()
+
+def extract(text, fields):
+    """fields: {name: instruction}. Returns {name: [extracted spans]}, [] when the text has no answer."""
+    enc = tok(list(fields.values()), [text] * len(fields), truncation="only_second", max_length=2048,
+              padding=True, return_offsets_mapping=True, return_tensors="pt")
+    offsets = enc.pop("offset_mapping")
+    with torch.no_grad():
+        tags = model(**enc).logits.argmax(-1).tolist()  # 0 = O, 1 = B, 2 = I
+    out = {}
+    for f, field in enumerate(fields):
+        spans, start, seq = [], None, enc.sequence_ids(f) + [None]
+        for i, tag in enumerate(tags[f] + [0]):
+            if start is not None and (seq[i] != 1 or tag != 2):
+                spans.append(text[offsets[f][start][0]:offsets[f][i - 1][1]].strip().rstrip(",;:"))
+                start = None
+            if seq[i] == 1 and tag in (1, 2) and start is None:
+                start = i
+        out[field] = spans
+    return out
+
+text = "Invoice INV-2291 from Acme Corp to Globex, issued 3 March 2025. Total due: $1,240.50. Contact: billing@acme.com."
+print(extract(text, {
+    "invoice_number": "The invoice number.",
+    "companies": "Every company mentioned in the text, as written.",
+    "total": "The total amount due.",
+    "phone": "The contact phone number.",
+}))
+# {'invoice_number': ['INV-2291'], 'companies': ['Acme Corp', 'Globex'], 'total': ['$1,240.50'], 'phone': []}
+```
+
+`trust_remote_code=True` is needed for the UltraBERT architecture.
+
+## Evaluation
+
+Test split of the extraction config (revision v2026.10.20.7), 100 rows per source, procedurally generated sources excluded. *text* fields (one answer or none) are scored by SQuAD F1, the empty answer counting as an answer. *list* fields (every span of a kind) are scored by micro F1 over normalized values. Scores are means over sources.
+
+| Model | Params | QA (F1) | NER / span lists (F1) |
+|---|---|---|---|
+| **UltraBERT-380M-clairic-extraction** | 380M | **0.756** | **0.750** |
+| [UltraBERT-150M-clairic-extraction](https://huggingface.co/tasksource/UltraBERT-150M-clairic-extraction) | 150M | 0.720 | 0.702 |
+| [gliner-multitask-large-v0.5](https://huggingface.co/knowledgator/gliner-multitask-large-v0.5) | 440M | 0.474 | 0.408 |
+| [deberta-v3-base-squad2](https://huggingface.co/deepset/deberta-v3-base-squad2) | 184M | 0.498 | 0.060 |
+
+The test sources are the training sources (their test splits), while the baselines see most of them zero-shot, so this measures performance on these tasks rather than on unseen ones.
+
+Out of domain (MRQA test sets the model never trained on, plus FinEntity), F1:
+
+| Model | BioASQ | RACE | RelationExtraction | TextbookQA | FinEntity | Mean |
+|---|---|---|---|---|---|---|
+| **UltraBERT-380M-clairic-extraction** | 0.626 | 0.412 | 0.870 | 0.482 | 0.717 | 0.622 |
+| [gliner-multitask-large-v0.5](https://huggingface.co/knowledgator/gliner-multitask-large-v0.5) | **0.643** | **0.458** | **0.899** | **0.545** | **0.913** | **0.692** |
+
+GLiNER-multitask stays ahead out of domain.
+
+## Training data
+
+The extraction config gathers extractive QA (SQuAD 2.0, Natural Questions, HotpotQA, NewsQA, DuoRC, QuAC, CUAD, Qasper, TimeQA, WebSRC, SWDE, and others) and span tagging (CoNLL-2003, OntoNotes, Few-NERD, WNUT, CrossNER, biomedical, legal and financial NER, skills, aspects, keyphrases, and others), plus procedurally generated documents. A few sources (e.g. GUM, PreCo) are licensed for non-commercial use only; check the dataset card for per-source licenses.
+
+## Citation
+
+```bibtex
+@inproceedings{sileo-2024-tasksource,
+    title = "tasksource: A Large Collection of {NLP} tasks with a Structured Dataset Preprocessing Framework",
+    author = "Sileo, Damien",
+    booktitle = "Proceedings of the 2024 Joint International Conference on Computational Linguistics, Language Resources and Evaluation (LREC-COLING 2024)",
+    year = "2024",
+    address = "Torino, Italia",
+    pages = "15655--15684",
+    url = "https://aclanthology.org/2024.lrec-main.1361",
+}
+```

config.json

@@ -0,0 +1,104 @@
+{
+  "architectures": [
+    "UltrabertForTokenClassification"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attention_windows": [
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_ultrabert.UltrabertConfig",
+    "AutoModel": "modeling_ultrabert.UltrabertModel",
+    "AutoModelForMaskedLM": "modeling_ultrabert.UltrabertForMaskedLM",
+    "AutoModelForQuestionAnswering": "modeling_ultrabert.UltrabertForQuestionAnswering",
+    "AutoModelForSequenceClassification": "modeling_ultrabert.UltrabertForSequenceClassification",
+    "AutoModelForTokenClassification": "modeling_ultrabert.UltrabertForTokenClassification"
+  },
+  "bos_token_id": 0,
+  "classifier_dropout": 0.0,
+  "classifiers_size": [
+    3
+  ],
+  "dtype": "float32",
+  "eos_token_id": 0,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 1024,
+  "id2label": {
+    "0": "O",
+    "1": "B-answer",
+    "2": "I-answer"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "label2id": {
+    "B-answer": 1,
+    "I-answer": 2,
+    "O": 0
+  },
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "ultrabert",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 8,
+  "pad_token_id": 0,
+  "prediction_head_size": 1024,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": [
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0
+  ],
+  "tasks": [
+    "jev-extraction"
+  ],
+  "tie_word_embeddings": false,
+  "transformers_version": "5.17.0",
+  "vocab_size": 49152
+}

configuration_ultrabert.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class UltrabertConfig(PretrainedConfig):
+    model_type = "ultrabert"
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=4096,
+        intermediate_size=11008,
+        prediction_head_size=None,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        attention_windows=None,
+        hidden_act="silu",
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        pad_token_id=0,
+        bos_token_id=0,
+        eos_token_id=0,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        attention_bias=False,
+        attention_dropout=0.0,
+        mlp_bias=False,
+        head_dim=None,
+        classifier_dropout=0.0,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.prediction_head_size = prediction_head_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads or num_attention_heads
+        self.attention_windows = (
+            attention_windows
+            if isinstance(attention_windows, Iterable)
+            else [attention_windows] * num_hidden_layers
+        )
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = (
+            rope_theta
+            if isinstance(rope_theta, Iterable)
+            else [rope_theta] * num_hidden_layers
+        )
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.mlp_bias = mlp_bias
+        self.head_dim = head_dim or self.hidden_size // self.num_attention_heads
+        self.classifier_dropout = classifier_dropout

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:188edece3e65295207c3137de5af06c4011afbe1f81c7091ca049ce0048b3ce3
+size 1510188284

modeling_ultrabert.py

@@ -0,0 +1,766 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from examples/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.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+from collections.abc import Callable
+
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    MaskedLMOutput,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+from transformers.utils.generic import merge_with_config_defaults
+from transformers.utils.output_capturing import capture_outputs
+
+from transformers.cache_utils import Cache
+from transformers.integrations import use_kernel_forward_from_hub, use_kernelized_func
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_rope_utils import dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.utils.deprecation import deprecate_kwarg
+from transformers.utils.generic import maybe_autocast
+from .configuration_ultrabert import UltrabertConfig
+
+
+class RotaryEmbedding(nn.Module):
+    @deprecate_kwarg("device", version="5.18")
+    def __init__(self, base: float, dim: int, device: torch.device | None = None):
+        super().__init__()
+        self.rope_type = "default"
+        self.attention_scaling = 1
+        self.base, self.dim = base, dim
+
+        inv_freq = torch.arange(0, dim, 2, dtype=torch.int64)
+        inv_freq = 1.0 / (base ** (inv_freq.to(device=device, dtype=torch.float) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    @staticmethod
+    @deprecate_kwarg("device", version="5.18")
+    def compute_default_rope_parameters(
+        config: UltrabertConfig, device=None, **kwargs
+    ) -> tuple[torch.Tensor, float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+        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).
+        """
+        base = config.rope_parameters["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):
+        inv_freq_expanded = (
+            self.inv_freq[None, :, None]
+            .expand(position_ids.shape[0], -1, 1)
+            .to(dtype=torch.float, device=x.device)
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = (
+            x.device.type
+            if isinstance(x.device.type, str) and x.device.type != "mps"
+            else "cpu"
+        )
+        # Disable any outside autocast context if any, to really force fp32
+        with maybe_autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.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,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        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
+        )
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        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,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+@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 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 UltrabertLayer(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 = "ultrabert"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["UltrabertLayer"]
+    _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": UltrabertLayer,
+        "attentions": UltrabertAttention,
+    }
+    config_class = UltrabertConfig
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, RotaryEmbedding):
+            inv_freq = torch.arange(0, module.dim, 2, dtype=torch.int64)
+            inv_freq = 1.0 / (
+                module.base
+                ** (
+                    inv_freq.to(device=module.inv_freq.device, dtype=torch.float)
+                    / module.dim
+                )
+            )
+            module.inv_freq.copy_(inv_freq)
+
+
+def create_local_attention_mask(
+    seq_len: int,
+    sliding_window: int | None = None,
+    device: torch.device | None = None,
+) -> torch.Tensor:
+    attention_mask = torch.ones(1, seq_len, seq_len, dtype=torch.bool, device=device)
+    if sliding_window:
+        attention_mask = attention_mask.tril(sliding_window).triu(-sliding_window)
+    return attention_mask
+
+
+@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(
+            [UltrabertLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = UltrabertRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = nn.ModuleDict()
+        self.gradient_checkpointing = False
+        for base in set(config.rope_theta):
+            name = str(base).replace(".", "_")
+            self.rotary_emb[name] = RotaryEmbedding(base, config.head_dim)
+
+        # 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("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.size(1))
+            position_ids = position_ids.to(inputs_embeds.device).unsqueeze(0)
+        if (
+            attention_mask is None
+            and self.config._attn_implementation != "flash_attention_2"
+        ):
+            attention_mask = inputs_embeds.new_ones(inputs_embeds.shape[:2])
+        varlen_kwargs = {
+            k: kwargs[k]
+            for k in ("cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k")
+            if kwargs.get(k) is not None
+        }
+
+        # Create attention masks for different local window sizes and store by window.
+        attention_mask_dict = {}
+        for local_window in list(set(self.config.attention_windows)) + [None]:
+            if self.config._attn_implementation == "flash_attention_2":
+                # For FlashAttention 2, rely on the `sliding_window` keyword instead of
+                # constructing masks manually.
+                attention_mask_dict[local_window] = attention_mask
+            else:
+                window_mask = create_local_attention_mask(
+                    seq_len=attention_mask.size(1),
+                    sliding_window=local_window,
+                    device=attention_mask.device,
+                )
+                window_mask = window_mask & attention_mask.bool().unsqueeze(1)
+                attention_mask_dict[local_window] = (
+                    torch.where(window_mask, 0.0, torch.finfo(inputs_embeds.dtype).min)
+                    .to(inputs_embeds.dtype)
+                    .unsqueeze(1)
+                )
+
+        # Create position embeddings to be shared across the decoder layers.
+        hidden_states, position_embeddings = inputs_embeds, {}
+        for base in set(self.config.rope_theta):
+            rotary_emb = self.rotary_emb[str(base).replace(".", "_")]
+            position_embeddings[base] = rotary_emb(hidden_states, position_ids)
+
+        for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
+            local_window = self.config.attention_windows[i]
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask_dict[local_window],
+                position_ids=position_ids,
+                position_embeddings=position_embeddings[self.config.rope_theta[i]],
+                sliding_window=(
+                    local_window + 1
+                    if local_window is not None
+                    and self.config._attn_implementation == "flash_attention_2"
+                    else local_window
+                ),
+                **varlen_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__()
+        if (dim := config.prediction_head_size) is not None:
+            self.linear1 = nn.Linear(config.hidden_size, dim, bias=False)
+            self.linear2 = nn.Linear(dim, config.hidden_size, bias=False)
+            self.act_fn = ACT2FN["gelu_pytorch_tanh"]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "linear1"):
+            return self.act_fn(self.linear2(self.act_fn(self.linear1(x))))
+        return x
+
+
+@auto_docstring
+class UltrabertForMaskedLM(UltrabertPreTrainedModel):
+    _tied_weights_keys = {"cls.weight": "ultrabert.embed_tokens.weight"}
+
+    def __init__(self, config: UltrabertConfig):
+        super().__init__(config)
+        self.ultrabert = UltrabertModel(config)
+        self.head = UltrabertPredictionHead(config)
+        self.cls = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.cls
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls = new_embeddings
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        token_type_ids: 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:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+            config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
+            loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
+        """
+        outputs = self.ultrabert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+        prediction_scores = self.cls(self.head(outputs[0]))
+
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()  # -100 index = padding token
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)
+            )
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    Ultrabert Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled
+    output) e.g. for GLUE tasks.
+    """
+)
+class UltrabertForSequenceClassification(UltrabertPreTrainedModel):
+    def __init__(self, config: UltrabertConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.ultrabert = UltrabertModel(config)
+        classifier_dropout = (
+            config.classifier_dropout
+            if config.classifier_dropout is not None
+            else config.hidden_dropout_prob
+        )
+        self.dropout = 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.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        token_type_ids: 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:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        outputs = self.ultrabert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+        logits = self.classifier(self.dropout(outputs[0][:, 0]))
+
+        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
+class UltrabertForQuestionAnswering(UltrabertPreTrainedModel):
+    def __init__(self, config: UltrabertConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.ultrabert = UltrabertModel(config)
+        self.qa_outputs = 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.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        token_type_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        start_positions: torch.Tensor | None = None,
+        end_positions: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor] | QuestionAnsweringModelOutput:
+        outputs = self.ultrabert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+        logits = self.qa_outputs(outputs[0])
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+
+            ignored_index = start_logits.size(1)
+            start_positions = start_positions.clamp(0, ignored_index)
+            end_positions = end_positions.clamp(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        return QuestionAnsweringModelOutput(
+            loss=total_loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@auto_docstring
+class UltrabertForTokenClassification(UltrabertPreTrainedModel):
+    def __init__(self, config: UltrabertConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.ultrabert = UltrabertModel(config)
+        classifier_dropout = (
+            config.classifier_dropout
+            if config.classifier_dropout is not None
+            else config.hidden_dropout_prob
+        )
+        self.dropout = 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.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        token_type_ids: 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:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
+        """
+        outputs = self.ultrabert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+        sequence_output = outputs[0]
+
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        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,
+        )

tokenizer_config.json

@@ -0,0 +1,36 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<repo_name>",
+    "<reponame>",
+    "<file_sep>",
+    "<filename>",
+    "<gh_stars>",
+    "<issue_start>",
+    "<issue_comment>",
+    "<issue_closed>",
+    "<jupyter_start>",
+    "<jupyter_text>",
+    "<jupyter_code>",
+    "<jupyter_output>",
+    "<|sep|>",
+    "<|mask|>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<|mask|>",
+  "model_max_length": 32768,
+  "pad_token": "<|endoftext|>",
+  "sep_token": "<|sep|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>",
+  "vocab_size": 49152
+}