robbert-2026 / model.py

Store weights in bfloat16 (1.10 GB -> 552 MB)

a406b1a verified about 2 months ago

22.2 kB

# From https://github.com/facebookresearch/llama/blob/main/llama/model.py

import torch
from torch import nn

from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
from torch.nn.functional import scaled_dot_product_attention

from typing import Optional, Tuple
import numpy as np

try:  # optional: xformers provides a fused SwiGLU when it is installed and importable
    from xformers.ops import SwiGLU
except Exception:  # not installed, or broken against the local torch/CUDA build
    from torch.nn.functional import silu as _silu

    class SwiGLU(nn.Module):
        """Pure-torch drop-in for ``xformers.ops.SwiGLU``.
        Same packed parameter layout as xformers, so checkpoints saved with the fused
        kernel load unchanged: ``w12`` is Linear(in, 2*hidden), split into (x1, x2);
        the activation is silu(x1) * x2; ``w3`` projects back to out_features.
        """

        def __init__(self, in_features, hidden_features, out_features=None, bias=True,
                     *, _pack_weights=True):
            super().__init__()
            out_features = in_features if out_features is None else out_features
            self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
            self.w3 = nn.Linear(hidden_features, out_features, bias=bias)

        def forward(self, x):
            x1, x2 = self.w12(x).chunk(2, dim=-1)
            return self.w3(_silu(x1) * x2)

try:
    from flash_attn.flash_attn_interface import flash_attn_varlen_func

    FLASH_ATTN_AVAILABLE = True
except ImportError:
    FLASH_ATTN_AVAILABLE = False

from transformers import (
    PreTrainedModel,
    PretrainedConfig,
    DataCollatorForLanguageModeling,
)
from transformers.modeling_outputs import (
    BaseModelOutput,
    MaskedLMOutput,
    SequenceClassifierOutput,
)

from .rotary import precompute_freqs_cis, apply_rotary_emb


class DataCollatorWithPacking(DataCollatorForLanguageModeling):
    def __init__(self, pack_sequences=False, **kwargs):
        super().__init__(**kwargs)
        self.pack_sequences = pack_sequences

    def __call__(self, batch):
        if self.pack_sequences:
            # Add position_ids if not present
            if "position_ids" not in batch[0]:
                for item in batch:
                    item["position_ids"] = list(range(len(item["input_ids"])))

            # Pack the sequences into a single list
            input_ids_list = [item["input_ids"] for item in batch]
            position_ids_list = [item["position_ids"] for item in batch]
            seqlens = np.array([0] + [len(ids) for ids in input_ids_list])

            packed_batch = {
                "position_ids": np.concatenate(position_ids_list, axis=0),
                "input_ids": np.concatenate(input_ids_list, axis=0),
                "cu_seqlens": np.cumsum(seqlens),
                "max_seqlen": max(seqlens),
            }

            batch = super().__call__([packed_batch])
            batch["cu_seqlens"] = batch["cu_seqlens"].to(torch.int32).squeeze()
        else:
            batch = super().__call__(batch)
            batch["attention_mask"] = batch["attention_mask"].to(torch.bool)

        return batch


