Upload modeling_eurobert.py

modeling_eurobert.py

@@ -0,0 +1,136 @@
+from typing import Unpack
+import torch
+from transformers import DataCollatorWithFlattening
+from transformers.masking_utils import create_bidirectional_mask
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.models.eurobert import (
+    EuroBertForMaskedLM,
+    EuroBertModel,
+    EuroBertForSequenceClassification,
+    EuroBertForTokenClassification
+)
+from transformers.utils import TransformersKwargs
+
+
+def _unpad_input(input_ids: torch.Tensor, attention_mask: torch.Tensor):
+    collator = DataCollatorWithFlattening(return_flash_attn_kwargs=True)
+    features = collator([{"input_ids": i[a.bool()].tolist()} for i, a in zip(input_ids, attention_mask)])
+    return features
+
+
+def _pad_output(inputs: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int,) -> torch.Tensor:
+    if inputs.dim() == 3:
+        inputs = inputs.squeeze()
+    if inputs.dim() == 1:
+        output = torch.zeros(batch * seqlen, dtype=inputs.dtype, device=inputs.device)
+        output[indices] = inputs
+        padded_inputs = output.view(batch, seqlen)
+    else:
+        _, *rest = inputs.shape
+        output = torch.zeros(batch * seqlen, *rest, dtype=inputs.dtype, device=inputs.device)
+        output[indices] = inputs
+        padded_inputs = output.view(batch, seqlen, *rest)
+    return padded_inputs
+
+
+class UnpadEuroBertModel(EuroBertModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+    def forward(
+            self,
+            input_ids: torch.LongTensor = None,
+            attention_mask: torch.Tensor | None = None,
+            position_ids: torch.LongTensor | None = None,
+            inputs_embeds: torch.FloatTensor | None = None,
+            **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | 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 input_ids is not None:
+            device = input_ids.device
+            seq_length = input_ids.shape[1]
+            batch_size = input_ids.size(0)
+        else:
+            device = inputs_embeds.device
+            seq_length = inputs_embeds.shape[1]
+            batch_size = inputs_embeds.size(0)
+
+        indices = None
+        if self.config._attn_implementation.startswith("flash_attention"):
+            if input_ids is None or attention_mask is None:
+                raise ValueError("Unpadding requires both input_ids and attention_mask")
+            with torch.no_grad():
+                indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+                features = _unpad_input(input_ids, attention_mask)
+                input_ids = features["input_ids"].to(device=device)
+                position_ids = features["position_ids"].to(device=device)
+                attention_mask = None
+                kwargs["cu_seq_lens_k"] = features["cu_seq_lens_k"].to(device=device)
+                kwargs["cu_seq_lens_q"] = features["cu_seq_lens_q"].to(device=device)
+                kwargs["max_length_k"] = features["max_length_k"]
+                kwargs["max_length_q"] = features["max_length_q"]
+
+        if inputs_embeds is None:
+            inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+        if position_ids is None:
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0)
+
+        bidirectional_mask = create_bidirectional_mask(
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+        )
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+        for encoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = encoder_layer(
+                hidden_states,
+                attention_mask=bidirectional_mask,
+                position_embeddings=position_embeddings,
+                position_ids=position_ids,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        if self.config._attn_implementation.startswith("flash_attention"):
+            hidden_states = _pad_output(
+                inputs=hidden_states, indices=indices, batch=batch_size, seqlen=seq_length
+            )
+
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+        )
+
+
+class UnpadEuroBertForMaskedLM(EuroBertForMaskedLM):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = UnpadEuroBertModel(config)
+        self.post_init()
+
+
+class UnpadEuroBertForSequenceClassification(EuroBertForSequenceClassification):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = UnpadEuroBertModel(config)
+        self.post_init()
+
+
+class UnpadEuroBertForTokenClassification(EuroBertForTokenClassification):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = UnpadEuroBertModel(config)
+        self.post_init()
+
+
+def enable_eurobert_unpadding():
+    EuroBertModel.forward = UnpadEuroBertModel.forward