File size: 4,914 Bytes
214d44d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | from torch import nn
from torch.nn import CrossEntropyLoss
from transformers import AutoModel
from transformers.models.modernvbert.modeling_modernvbert import (
ModernVBertConnector,
ModernVBertForSequenceClassification,
ModernVBertForTokenClassification,
ModernVBertMaskedLMOutput,
ModernVBertModel,
ModernVBertPreTrainedModel,
)
from transformers.models.ultrabert.modeling_ultrabert import UltrabertPredictionHead, UltrabertRMSNorm
from .configuration_ultravbert import UltraVBertConfig
class UltraVBertConnector(ModernVBertConnector):
def __init__(self, config: UltraVBertConfig):
super().__init__(config)
# UltraBERT's first layer takes embeddings unnormalized; raw projections would swamp its token embeddings.
self.norm = (
UltrabertRMSNorm(config.text_config.hidden_size, eps=config.text_config.rms_norm_eps)
if config.connector_norm
else nn.Identity()
)
def forward(self, image_hidden_states):
return self.norm(super().forward(image_hidden_states))
class UltraVBertModel(ModernVBertModel):
config_class = UltraVBertConfig
_keys_to_ignore_on_load_unexpected = [r"vision_model\.head\."]
def __init__(self, config: UltraVBertConfig):
# ModernVBertModel.__init__ builds ModernVBertConnector; the rest of its body follows unchanged.
ModernVBertPreTrainedModel.__init__(self, config)
self.padding_idx = config.text_config.pad_token_id
self.vocab_size = config.text_config.vocab_size
self.vision_model = AutoModel.from_config(config.vision_config)
self.connector = UltraVBertConnector(config)
self.text_model = AutoModel.from_config(config.text_config)
self.image_seq_len = int(
((config.vision_config.image_size // config.vision_config.patch_size) ** 2)
/ (config.pixel_shuffle_factor**2)
)
self.image_token_id = config.image_token_id
self.post_init()
class UltraVBertForMaskedLM(ModernVBertPreTrainedModel):
config_class = UltraVBertConfig
_keys_to_ignore_on_load_unexpected = [r"vision_model\.head\."]
_tied_weights_keys = {"head.decoder.weight": "model.text_model.embed_tokens.weight"}
def __init__(self, config: UltraVBertConfig):
super().__init__(config)
self.vocab_size = config.text_config.vocab_size
self.model = UltraVBertModel(config)
self.head = UltrabertPredictionHead(config.text_config)
self.post_init()
def get_output_embeddings(self):
return self.head.decoder
def set_output_embeddings(self, new_embeddings):
self.head.decoder = new_embeddings
def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None, pixel_values=None,
pixel_attention_mask=None, image_hidden_states=None, labels=None, **kwargs):
outputs = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
inputs_embeds=inputs_embeds,
pixel_values=pixel_values,
pixel_attention_mask=pixel_attention_mask,
image_hidden_states=image_hidden_states,
**kwargs,
)
logits = self.head(outputs.last_hidden_state)
loss = None
if labels is not None:
loss = CrossEntropyLoss()(logits.view(-1, self.vocab_size), labels.view(-1))
return ModernVBertMaskedLMOutput(
loss=loss,
logits=logits,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
image_hidden_states=outputs.image_hidden_states,
)
def _build_classifier(model, config: UltraVBertConfig):
# ModernVBert's own __init__ builds a ModernBERT prediction head that UltraBERT's text config cannot configure.
ModernVBertPreTrainedModel.__init__(model, config)
model.num_labels = config.num_labels
model.model = UltraVBertModel(config)
model.head = nn.Identity()
model.drop = nn.Dropout(config.classifier_dropout)
model.classifier = nn.Linear(config.text_config.hidden_size, config.num_labels)
model.post_init()
class UltraVBertForSequenceClassification(ModernVBertForSequenceClassification):
config_class = UltraVBertConfig
_keys_to_ignore_on_load_unexpected = [r"vision_model\.head\."]
def __init__(self, config: UltraVBertConfig):
_build_classifier(self, config)
class UltraVBertForTokenClassification(ModernVBertForTokenClassification):
config_class = UltraVBertConfig
_keys_to_ignore_on_load_unexpected = [r"vision_model\.head\."]
def __init__(self, config: UltraVBertConfig):
_build_classifier(self, config)
__all__ = [
"UltraVBertModel",
"UltraVBertForMaskedLM",
"UltraVBertForSequenceClassification",
"UltraVBertForTokenClassification",
]
|