class NeoBERTConfig(PretrainedConfig):
    model_type = "neobert"

    # All config parameters must have a default value.
    def __init__(
        self,
        hidden_size: int = 768,
        num_hidden_layers: int = 28,
        num_attention_heads: int = 12,
        intermediate_size: int = 3072,
        embedding_init_range: float = 0.02,
        decoder_init_range: float = 0.02,
        norm_eps: float = 1e-06,
        vocab_size: int = 30522,
        pad_token_id: int = 0,
        max_length: int = 1024,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        if hidden_size % num_attention_heads != 0:
            raise ValueError("Hidden size must be divisible by the number of heads.")
        self.dim_head = hidden_size // num_attention_heads
        self.intermediate_size = intermediate_size
        self.embedding_init_range = embedding_init_range
        self.decoder_init_range = decoder_init_range
        self.norm_eps = norm_eps
        self.vocab_size = vocab_size
        self.pad_token_id = pad_token_id
        self.max_length = max_length
        self.kwargs = kwargs


class EncoderBlock(nn.Module):
    """Transformer encoder block."""

    def __init__(self, config: NeoBERTConfig):
        super().__init__()

        self.config = config

        # Attention
        self.qkv = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size * 3, bias=False)
        self.wo = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size, bias=False)

        # Feedforward network
        multiple_of = 8
        intermediate_size = int(2 * config.intermediate_size / 3)
        intermediate_size = multiple_of * ((intermediate_size + multiple_of - 1) // multiple_of)
        self.ffn = SwiGLU(config.hidden_size, intermediate_size, config.hidden_size, bias=False)

        # Layer norms
        self.attention_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)
        self.ffn_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: torch.Tensor,
        freqs_cis: torch.Tensor,
        output_attentions: bool,
        max_seqlen: int = None,
        cu_seqlens: torch.Tensor = None,
    ):
        # Attention
        attn_output, attn_weights = self._att_block(
            self.attention_norm(x), attention_mask, freqs_cis, output_attentions, max_seqlen, cu_seqlens
        )

        # Residual
        x = x + attn_output

        # Feed-forward
        x = x + self.ffn(self.ffn_norm(x))

        return x, attn_weights

    def _att_block(
        self,
        x: torch.Tensor,
        attention_mask: torch.Tensor,
        freqs_cis: torch.Tensor,
        output_attentions: bool,
        max_seqlen: int = None,
        cu_seqlens: torch.Tensor = None,
    ):
        batch_size, seq_len, _ = x.shape

        xq, xk, xv = self.qkv(x).view(batch_size, seq_len, self.config.num_attention_heads, self.config.dim_head * 3).chunk(3, axis=-1)

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis)

        # Attn block
        attn_weights = None

        # Flash attention if the tensors are packed
        if cu_seqlens is not None:
            attn = flash_attn_varlen_func(
                q=xq.squeeze(0),
                k=xk.squeeze(0),
                v=xv.squeeze(0),
                cu_seqlens_q=cu_seqlens,
                cu_seqlens_k=cu_seqlens,
                max_seqlen_q=max_seqlen,
                max_seqlen_k=max_seqlen,
                dropout_p=0.0,
                causal=False,
            )
        # Eager attention if attention weights are needed in the output
        elif output_attentions:
            attn_weights = xq.permute(0, 2, 1, 3) @ xk.permute(0, 2, 3, 1) / (xq.size(-1) ** 0.5)
            if attention_mask is not None:
                attn_weights = attn_weights * attention_mask
            attn_weights = attn_weights.softmax(-1)
            attn = attn_weights @ xv.permute(0, 2, 1, 3)
            attn = attn.transpose(1, 2)
        # Fall back to SDPA otherwise
        else:
            attn = scaled_dot_product_attention(
                query=xq.transpose(1, 2),
                key=xk.transpose(1, 2),
                value=xv.transpose(1, 2),
                attn_mask=attention_mask.bool(),
                dropout_p=0,
            ).transpose(1, 2)

        return self.wo(attn.reshape(batch_size, seq_len, self.config.num_attention_heads * self.config.dim_head)), attn_weights


class NeoBERTPreTrainedModel(PreTrainedModel):
    config_class = NeoBERTConfig
    base_model_prefix = "model"
    _supports_cache_class = True

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.uniform_(-self.config.decoder_init_range, self.config.decoder_init_range)
        elif isinstance(module, nn.Embedding):
            module.weight.data.uniform_(-self.config.embedding_init_range, self.config.embedding_init_range)


