| 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) |
| |
| 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): |
| |
| 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"] |
|
|
