testing-dummy-UltraVBERT / modeling_ultravbert.py

AntonV HF Staff

Duplicate from athrael-soju/UltraVBERT

214d44d 19 days ago

4.91 kB

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",
]

	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",
	]