class NeoBERT(NeoBERTPreTrainedModel):
    config_class = NeoBERTConfig

    def __init__(self, config: NeoBERTConfig):
        super().__init__(config)

        self.config = config

        self.encoder = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)

        # Ensures freqs_cis is moved to the same devices as the model. Non-persistent buffers are not saved in the state_dict.
        freqs_cis = precompute_freqs_cis(config.hidden_size // config.num_attention_heads, config.max_length)
        # Persisted as a REAL view: transformers' safetensors loader rejects complex dtypes
        # ("Cannot load safetensors of unknown dtype C64"), so we store float32 [L, D/2, 2]
        # and expose the complex tensor through the `freqs_cis` property below. Persisting it
        # matters because transformers>=5 materialises modules without re-running this
        # computation, which previously left the buffer uninitialised -> NaN logits.
        self.register_buffer("freqs_cis_real", torch.view_as_real(freqs_cis).contiguous(), persistent=True)

        self.transformer_encoder = nn.ModuleList()
        for _ in range(config.num_hidden_layers):
            self.transformer_encoder.append(EncoderBlock(config))

        self.layer_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)

        # Initialize weights and apply final processing
        self.post_init()

    @property
    def freqs_cis(self):
        # Complex view over the persisted real buffer (see __init__).
        # .float(): view_as_complex rejects bfloat16, and from_pretrained may cast
        # buffers to the config dtype. RoPE frequencies stay in fp32 regardless.
        return torch.view_as_complex(self.freqs_cis_real.float())

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        position_ids: torch.Tensor = None,
        max_seqlen: int = None,
        cu_seqlens: torch.Tensor = None,
        attention_mask: torch.Tensor = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        output_hidden_states: bool = False,
        output_attentions: bool = False,
        **kwargs,
    ):
        # Initialize
        hidden_states, attentions = [], []

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        # Expand and repeat: (Batch, Length) -> (Batch, Heads, Length, Length)
        if attention_mask is not None:
            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).repeat(1, self.config.num_attention_heads, attention_mask.size(-1), 1)

        # Checks to be done if inputs are packed sequences
        if cu_seqlens is not None:
            assert (
                FLASH_ATTN_AVAILABLE
            ), "Flash-attention is not available. Please ''pip install flash_attn'', or provide un-packed sequences."
            assert not output_attentions, "Output attentions is not supported when sequences are packed."
            assert max_seqlen is not None, "Missing max_seqlen. It must be provided when cu_seqlens are not None."
            assert (input_ids if input_ids is not None else inputs_embeds).shape[
                0
            ] == 1, "Cumulative sequence lengths are provided but inputs are not packed."
            assert (
                input_ids if input_ids is not None else inputs_embeds
            ).is_cuda, "Packing uses an implementation of flash-attention and is only supported on GPU."

        # RoPE. freqs_cis is a NON-persistent buffer, so it is not in the checkpoint;
        # transformers>=5 materialises the module without re-running the __init__ computation
        # that fills it, which leaves it uninitialised and yields NaN logits. Rebuild once.
        if not getattr(self, "_freqs_cis_checked", False):
            _fc = getattr(self, "freqs_cis_real", None)
            if _fc is None or _fc.numel() == 0 or not torch.isfinite(_fc).all():
                self.register_buffer(
                    "freqs_cis_real",
                    torch.view_as_real(
                        precompute_freqs_cis(
                            self.config.hidden_size // self.config.num_attention_heads,
                            self.config.max_length,
                        )
                    ).contiguous().to(self.encoder.weight.device),
                    persistent=True,
                )
            self._freqs_cis_checked = True

        freqs_cis = (
            self.freqs_cis[position_ids]
            if position_ids is not None
            else self.freqs_cis[: (input_ids if input_ids is not None else inputs_embeds).shape[1]].unsqueeze(0)
        )

        # Embedding
        x = self.encoder(input_ids) if input_ids is not None else inputs_embeds

        # Transformer encoder
        for layer in self.transformer_encoder:
            x, attn = layer(x, attention_mask, freqs_cis, output_attentions, max_seqlen, cu_seqlens)
            if output_hidden_states:
                hidden_states.append(x)
            if output_attentions:
                attentions.append(attn)

        # Final normalization layer
        x = self.layer_norm(x)

        # Return the output of the last hidden layer
        return BaseModelOutput(
            last_hidden_state=x,
            hidden_states=hidden_states if output_hidden_states else None,
            attentions=attentions if output_attentions else None,
        )


class NeoBERTLMHead(NeoBERTPreTrainedModel):
    config_class = NeoBERTConfig

    def __init__(self, config: NeoBERTConfig):
        super().__init__(config)

        self.config = config

        self.model = NeoBERT(config)
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size)

        self.post_init()

    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor = None,
        max_seqlen: int = None,
        cu_seqlens: torch.Tensor = None,
        attention_mask: torch.Tensor = None,
        output_hidden_states: bool = False,
        output_attentions: bool = False,
        **kwargs,
    ):

        output = self.model.forward(
            input_ids=input_ids,
            position_ids=position_ids,
            max_seqlen=max_seqlen,
            cu_seqlens=cu_seqlens,
            attention_mask=attention_mask,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
        )
        logits = self.decoder(output.last_hidden_state)

        return MaskedLMOutput(
            hidden_states=output.hidden_states if output_hidden_states else None,
            attentions=output.attentions if output_attentions else None,
            logits=logits,
        )


