File size: 1,850 Bytes
e730d84 | 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 | from torch import nn
from transformers import AutoModel
from transformers.models.modernvbert.modeling_modernvbert import (
ModernVBertConnector,
ModernVBertModel,
ModernVBertPreTrainedModel,
)
from transformers.models.ultrabert.modeling_ultrabert import 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()
__all__ = ["UltraVBertModel"]
|