class NeoBERTForSequenceClassification(NeoBERTPreTrainedModel):
    config_class = NeoBERTConfig

    def __init__(self, config: NeoBERTConfig):
        super().__init__(config)

        self.config = config

        self.num_labels = getattr(config, "num_labels", 2)
        self.classifier_dropout = getattr(config, "classifier_dropout", 0.1)
        self.classifier_init_range = getattr(config, "classifier_init_range", 0.02)

        self.model = NeoBERT(config)

        self.dense = nn.Linear(self.config.hidden_size, self.config.hidden_size)
        self.dropout = nn.Dropout(self.classifier_dropout)
        self.classifier = nn.Linear(self.config.hidden_size, self.num_labels)

        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.classifier_init_range)
            if module.bias is not None:
                module.bias.data.zero_()

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        position_ids: torch.Tensor = None,
        max_seqlen: int = None,
        cu_seqlens: torch.Tensor = None,
        attention_mask: torch.Tensor = None,
        output_hidden_states: bool = False,
        output_attentions: bool = False,
        labels: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        
        output = self.model.forward(
            input_ids=input_ids,
            position_ids=position_ids,
            max_seqlen=max_seqlen,
            cu_seqlens=cu_seqlens,
            attention_mask=attention_mask,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
        )
        hidden_states = output.last_hidden_state

        x = hidden_states[:, 0, :]
        x = self.dropout(x)
        x = self.dense(x)
        x = torch.tanh(x)
        x = self.dropout(x)

        logits = self.classifier(x)

        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)

        if not return_dict:
            result = (logits,)
            return ((loss,) + result) if loss is not None else result

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=output.hidden_states if output_hidden_states else None,
            attentions=output.attentions if output_attentions else None,
        )


# ---------------------------------------------------------------------------
# Span-prediction head, added by neobert-nl/dumb_eval/add_qa_head.py.
# Mirrors NeoBERTForSequenceClassification's structure and BertForQuestionAnswering's
# loss so DUMB's squadnl scores are comparable to the baselines' out of the box.
# ---------------------------------------------------------------------------
from transformers.modeling_outputs import QuestionAnsweringModelOutput


class NeoBERTForQuestionAnswering(NeoBERTPreTrainedModel):
    config_class = NeoBERTConfig

    def __init__(self, config: NeoBERTConfig):
        super().__init__(config)

        self.config = config
        self.classifier_init_range = getattr(config, "classifier_init_range", 0.02)

        self.model = NeoBERT(config)
        self.qa_outputs = nn.Linear(config.hidden_size, 2)

        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.classifier_init_range)
            if module.bias is not None:
                module.bias.data.zero_()

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        position_ids: torch.Tensor = None,
        max_seqlen: int = None,
        cu_seqlens: torch.Tensor = None,
        attention_mask: torch.Tensor = None,
        output_hidden_states: bool = False,
        output_attentions: bool = False,
        start_positions: Optional[torch.Tensor] = None,
        end_positions: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        output = self.model.forward(
            input_ids=input_ids,
            position_ids=position_ids,
            max_seqlen=max_seqlen,
            cu_seqlens=cu_seqlens,
            attention_mask=attention_mask,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
        )

        logits = self.qa_outputs(output.last_hidden_state)
        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 start_positions.dim() > 1:
                start_positions = start_positions.squeeze(-1)
            if end_positions.dim() > 1:
                end_positions = end_positions.squeeze(-1)
            # Positions outside the window (unanswerable / truncated span) are ignored.
            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)
            total_loss = (loss_fct(start_logits, start_positions) + loss_fct(end_logits, end_positions)) / 2

        if not return_dict:
            result = (start_logits, end_logits)
            return ((total_loss,) + result) if total_loss is not None else result

        return QuestionAnsweringModelOutput(
            loss=total_loss,
            start_logits=start_logits,
            end_logits=end_logits,
            hidden_states=output.hidden_states if output_hidden_states else None,
            attentions=output.attentions if output_attentions else None,
        )

	# From https://github.com/facebookresearch/llama/blob/main/llama/model.py

	import torch
	from torch import nn

	from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
	from torch.nn.functional import scaled_dot_product_attention

	from typing import Optional, Tuple
	import numpy as np

	try: # optional: xformers provides a fused SwiGLU when it is installed and importable
	from xformers.ops import SwiGLU
	except Exception: # not installed, or broken against the local torch/CUDA build
	from torch.nn.functional import silu as _silu

	class SwiGLU(nn.Module):
	"""Pure-torch drop-in for ``xformers.ops.SwiGLU``.

	Same packed parameter layout as xformers, so checkpoints saved with the fused
	kernel load unchanged: ``w12`` is Linear(in, 2*hidden), split into (x1, x2);
	the activation is silu(x1) * x2; ``w3`` projects back to out_features.
	"""

	def __init__(self, in_features, hidden_features, out_features=None, bias=True,
	*, _pack_weights=True):
	super().__init__()
	out_features = in_features if out_features is None else out_features
	self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
	self.w3 = nn.Linear(hidden_features, out_features, bias=bias)

	def forward(self, x):
	x1, x2 = self.w12(x).chunk(2, dim=-1)
	return self.w3(_silu(x1) * x2)

	try:
	from flash_attn.flash_attn_interface import flash_attn_varlen_func

	FLASH_ATTN_AVAILABLE = True
	except ImportError:
	FLASH_ATTN_AVAILABLE = False

	from transformers import (
	PreTrainedModel,
	PretrainedConfig,
	DataCollatorForLanguageModeling,
	)
	from transformers.modeling_outputs import (
	BaseModelOutput,
	MaskedLMOutput,
	SequenceClassifierOutput,
	)

	from .rotary import precompute_freqs_cis, apply_rotary_emb


	class DataCollatorWithPacking(DataCollatorForLanguageModeling):
	def __init__(self, pack_sequences=False, **kwargs):
	super().__init__(**kwargs)
	self.pack_sequences = pack_sequences

	def __call__(self, batch):
	if self.pack_sequences:
	# Add position_ids if not present
	if "position_ids" not in batch[0]:
	for item in batch:
	item["position_ids"] = list(range(len(item["input_ids"])))

	# Pack the sequences into a single list
	input_ids_list = [item["input_ids"] for item in batch]
	position_ids_list = [item["position_ids"] for item in batch]
	seqlens = np.array([0] + [len(ids) for ids in input_ids_list])

	packed_batch = {
	"position_ids": np.concatenate(position_ids_list, axis=0),
	"input_ids": np.concatenate(input_ids_list, axis=0),
	"cu_seqlens": np.cumsum(seqlens),
	"max_seqlen": max(seqlens),
	}

	batch = super().__call__([packed_batch])
	batch["cu_seqlens"] = batch["cu_seqlens"].to(torch.int32).squeeze()
	else:
	batch = super().__call__(batch)
	batch["attention_mask"] = batch["attention_mask"].to(torch.bool)

	return batch


	class NeoBERTConfig(PretrainedConfig):
	model_type = "neobert"

	# All config parameters must have a default value.
	def __init__(
	self,
	hidden_size: int = 768,
	num_hidden_layers: int = 28,
	num_attention_heads: int = 12,
	intermediate_size: int = 3072,
	embedding_init_range: float = 0.02,
	decoder_init_range: float = 0.02,
	norm_eps: float = 1e-06,
	vocab_size: int = 30522,
	pad_token_id: int = 0,
	max_length: int = 1024,
	**kwargs,
	):
	super().__init__(**kwargs)

	self.hidden_size = hidden_size
	self.num_hidden_layers = num_hidden_layers
	self.num_attention_heads = num_attention_heads
	if hidden_size % num_attention_heads != 0:
	raise ValueError("Hidden size must be divisible by the number of heads.")
	self.dim_head = hidden_size // num_attention_heads
	self.intermediate_size = intermediate_size
	self.embedding_init_range = embedding_init_range
	self.decoder_init_range = decoder_init_range
	self.norm_eps = norm_eps
	self.vocab_size = vocab_size
	self.pad_token_id = pad_token_id
	self.max_length = max_length
	self.kwargs = kwargs


	class EncoderBlock(nn.Module):
	"""Transformer encoder block."""

	def __init__(self, config: NeoBERTConfig):
	super().__init__()

	self.config = config

	# Attention
	self.qkv = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size * 3, bias=False)
	self.wo = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size, bias=False)

	# Feedforward network
	multiple_of = 8
	intermediate_size = int(2 * config.intermediate_size / 3)
	intermediate_size = multiple_of * ((intermediate_size + multiple_of - 1) // multiple_of)
	self.ffn = SwiGLU(config.hidden_size, intermediate_size, config.hidden_size, bias=False)

	# Layer norms
	self.attention_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)
	self.ffn_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)

	def forward(
	self,
	x: torch.Tensor,
	attention_mask: torch.Tensor,
	freqs_cis: torch.Tensor,
	output_attentions: bool,
	max_seqlen: int = None,
	cu_seqlens: torch.Tensor = None,
	):
	# Attention
	attn_output, attn_weights = self._att_block(
	self.attention_norm(x), attention_mask, freqs_cis, output_attentions, max_seqlen, cu_seqlens
	)

	# Residual
	x = x + attn_output

	# Feed-forward
	x = x + self.ffn(self.ffn_norm(x))

	return x, attn_weights

	def _att_block(
	self,
	x: torch.Tensor,
	attention_mask: torch.Tensor,
	freqs_cis: torch.Tensor,
	output_attentions: bool,
	max_seqlen: int = None,
	cu_seqlens: torch.Tensor = None,
	):
	batch_size, seq_len, _ = x.shape

	xq, xk, xv = self.qkv(x).view(batch_size, seq_len, self.config.num_attention_heads, self.config.dim_head * 3).chunk(3, axis=-1)

	xq, xk = apply_rotary_emb(xq, xk, freqs_cis)

	# Attn block
	attn_weights = None

	# Flash attention if the tensors are packed
	if cu_seqlens is not None:
	attn = flash_attn_varlen_func(
	q=xq.squeeze(0),
	k=xk.squeeze(0),
	v=xv.squeeze(0),
	cu_seqlens_q=cu_seqlens,
	cu_seqlens_k=cu_seqlens,
	max_seqlen_q=max_seqlen,
	max_seqlen_k=max_seqlen,
	dropout_p=0.0,
	causal=False,
	)
	# Eager attention if attention weights are needed in the output
	elif output_attentions:
	attn_weights = xq.permute(0, 2, 1, 3) @ xk.permute(0, 2, 3, 1) / (xq.size(-1) ** 0.5)
	if attention_mask is not None:
	attn_weights = attn_weights * attention_mask
	attn_weights = attn_weights.softmax(-1)
	attn = attn_weights @ xv.permute(0, 2, 1, 3)
	attn = attn.transpose(1, 2)
	# Fall back to SDPA otherwise
	else:
	attn = scaled_dot_product_attention(
	query=xq.transpose(1, 2),
	key=xk.transpose(1, 2),
	value=xv.transpose(1, 2),
	attn_mask=attention_mask.bool(),
	dropout_p=0,
	).transpose(1, 2)

	return self.wo(attn.reshape(batch_size, seq_len, self.config.num_attention_heads * self.config.dim_head)), attn_weights


	class NeoBERTPreTrainedModel(PreTrainedModel):
	config_class = NeoBERTConfig
	base_model_prefix = "model"
	_supports_cache_class = True

	def _init_weights(self, module):
	if isinstance(module, nn.Linear):
	module.weight.data.uniform_(-self.config.decoder_init_range, self.config.decoder_init_range)
	elif isinstance(module, nn.Embedding):
	module.weight.data.uniform_(-self.config.embedding_init_range, self.config.embedding_init_range)


	class NeoBERT(NeoBERTPreTrainedModel):
	config_class = NeoBERTConfig

	def __init__(self, config: NeoBERTConfig):
	super().__init__(config)

	self.config = config

	self.encoder = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)

	# Ensures freqs_cis is moved to the same devices as the model. Non-persistent buffers are not saved in the state_dict.
	freqs_cis = precompute_freqs_cis(config.hidden_size // config.num_attention_heads, config.max_length)
	# Persisted as a REAL view: transformers' safetensors loader rejects complex dtypes
	# ("Cannot load safetensors of unknown dtype C64"), so we store float32 [L, D/2, 2]
	# and expose the complex tensor through the `freqs_cis` property below. Persisting it
	# matters because transformers>=5 materialises modules without re-running this
	# computation, which previously left the buffer uninitialised -> NaN logits.
	self.register_buffer("freqs_cis_real", torch.view_as_real(freqs_cis).contiguous(), persistent=True)

	self.transformer_encoder = nn.ModuleList()
	for _ in range(config.num_hidden_layers):
	self.transformer_encoder.append(EncoderBlock(config))

	self.layer_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)

	# Initialize weights and apply final processing
	self.post_init()

	@property
	def freqs_cis(self):
	# Complex view over the persisted real buffer (see __init__).
	# .float(): view_as_complex rejects bfloat16, and from_pretrained may cast
	# buffers to the config dtype. RoPE frequencies stay in fp32 regardless.
	return torch.view_as_complex(self.freqs_cis_real.float())

	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	position_ids: torch.Tensor = None,
	max_seqlen: int = None,
	cu_seqlens: torch.Tensor = None,
	attention_mask: torch.Tensor = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	output_hidden_states: bool = False,
	output_attentions: bool = False,
	**kwargs,
	):
	# Initialize
	hidden_states, attentions = [], []

	if (input_ids is None) ^ (inputs_embeds is not None):
	raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

	# Expand and repeat: (Batch, Length) -> (Batch, Heads, Length, Length)
	if attention_mask is not None:
	attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).repeat(1, self.config.num_attention_heads, attention_mask.size(-1), 1)

	# Checks to be done if inputs are packed sequences
	if cu_seqlens is not None:
	assert (
	FLASH_ATTN_AVAILABLE
	), "Flash-attention is not available. Please ''pip install flash_attn'', or provide un-packed sequences."
	assert not output_attentions, "Output attentions is not supported when sequences are packed."
	assert max_seqlen is not None, "Missing max_seqlen. It must be provided when cu_seqlens are not None."
	assert (input_ids if input_ids is not None else inputs_embeds).shape[
	0
	] == 1, "Cumulative sequence lengths are provided but inputs are not packed."
	assert (
	input_ids if input_ids is not None else inputs_embeds
	).is_cuda, "Packing uses an implementation of flash-attention and is only supported on GPU."

	# RoPE. freqs_cis is a NON-persistent buffer, so it is not in the checkpoint;
	# transformers>=5 materialises the module without re-running the __init__ computation
	# that fills it, which leaves it uninitialised and yields NaN logits. Rebuild once.
	if not getattr(self, "_freqs_cis_checked", False):
	_fc = getattr(self, "freqs_cis_real", None)
	if _fc is None or _fc.numel() == 0 or not torch.isfinite(_fc).all():
	self.register_buffer(
	"freqs_cis_real",
	torch.view_as_real(
	precompute_freqs_cis(
	self.config.hidden_size // self.config.num_attention_heads,
	self.config.max_length,
	)
	).contiguous().to(self.encoder.weight.device),
	persistent=True,
	)
	self._freqs_cis_checked = True

	freqs_cis = (
	self.freqs_cis[position_ids]
	if position_ids is not None
	else self.freqs_cis[: (input_ids if input_ids is not None else inputs_embeds).shape[1]].unsqueeze(0)
	)

	# Embedding
	x = self.encoder(input_ids) if input_ids is not None else inputs_embeds

	# Transformer encoder
	for layer in self.transformer_encoder:
	x, attn = layer(x, attention_mask, freqs_cis, output_attentions, max_seqlen, cu_seqlens)
	if output_hidden_states:
	hidden_states.append(x)
	if output_attentions:
	attentions.append(attn)

	# Final normalization layer
	x = self.layer_norm(x)

	# Return the output of the last hidden layer
	return BaseModelOutput(
	last_hidden_state=x,
	hidden_states=hidden_states if output_hidden_states else None,
	attentions=attentions if output_attentions else None,
	)


	class NeoBERTLMHead(NeoBERTPreTrainedModel):
	config_class = NeoBERTConfig

	def __init__(self, config: NeoBERTConfig):
	super().__init__(config)

	self.config = config

	self.model = NeoBERT(config)
	self.decoder = nn.Linear(config.hidden_size, config.vocab_size)

	self.post_init()

	def forward(
	self,
	input_ids: torch.Tensor,
	position_ids: torch.Tensor = None,
	max_seqlen: int = None,
	cu_seqlens: torch.Tensor = None,
	attention_mask: torch.Tensor = None,
	output_hidden_states: bool = False,
	output_attentions: bool = False,
	**kwargs,
	):

	output = self.model.forward(
	input_ids=input_ids,
	position_ids=position_ids,
	max_seqlen=max_seqlen,
	cu_seqlens=cu_seqlens,
	attention_mask=attention_mask,
	output_hidden_states=output_hidden_states,
	output_attentions=output_attentions,
	)
	logits = self.decoder(output.last_hidden_state)

	return MaskedLMOutput(
	hidden_states=output.hidden_states if output_hidden_states else None,
	attentions=output.attentions if output_attentions else None,
	logits=logits,
	)


	class NeoBERTForSequenceClassification(NeoBERTPreTrainedModel):
	config_class = NeoBERTConfig

	def __init__(self, config: NeoBERTConfig):
	super().__init__(config)

	self.config = config

	self.num_labels = getattr(config, "num_labels", 2)
	self.classifier_dropout = getattr(config, "classifier_dropout", 0.1)
	self.classifier_init_range = getattr(config, "classifier_init_range", 0.02)

	self.model = NeoBERT(config)

	self.dense = nn.Linear(self.config.hidden_size, self.config.hidden_size)
	self.dropout = nn.Dropout(self.classifier_dropout)
	self.classifier = nn.Linear(self.config.hidden_size, self.num_labels)

	self.post_init()

	def _init_weights(self, module):
	if isinstance(module, nn.Linear):
	module.weight.data.normal_(mean=0.0, std=self.classifier_init_range)
	if module.bias is not None:
	module.bias.data.zero_()

	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	position_ids: torch.Tensor = None,
	max_seqlen: int = None,
	cu_seqlens: torch.Tensor = None,
	attention_mask: torch.Tensor = None,
	output_hidden_states: bool = False,
	output_attentions: bool = False,
	labels: Optional[torch.Tensor] = None,
	return_dict: Optional[bool] = None,
	):
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	output = self.model.forward(
	input_ids=input_ids,
	position_ids=position_ids,
	max_seqlen=max_seqlen,
	cu_seqlens=cu_seqlens,
	attention_mask=attention_mask,
	output_hidden_states=output_hidden_states,
	output_attentions=output_attentions,
	)
	hidden_states = output.last_hidden_state

	x = hidden_states[:, 0, :]
	x = self.dropout(x)
	x = self.dense(x)
	x = torch.tanh(x)
	x = self.dropout(x)

	logits = self.classifier(x)

	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)

	if not return_dict:
	result = (logits,)
	return ((loss,) + result) if loss is not None else result

	return SequenceClassifierOutput(
	loss=loss,
	logits=logits,
	hidden_states=output.hidden_states if output_hidden_states else None,
	attentions=output.attentions if output_attentions else None,
	)


	# ---------------------------------------------------------------------------
	# Span-prediction head, added by neobert-nl/dumb_eval/add_qa_head.py.
	# Mirrors NeoBERTForSequenceClassification's structure and BertForQuestionAnswering's
	# loss so DUMB's squadnl scores are comparable to the baselines' out of the box.
	# ---------------------------------------------------------------------------
	from transformers.modeling_outputs import QuestionAnsweringModelOutput


	class NeoBERTForQuestionAnswering(NeoBERTPreTrainedModel):
	config_class = NeoBERTConfig

	def __init__(self, config: NeoBERTConfig):
	super().__init__(config)

	self.config = config
	self.classifier_init_range = getattr(config, "classifier_init_range", 0.02)

	self.model = NeoBERT(config)
	self.qa_outputs = nn.Linear(config.hidden_size, 2)

	self.post_init()

	def _init_weights(self, module):
	if isinstance(module, nn.Linear):
	module.weight.data.normal_(mean=0.0, std=self.classifier_init_range)
	if module.bias is not None:
	module.bias.data.zero_()

	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	position_ids: torch.Tensor = None,
	max_seqlen: int = None,
	cu_seqlens: torch.Tensor = None,
	attention_mask: torch.Tensor = None,
	output_hidden_states: bool = False,
	output_attentions: bool = False,
	start_positions: Optional[torch.Tensor] = None,
	end_positions: Optional[torch.Tensor] = None,
	return_dict: Optional[bool] = None,
	**kwargs,
	):
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	output = self.model.forward(
	input_ids=input_ids,
	position_ids=position_ids,
	max_seqlen=max_seqlen,
	cu_seqlens=cu_seqlens,
	attention_mask=attention_mask,
	output_hidden_states=output_hidden_states,
	output_attentions=output_attentions,
	)

	logits = self.qa_outputs(output.last_hidden_state)
	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 start_positions.dim() > 1:
	start_positions = start_positions.squeeze(-1)
	if end_positions.dim() > 1:
	end_positions = end_positions.squeeze(-1)
	# Positions outside the window (unanswerable / truncated span) are ignored.
	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)
	total_loss = (loss_fct(start_logits, start_positions) + loss_fct(end_logits, end_positions)) / 2

	if not return_dict:
	result = (start_logits, end_logits)
	return ((total_loss,) + result) if total_loss is not None else result

	return QuestionAnsweringModelOutput(
	loss=total_loss,
	start_logits=start_logits,
	end_logits=end_logits,
	hidden_states=output.hidden_states if output_hidden_states else None,
	attentions=output.attentions if output_attentions else None,
	)